import gradio as gr
import torch
from transformers import BertTokenizer, BertForSequenceClassification

# Load pre-trained model tokenizer (vocabulary)
tokenizer = BertTokenizer.from_pretrained('ProsusAI/finbert')

# Load pre-trained model
model = BertForSequenceClassification.from_pretrained('ProsusAI/finbert')

def get_sentiment(sec_text):  # Ensure the parameter name matches the placeholder name.
    # Encode the text
    tokens = tokenizer.encode_plus(sec_text, add_special_tokens=True, return_tensors="pt")
    
    # Make prediction
    with torch.no_grad():
        outputs = model(**tokens)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)

    # Convert predictions to sentiment labels
    labels = ['Positive', 'Neutral', 'Negative']
    sentiment = labels[torch.argmax(predictions)]
    
    # Return the sentiment analysis result
    return f"{sentiment} Sentiment"

# Custom CSS to center the title
custom_css = """
    .title {
        text-align: center;
    }
"""

# Define the Gradio interface
gr_interface = gr.Interface(
    fn=get_sentiment,
    inputs=gr.Textbox(lines=1, placeholder=""),
    outputs="text",
    title="Sentiment Analysis",
    css=custom_css  # Add the custom CSS to the Interface
)

# Launch the interface
gr_interface.launch()