import gradio as gr
from transformers import pipeline

# Load the text-generation pipeline from Hugging Face
code_gen = pipeline("text-generation", model="bigcode/starcoder", device_map="auto")

def generate_code(prompt, language):
    header = f"# Language: {language}\n"
    full_prompt = header + prompt

    output = code_gen(full_prompt, max_new_tokens=200, temperature=0.3, do_sample=True)[0]['generated_text']

    # Trim only the newly generated code part
    return output.replace(full_prompt, "").strip()

gr.Interface(
    fn=generate_code,
    inputs=[
        gr.Textbox(lines=4, label="Enter your code task or prompt"),
        gr.Dropdown(["Python", "JavaScript", "C++", "Java", "Bash", "HTML"], label="Target Language", value="Python"),
    ],
    outputs=gr.Code(label="Generated Code"),
    title="AI Code Generator",
    description="Generate code snippets using a Hugging Face LLM like StarCoder. Try writing prompts like 'sort a list of integers' or 'read a CSV and print the average column value'."
).launch()