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

checkpoint = "mr4/bert-base-jp-sentiment-analysis"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint)

def classify_comment(comment):
    raw_inputs = [comment]
    inputs = tokenizer(raw_inputs, padding=True,
                    truncation=True, return_tensors="pt")
    outputs = model(**inputs)
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    print(predictions)
    results = []
    for i, prediction in enumerate(predictions):
        for j, value in enumerate(prediction):
            results.append({
                'label': model.config.id2label[j],
                'score': value.item(),
            })
    print(results)
    sorted_data = sorted(results, key=lambda x: x['score'])
    return f"{"❌" if sorted_data[0]['label'].lower() in ["negative"] else '✅'} {comment}\n{sorted_data[0]['label']} ({sorted_data[0]['score']*100:.5f}%)\n{sorted_data[1]['label']} ({sorted_data[1]['score']*100:.5f}%)"

iface = gr.Interface(
    fn=classify_comment,
    inputs=gr.Textbox(lines=3, placeholder="コメントを入力してください"),
    outputs="text",
    title="日本の感情分析",
    description="日本の感情分析"
)

iface.launch()