Update app.py

app.py

@@ -3,10 +3,9 @@ import torch
 import re
 from transformers import AutoTokenizer
 from peft import AutoPeftModelForCausalLM
-import json
 
 # Carrega modelo e tokenizer da Hugging Face - LoRA fine-tuned
-model_name = "habulaj/filterinstruct"
+model_name = "habulaj/filter"
 print("Carregando tokenizer e modelo (CPU)...")
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -29,6 +28,7 @@ except Exception as e:
 # -------- FASTAPI --------
 app = FastAPI(title="News Filter JSON API")
 
+# -------- ROOT ENDPOINT --------
 @app.get("/")
 def read_root():
     return {"message": "News Filter JSON API is running!", "docs": "/docs"}
@@ -43,37 +43,44 @@ Title: "{title}"
 Content: "{content}"
 """
 
+    # Otimizações de tokenização
     inputs = tokenizer(
-        prompt,
+        prompt, 
         return_tensors="pt",
         truncation=True,
-        max_length=512,
-        padding=False
+        max_length=512,  # Limita tamanho do input
+        padding=False    # Não faz padding desnecessário
     )
     input_ids = inputs.input_ids.to("cpu")
-
+    
     with torch.no_grad():
+        # Configurações otimizadas para velocidade
         outputs = model.generate(
             input_ids=input_ids,
-            max_new_tokens=100,
-            temperature=1.0,
+            max_new_tokens=100,      # Reduzido de 128 para 100
+            temperature=1.0,         # Reduzido para ser mais determinístico
             do_sample=True,
             top_p=0.9,
-            num_beams=1,
-            early_stopping=True,
+            num_beams=1,            # Beam search = 1 (greedy) é mais rápido
+            early_stopping=True,    # Para quando encontrar EOS
             eos_token_id=tokenizer.eos_token_id,
             pad_token_id=tokenizer.eos_token_id,
         )
-
+    
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Remove prompt do output
     generated = decoded[len(prompt):].strip()
-
+    
+    # Extrai JSON
     match = re.search(r"\{.*\}", generated, re.DOTALL)
     if match:
-        return match.group(0)
+        result = match.group(0)
+        return result
     else:
         return "⚠️ Failed to extract JSON. Output:\n" + generated
 
+# -------- API ROUTE --------
 @app.get("/filter")
 def get_filter(
     title: str = Query(..., description="Title of the news"),
@@ -81,8 +88,11 @@ def get_filter(
 ):
     try:
         json_output = infer_filter(title, content)
+        import json
+        # Retorna como dados brutos (parse do JSON)
         return json.loads(json_output)
     except json.JSONDecodeError:
+        # Se não conseguir fazer parse, retorna como string
         return {"raw_output": json_output}
     except Exception as e:
         raise HTTPException(status_code=422, detail=str(e))