from fastapi import FastAPI, Query, HTTPException import torch import re import time import logging import os from transformers import AutoTokenizer, GenerationConfig from peft import AutoPeftModelForCausalLM import gc # -------- CONFIGURAÇÕES DE OTIMIZAÇÃO -------- os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["OMP_NUM_THREADS"] = "2" os.environ["MKL_NUM_THREADS"] = "2" torch.set_num_threads(2) torch.set_num_interop_threads(1) # -------- LOGGING CONFIG -------- logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") log = logging.getLogger("news-filter") # -------- LOAD MODEL -------- model_name = "habulaj/filterinstruct180" log.info("🚀 Carregando modelo e tokenizer...") tokenizer = AutoTokenizer.from_pretrained( model_name, use_fast=True, padding_side="left" ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoPeftModelForCausalLM.from_pretrained( model_name, device_map="cpu", torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, use_cache=True, trust_remote_code=True ) model.eval() log.info("✅ Modelo carregado (eval mode).") generation_config = GenerationConfig( max_new_tokens=128, temperature=1.0, do_sample=False, num_beams=1, use_cache=True, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.eos_token_id, no_repeat_ngram_size=2, repetition_penalty=1.1, length_penalty=1.0 ) # -------- FASTAPI INIT -------- app = FastAPI(title="News Filter JSON API") @app.get("/") def read_root(): return {"message": "News Filter JSON API is running!", "docs": "/docs"} # -------- INFERÊNCIA -------- def infer_filter(title, content): log.info(f"🧠 Inferência iniciada para: {title}") start_time = time.time() chat_prompt = build_chat_prompt(title, content) inputs = tokenizer( chat_prompt, return_tensors="pt", truncation=True, max_length=512, padding=False, add_special_tokens=False ) input_ids = inputs.input_ids attention_mask = inputs.attention_mask with torch.no_grad(), torch.inference_mode(): outputs = model.generate( input_ids=input_ids, attention_mask=attention_mask, generation_config=generation_config, num_return_sequences=1, output_scores=False, return_dict_in_generate=False ) generated_tokens = outputs[0][len(input_ids[0]):] generated = tokenizer.decode( generated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=True ) log.info("📤 Resultado gerado:") log.info(generated) json_result = extract_json(generated) duration = time.time() - start_time log.info(f"✅ JSON extraído em {duration:.2f}s") # Limpeza de memória del outputs, generated_tokens, inputs gc.collect() if json_result: return json_result else: raise HTTPException(status_code=404, detail="Unable to extract JSON from model output.") def build_chat_prompt(title: str, content: str) -> str: return f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|> 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}"<|eot_id|><|start_header_id|>assistant<|end_header_id|>""" def extract_json(text): match = re.search(r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', text, re.DOTALL) if match: json_text = match.group(0) # Conversões comuns json_text = re.sub(r"'", '"', json_text) json_text = re.sub(r'\bTrue\b', 'true', json_text) json_text = re.sub(r'\bFalse\b', 'false', json_text) json_text = re.sub(r",\s*}", "}", json_text) json_text = re.sub(r",\s*]", "]", json_text) return json_text.strip() return text # -------- API ROUTE -------- @app.get("/filter") def get_filter( title: str = Query(..., description="News title"), content: str = Query(..., description="News content") ): try: json_output = infer_filter(title, content) import json try: parsed = json.loads(json_output) return {"result": parsed} except json.JSONDecodeError as e: log.error(f"❌ Erro ao parsear JSON: {e}") return {"result": json_output, "warning": "JSON returned as string due to parsing error"} except HTTPException as e: raise e except Exception as e: log.exception("❌ Erro inesperado:") raise HTTPException(status_code=500, detail="Internal server error during inference.") @app.on_event("startup") async def warmup(): log.info("🔥 Executando warmup...") try: infer_filter("Test title", "Test content") log.info("✅ Warmup concluído.") except Exception as e: log.warning(f"⚠️ Warmup falhou: {e}")