from fastapi import FastAPI
from transformers import pipeline

app = FastAPI()

# Carregando o modelo correto: ptt5 (text2text-generation)
ptt5 = pipeline(
    "text2text-generation",
    model="unicamp-dl/ptt5-base-portuguese-text2text",
    tokenizer="unicamp-dl/ptt5-base-portuguese-text2text"
)

@app.get("/generate")
async def generate_question(theme: str, difficulty: str):
    prompt = (
        f"Gere uma questão de múltipla escolha sobre o tema '{theme}', "
        f"com dificuldade {difficulty}, no estilo da prova de residência médica da USP. "
        "Inclua: enunciado clínico, 4 alternativas (A, B, C, D), gabarito e explicação."
    )
    response = ptt5(prompt, max_new_tokens=512, temperature=0.7, top_p=0.9)[0]['generated_text']
    return {"question": response}

@app.get("/simulado")
async def get_simulado(num_questions: int = 5):
    simulado = []
    for _ in range(num_questions):
        prompt = (
            "Gere uma questão de múltipla escolha sobre clínica médica, nível médio, "
            "no estilo da prova de residência da USP. "
            "Inclua: enunciado clínico, 4 alternativas (A, B, C, D), gabarito e explicação."
        )
        result = ptt5(prompt, max_new_tokens=512, temperature=0.7, top_p=0.9)[0]['generated_text']
        simulado.append({"question": result})
    return {"simulado": simulado}