import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F

# Load Pretrained Model & Tokenizer
MODEL_NAME = "nlptown/bert-base-multilingual-uncased-sentiment"  # Pretrained sentiment model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# Function for sentiment analysis
def analyze_sentiment(user_input):
    if user_input.strip():
        # Tokenize and Convert Input to Tensor
        inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True)
        
        # Perform Sentiment Analysis
        with torch.no_grad():
            outputs = model(**inputs)
            scores = F.softmax(outputs.logits, dim=-1).squeeze().tolist()
        
        # Interpret Results
        labels = ["非常に否定的 😡", "否定的 😞", "中立的 😐", "肯定的 🙂", "非常に肯定的 😍"]
        sentiment = labels[scores.index(max(scores))]  # Select sentiment with highest score
        
        # Format Confidence Scores
        confidence = "\n".join([f"{label}: {score:.2%}" for label, score in zip(labels, scores)])
        
        return f"予測された感情: {sentiment}", confidence, sentiment  # Returning sentiment separately
    else:
        return "⚠️ Please enter text before analyzing.", "", ""

# Gradio Blocks interface
with gr.Blocks() as app:
    gr.Markdown("## 🔥 感情分析")
    gr.Markdown("分析するテキストを入力してください")
    
    # Text input
    user_input = gr.Textbox(label="分析するテキスト", info="シンガポールで開かれているアジア安全保障会議で演説したアメリカのヘグセス国防長官が台湾情勢などをめぐり中国を名指しして繰り返し非難したことを受け、中国外務省はアメリカ側に抗議したことを明らかにしました。")

    # Outputs
    sentiment_output = gr.Textbox(label="予測された感情", interactive=False)
    confidence_output = gr.Textbox(label="確信度", interactive=False)
    top_sentiment_output = gr.Textbox(label="最も高い感情", interactive=False)  # New output for top sentiment

    # Button to trigger analysis
    analyze_button = gr.Button("感情分析")

    # Connect button to function
    analyze_button.click(
        fn=analyze_sentiment,
        inputs=[user_input],
        outputs=[sentiment_output, confidence_output, top_sentiment_output]  # Include new output
    )

# Launch the Gradio app
app.launch()