import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Load the small code generation model
MODEL_NAME = "Salesforce/codegen-350M-multi"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Gradio interface
def generate_code(prompt):
    result = generator(prompt, max_length=200, do_sample=True)
    return result[0]['generated_text']

iface = gr.Interface(fn=generate_code, inputs="text", outputs="text",
                     title="CodeGen Demo", description="Generate code from prompts")
iface.launch()