import gradio as gr
from transformers import pipeline

# Load the model using CPU (optimized for 8GB RAM)
generator = pipeline("text2text-generation", model="Salesforce/codet5p-220m", device=-1)

def generate_tests(code, instruction):
    prompt = f"Input:\n{code}\n\nInstruction:\n{instruction}"
    result = generator(prompt)[0]["generated_text"]
    return result

gr.Interface(
    fn=generate_tests,
    inputs=[
        gr.Textbox(label="Your Code", lines=10, placeholder="Paste your function here..."),
        gr.Textbox(label="Instruction", placeholder="e.g., Generate unit tests using Python unittest.")
    ],
    outputs=gr.Code(language="python"),
    title="🧪 Unit Test Generator (Lite)",
    description="Paste your function and get test cases using a small AI model. Optimized for 8GB RAM."
).launch()