import gradio as gr
from transformers import pipeline

# Load your model (will download first time, then cache forever)
generator = pipeline(
    "text-generation",
    model="suneeldk/DesiFamily-2B",
    max_new_tokens=180,
    do_sample=True,
    temperature=0.8,
    top_p=0.9,
    device_map="cpu"
)

def roast(statement, character):
    if not statement.strip():
        return "Write something first, lazy human!"
    
    prompt = f"Reply exactly like a {character} in pure English to this statement: {statement}"
    result = generator(prompt)[0]["generated_text"]
    reply = result.split("### Response:")[-1].strip()
    if not reply or len(reply) < 10:
        reply = result[len(prompt):].strip()
    return reply

# THIS IS THE ONLY WORKING VERSION FOR DEC 2025 SPACES
gr.Interface(
    fn=roast,
    inputs=[
        gr.Textbox(label="What did you just say/do?", placeholder="Ex: I want to quit my job and become a rapper...", lines=4),
        gr.Dropdown(
            choices=["Indian Mom", "Indian Dad", "Best Friend", "Random Uncle", 
                     "Jealous Aunty", "Annoying Cousin", "Supportive Sister", "Chill Grandpa"],
            label="Who is replying?", value="Indian Mom"
        )
    ],
    outputs=gr.Textbox(label="The Savage Reply", lines=8),
    title="DesiFamily-2B Roaster by suneeldk",
    description="Pure English Indian family roasts — 100% accurate, 0% mercy 😂",
    examples=[
        ["I failed my exam", "Indian Mom"],
        ["I want to marry my girlfriend", "Indian Dad"],
        ["I'm starting a startup", "Best Friend"],
        ["I'm moving to Canada", "Random Uncle"]
    ],
    # ←←←←← REMOVED allow_flagging (this was the only bug)
).launch()