app.py

import os
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
from transformers import pipeline

MODEL_NAME = "will702/indo-roBERTa-financial-sentiment-v2"
API_KEY = os.getenv("API_KEY")

# Label mapping — flipped: 0=Positive, 1=Neutral, 2=Negative
LABEL_MAP = {
    "label_0": "positive",
    "label_1": "neutral",
    "label_2": "negative",
    "positive": "positive",
    "neutral": "neutral",
    "negative": "negative",
}

classifier = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global classifier
    print(f"Loading model: {MODEL_NAME}")
    classifier = pipeline("text-classification", model=MODEL_NAME)
    print("Model loaded.")
    yield


app = FastAPI(title="StockPro Sentiment", lifespan=lifespan)


class PredictRequest(BaseModel):
    texts: list[str]


@app.post("/predict")
async def predict(body: PredictRequest, request: Request):
    if API_KEY:
        key = request.headers.get("X-API-Key")
        if key != API_KEY:
            raise HTTPException(status_code=401, detail="Invalid API key")

    texts = body.texts
    if not texts:
        raise HTTPException(status_code=400, detail="texts must not be empty")
    if len(texts) > 20:
        raise HTTPException(status_code=400, detail="Maximum 20 texts per request")

    if classifier is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    predictions = classifier(texts, truncation=True, max_length=512)

    results = []
    for text, pred in zip(texts, predictions):
        label = LABEL_MAP.get(pred["label"].lower(), "neutral")
        results.append({
            "text": text,
            "sentiment": label,
            "score": round(pred["score"], 4),
        })

    return {"results": results}


@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": classifier is not None}