import gradio as gr
import os
from transformers import pipeline

access_token = os.getenv("hf_token")


code_generator = pipeline('text-generation', model='google/codegemma-2b', token=access_token)

def generate_code(prompt):

    generated_code = code_generator(prompt, max_length=100, do_sample=True, temperature=0.7)
    return generated_code[0]['generated_text']

# Create the Gradio interface
iface = gr.Interface(
    #flagging of bc of bug
    allow_flagging="never",
    fn=generate_code,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs=gr.Textbox(label="Generated Code"),
    title="Gemma Code Generator",
    description="Enter a prompt to generate code using the little gemma model.",
)

# Launch the Gradio interface
iface.launch()