prueba_api / main.py
RaulCtm's picture
Update main.py
0c0fc2e verified
Raw
History Blame Contribute Delete
1.39 kB
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
#Inicializamos la aplicación FastAPI
app = FastAPI(
title="ToxicFilterAI API",
description="Microservicio de inferencia para detección de toxicidad en comentarios",
version="1.0.0"
)
#Cargamos el modelo al arrancar el contenedor
print("⏳ Descargando y montando modelo desde Hugging Face Hub...")
clasificador = pipeline("text-classification", modelo="ToxicFilterAI/distilbert-toxicidad")
print("✅ Modelo cargado en memoria.")
#Definimos el contrato de datos (Lo que esperamos recibir)
class PeticionComentario(BaseModel):
texto: str
#Endpoint de Health Check (Buenas prácticas de infraestructura)
@app.get("/")
def health_check():
return {"status": "online", "message": "Motor de inferencia operativo"}
#Endpoint principal de predicción
@app.post("/api/v1/predict")
def predecir_toxicidad(peticion: PeticionComentario):
if not peticion.texto.strip():
raise HTTPException(status_code=400, detail="El texto no puede estar vacío")
resultado = clasificador(peticion.texto)[0]
return {
"texto_analizado": peticion.texto,
"etiqueta_modelo": resultado['label'],
"score_confianza": round(float(resultado['score']), 4),
"es_toxico": True if resultado['label'] in ["Toxic", "NEGATIVE"] else False
}