import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_name = "yiyanghkust/finbert-tone"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

if torch.backends.mps.is_available():
    device = "mps"
elif torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"
print(f"Using device => {device}")
model.to(device)

# FinBERT 的 label 通常是這三種
label_names = ["positive", "negative", "neutral"]

def predict_sentiment(text):
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding="max_length",
        max_length=128  
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    with torch.no_grad():
        outputs = model(**inputs)
        pred_label = outputs.logits.argmax(dim=-1).item()

    return label_names[pred_label]

demo_description = """
    **This Space uses the FinBERT model for 3-class financial sentiment classification (positive, negative, neutral).**Simply input a financial news headline or sentence to see its sentiment classification.    
    
    **How to Use**:
    1. Enter text: Type or paste a financial news headline (or any short text) into the text box.
    2. Submit: Click the Submit button.
    3. View result: The predicted sentiment label—negative, neutral, or positive
    
    **Sample Questions**:
    1. The 2015 target for net sales has been set at GBP 1bn and the target for return on investment at over 20 % .
    2. The agreement was signed with Biohit Healthcare Ltd , the UK-based subsidiary of Biohit Oyj , a Finnish public company which develops , manufactures and markets liquid handling products and diagnostic test systems .
    """

demo = gr.Interface(
    fn=predict_sentiment,
    inputs="text",
    outputs="text",
    title="FinBERT Financial News Headline Sentiment Demo",
    description=demo_description,
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()