usando unicamp model

app.py

@@ -1,27 +1,34 @@
 from fastapi import FastAPI
 from transformers import pipeline
-import json
 
 app = FastAPI()
 
-# Carregando o MedGemma direto do Hugging Face
-medgemma = pipeline(
-    "text-generation",
-    model="google/medgemma-4b-it",
-    tokenizer="google/medgemma-4b-it"
+# 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 sobre {theme}, nível {difficulty}, com 4 alternativas (A, B, C, D), no estilo da USP."
-    response = medgemma(prompt, max_new_tokens=300, temperature=0.7, top_p=0.95)[0]['generated_text']
+    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 médica estilo USP com 4 alternativas (A, B, C, D), nível médio."
-        question = medgemma(prompt, max_new_tokens=300)[0]['generated_text']
-        simulado.append({"question": question})
+        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}