Create app.py

app.py

@@ -0,0 +1,125 @@
+from fastapi import FastAPI, Query, HTTPException
+from pydantic import BaseModel
+from transformers import AutoTokenizer
+from peft import AutoPeftModelForCausalLM
+import torch
+import re
+import json
+import os
+
+# -------- CONFIG --------
+MODEL_NAME = "habulaj/filter"
+DEVICE = "cpu"
+DTYPE = torch.float32
+
+# -------- LOAD MODEL --------
+print("🔁 Loading model and tokenizer...")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    
+    # Set pad_token if not exists
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    model = AutoPeftModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map=DEVICE,
+        torch_dtype=DTYPE,
+        trust_remote_code=True,
+    )
+    model.eval()
+    print("✅ Model loaded successfully.")
+    
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    model = None
+    tokenizer = None
+
+# -------- FASTAPI --------
+app = FastAPI(title="News Filter JSON API")
+
+# -------- ROOT ENDPOINT --------
+@app.get("/")
+def read_root():
+    return {
+        "message": "News Filter JSON API is running!",
+        "model_loaded": model is not None,
+        "docs": "/docs",
+        "endpoints": ["/filter", "/health"]
+    }
+
+# -------- INFERENCE FUNCTION --------
+def generate_json_filter(title: str, content: str) -> str:
+    if model is None or tokenizer is None:
+        raise ValueError("Model not loaded")
+    
+    prompt = f"""Analyze the news title and content, and return the filters in JSON format with the defined fields.
+
+Please respond ONLY with the JSON filter, do NOT add any explanations, system messages, or extra text.
+
+Title: "{title}"
+Content: "{content}"
+"""
+
+    try:
+        inputs = tokenizer(
+            prompt, 
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512
+        )
+        input_ids = inputs["input_ids"].to(DEVICE)
+        attention_mask = inputs["attention_mask"].to(DEVICE)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=128,
+                temperature=1.2,
+                top_p=0.9,
+                do_sample=True,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+
+        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        generated = decoded[len(prompt):].strip()
+
+        # Extrai JSON
+        match = re.search(r"\{.*\}", generated, re.DOTALL)
+        if match:
+            return match.group(0)
+        
+        # Fallback: retorna um JSON simples se não encontrar
+        return '{"status": "processed", "title": "' + title[:50] + '", "content_length": ' + str(len(content)) + '}'
+        
+    except Exception as e:
+        raise ValueError(f"Error during generation: {str(e)}")
+
+# -------- API ROUTE --------
+@app.get("/filter")
+def get_filter(
+    title: str = Query(..., description="Title of the news"),
+    content: str = Query(..., description="Content of the news")
+):
+    try:
+        json_output = generate_json_filter(title, content)
+        # Use json.loads instead of eval for safety
+        parsed_json = json.loads(json_output)
+        return {"filter": parsed_json}
+    except json.JSONDecodeError as e:
+        return {"error": "Invalid JSON generated", "raw_output": json_output}
+    except Exception as e:
+        raise HTTPException(status_code=422, detail=str(e))
+
+# -------- HEALTH CHECK --------
+@app.get("/health")
+def health_check():
+    return {
+        "status": "healthy",
+        "model_loaded": model is not None,
+        "device": DEVICE,
+        "torch_version": torch.__version__
+    }