small model

app.py

@@ -4,6 +4,7 @@ import json
 import os
 import logging
 import time
+import gc
 
 # Configura logging para capturar mais detalhes
 logging.basicConfig(level=logging.INFO)
@@ -27,13 +28,13 @@ def get_model():
         start_time = time.time()
         try:
             tokenizer = AutoTokenizer.from_pretrained(
-                "unicamp-dl/ptt5-base-portuguese-vocab",
+                "unicamp-dl/ptt5-small-portuguese-vocab",
                 legacy=False,
                 clean_up_tokenization_spaces=True
             )
             logger.info(f"Tokenizer carregado em {time.time() - start_time:.2f} segundos.")
             model = T5ForConditionalGeneration.from_pretrained(
-                "unicamp-dl/ptt5-base-portuguese-vocab",
+                "unicamp-dl/ptt5-small-portuguese-vocab",
                 device_map="cpu"  # Força uso da CPU
             )
             logger.info(f"Modelo carregado em {time.time() - start_time:.2f} segundos.")
@@ -68,23 +69,21 @@ def generate_question_from_prompt(theme, difficulty, example_question=None):
             f"Explicação: {example_question['explanation']}"
         )
         prompt = (
-            f"Baseado no exemplo: '{example_text}', gere uma nova questão de múltipla escolha "
-            f"sobre o tema '{theme}', com dificuldade {difficulty}, no estilo da prova de residência médica da USP. "
-            "Use o formato: 'Enunciado clínico: [texto detalhado com dados do paciente, exames e pergunta]. "
-            "Alternativas: A) [opção], B) [opção], C) [opção], D) [opção]. Gabarito: [letra]. "
-            "Explicação: [texto].'"
+            f"Baseado no exemplo: '{example_text}', gere uma questão de múltipla escolha "
+            f"sobre '{theme}', dificuldade {difficulty}, estilo prova de residência médica da USP. "
+            f"Formato: 'Enunciado clínico: [texto]. Alternativas: A) [opção], B) [opção], C) [opção], D) [opção]. "
+            f"Gabarito: [letra]. Explicação: [texto].'"
         )
     else:
         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. "
-            "Use o formato: 'Enunciado clínico: [texto detalhado com dados do paciente, exames e pergunta]. "
-            "Alternativas: A) [opção], B) [opção], C) [opção], D) [opção]. Gabarito: [letra]. "
-            "Explicação: [texto].'"
+            f"Gere uma questão de múltipla escolha sobre '{theme}', dificuldade {difficulty}, "
+            f"estilo prova de residência médica da USP. Formato: "
+            f"'Enunciado clínico: [texto]. Alternativas: A) [opção], B) [opção], C) [opção], D) [opção]. "
+            f"Gabarito: [letra]. Explicação: [texto].'"
         )
     try:
         inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
-        outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7, top_p=0.9)
+        outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.7, top_p=0.9)
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         parts = response.split("Alternativas:")
         if len(parts) > 1:
@@ -94,14 +93,20 @@ def generate_question_from_prompt(theme, difficulty, example_question=None):
             explanation_part = parts[1].split("Explicação:")[1].strip() if "Explicação:" in parts[1] else "Explicação padrão"
             options = [opt.strip() for opt in options_part.split(",")]
             if len(options) >= 4:
-                return {
+                result = {
                     "question": f"Enunciado clínico: {question_part}",
                     "options": [f"A) {options[0]}", f"B) {options[1]}", f"C) {options[2]}", f"D) {options[3]}"],
                     "answer": answer_part,
                     "explanation": explanation_part
                 }
-        return {"question": response, "options": [], "answer": "", "explanation": "Explicação padrão"}
+                # Libera memória
+                gc.collect()
+                return result
+        result = {"question": response, "options": [], "answer": "", "explanation": "Explicação padrão"}
+        gc.collect()
+        return result
     except Exception as e:
+        gc.collect()
         return {"question": f"Erro na geração: {e}", "options": [], "answer": "", "explanation": "Tente novamente."}
 
 @app.get("/generate")