# 추론
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망언분류_acc9304.bin", map_location=torch.device('cpu')))
# 모델을 평가 모드로 설정
model.eval()

# 입력 텍스트 예시
class_labels = ["문제없음/관련없음", "부적절(518 망언 가능)"]
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]
    unpredicted_label = class_labels[1-predicted_class]
    probability = probs[0][predicted_class].item()
    return f"{predicted_label}:{probability*100:.2f}%, {unpredicted_label}:{((1-probability)*100):.2f}%"

# Streamlit interface
st.markdown('### 5·18 민주화운동 관련 부적절한 발언 판단(베타버전)')
st.markdown('<small style="color:grey;">5·18 민주화운동과 관련해 무장 폭동, 북한군 개입, 가짜 유공자 등 부적절한 언급과 지역-이념에 대한 혐오성 발언이 문제되고 있습니다. 아래에 문장을 입력하면 이러한 내용을 중심으로 "문제없음/관련없음" 또는 "부적절(518 망언 가능)"로 판단해 드립니다. 예측 모델의 정확도는 93.04%로, 일부 부정확한 결과가 나올 수 있습니다 </small>', unsafe_allow_html=True)
user_input = st.text_area("이 곳에 문장 입력(100자 이하 권장, 너무 길면 분석 불가):")
if st.button('시작'):
    result = inference(user_input)
    st.write(result)