from fastapi import FastAPI
from transformers import BertForSequenceClassification, AutoTokenizer
import torch

# ✅ Load trained FinBERT model
MODEL_PATH = "hariharan220/finbert-stock-sentiment"
model = BertForSequenceClassification.from_pretrained(MODEL_PATH)
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

# ✅ Define sentiment labels
labels = ["Negative", "Neutral", "Positive"]

# ✅ Create FastAPI app
app = FastAPI()

@app.get("/")
async def home():
    return {"message": "Stock Sentiment Analysis API is running!"}

@app.post("/predict")
async def predict_sentiment(text: str):
    """Predicts sentiment of stock-related text using FinBERT"""
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    
    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    prediction = torch.argmax(logits, dim=1).item()
    sentiment = labels[prediction]

    return {"text": text, "sentiment": sentiment}