import gradio as gr
from transformers import pipeline

# Initialize the code generation pipeline
generator = pipeline('code-generation', model='Salesforce/codegen-350M-mono')

def generate_code(prompt, max_length=128):
  """Generates code based on the given prompt."""
  result = generator(prompt, max_length=max_length)
  return result[0]['generated_text']

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_code,
    inputs=[
        gr.Textbox(lines=5, placeholder="Enter your code prompt here..."),
        gr.Slider(minimum=64, maximum=512, step=8, value=128, label="Max Length")
    ],
    outputs=gr.Code(language="python"),
    title="🤖 Code Generation Demo",
    description="Generate code using an open-source language model."
)

iface.launch()