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}