hariharan220's picture
Create app.py
b4c83a0 verified
raw
history blame
976 Bytes
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}