import gradio as gr
from transformers import pipeline

# You can change the model to any light-weight code model
pipe = pipeline("text-generation", model="Salesforce/codegen-350M-mono")

def generate(prompt):
    result = pipe(prompt, max_new_tokens=200)
    return result[0]['generated_text']

gr.Interface(
    fn=generate,
    inputs=gr.Textbox(lines=4, placeholder="Enter code prompt here..."),
    outputs="text",
    title="Code Helper AI"
).launch()