MarIA / app.py
ydyue's picture
Upload 10 files
3a86801
raw
history blame contribute delete
693 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
# Cargar modelo desde carpeta local
pipe = pipeline("text-classification", model="./MarIA", tokenizer="./MarIA")
label_map = {
"0": "Negativo",
"1": "Neutral",
"2": "Positivo"
}
# FastAPI app
app = FastAPI()
# Entrada de la API
class TextInput(BaseModel):
text: str
@app.post("/predict")
def predict(input: TextInput):
result = pipe(input.text)[0]
label_id = result["label"]
label_num = label_id.split("_")[-1] if "_" in label_id else label_id
sentiment = label_map.get(label_num, "Desconocido")
return {"sentiment": sentiment}