# 추론
from transformers import RobertaTokenizer, RobertaForSequenceClassification
import torch
import streamlit as st

# 전이학습에 사용한 토크나이저와 모델 로드 & 가중치 로드
tokenizer = RobertaTokenizer.from_pretrained('beomi/KcBERT-v2023')
model = RobertaForSequenceClassification.from_pretrained('beomi/KcBERT-v2023', num_labels=2)
model.load_state_dict(torch.load("pytorchmodel_518망언분류_acc9308.bin"))
# 모델을 평가 모드로 설정
model.eval()

# 입력 텍스트 예시
class_labels = ["적절(518망언_NO)", "부적절(518망언_YES)"]
def inference(new_text):
    inputs = tokenizer(new_text, return_tensors="pt")
    # 추론 수행 (CPU 사용)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=-1)
    predicted_class = torch.argmax(probs, dim=1).item()
    predicted_label = class_labels[predicted_class]
    probability = probs[0][predicted_class].item()
    return f"예측: {predicted_label}, 확률: {probability:.4f}"

# Streamlit interface
st.title('5·18 민주화운동 관련 부적절한 발언 탐지')
st.markdown('<small style="color:grey;">5·18 민주화운동과 관련해 무장 폭동, 북한군 개입, 가짜 유공자 등 부적절한 언급과 지역-이념에 대한 혐오성 발언이 문제되고 있습니다. 아래에 문장을 입력하면 이러한 내용을 중심으로 문장의 부적절성 여부를 확률과 함께 판단해 드립니다. </small>', unsafe_allow_html=True)
user_input = st.text_area("이 곳에 문장 입력(100자 이하 권장):")
if st.button('시작'):
    result = inference(user_input)
    st.write(result)