import torch
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# -----------------------------
# CONFIG
# -----------------------------

MODEL_PATH = "./finbert_financial_model"

id2label = {0: "negative", 1: "neutral", 2: "positive"}

# -----------------------------
# LOAD MODEL
# -----------------------------

@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
    model.eval()
    return tokenizer, model

tokenizer, model = load_model()

# -----------------------------
# PREDICTION FUNCTION
# -----------------------------

def predict(text):
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=128
    )

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=1)
        pred_class = torch.argmax(probs, dim=1).item()

    return id2label[pred_class], probs[0].tolist()

# -----------------------------
# UI
# -----------------------------

st.set_page_config(page_title="Financial Sentiment Analyzer", layout="centered")

st.title("💰 Financial Sentiment Analyzer")
st.write("Analyze sentiment of financial text using your fine-tuned FinBERT model.")

user_input = st.text_area("Enter financial text:", height=150)

if st.button("Analyze Sentiment"):
    if user_input.strip() == "":
        st.warning("Please enter some text.")
    else:
        label, probs = predict(user_input)

        st.subheader("Prediction:")
        st.success(f"Sentiment: **{label.upper()}**")

        st.subheader("Confidence Scores:")
        st.write(f"Negative: {probs[0]:.4f}")
        st.write(f"Neutral:  {probs[1]:.4f}")
        st.write(f"Positive: {probs[2]:.4f}")