import streamlit as st
import torch
from transformers import BertConfig, BertForSequenceClassification, BertTokenizer
import numpy as np


# Load the model and tokenizer
def load_model():
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=7)
    model_state_dict = torch.load('sentiment7_model_acc8878.pth', map_location=torch.device('cpu')) # cpu 사용
    model.load_state_dict(model_state_dict)
    model.eval()
    return model, tokenizer

model, tokenizer = load_model()

# Define the inference function
def inference(input_doc):
    inputs = tokenizer(input_doc, return_tensors='pt')
    outputs = model(**inputs)
    probs = torch.softmax(outputs.logits, dim=1).squeeze().tolist()
    class_idx = {'공포': 0, '놀람': 1, '분노': 2, '슬픔': 3, '중립': 4, '행복': 5, '혐오': 6}
    return {class_name: prob for class_name, prob in zip(class_idx, probs)}

# Set up the Streamlit interface
st.title('Sentiment Analysis with BERT')
user_input = st.text_area("Enter text here:")
if st.button('Analyze'):
    result = inference(user_input)
    st.write(result)