# 추론
import torch
from transformers import RobertaForSequenceClassification, RobertaTokenizer

# 전이학습에 사용한 토크나이저와 모델 로드 & 가중치 로드


class ToxcitiyClassifier:
    def __init__(self):
        self.tokenizer = RobertaTokenizer.from_pretrained("beomi/KcBERT-v2023")
        self.model = RobertaForSequenceClassification.from_pretrained(
            "beomi/KcBERT-v2023", num_labels=2
        )
        self.model.load_state_dict(
            torch.load(
                "./src/ai/pytorchmodel_518망언분류_acc8583.bin",
                map_location=torch.device("cpu"),
                weights_only=True,
            )
        )
        # 모델을 평가 모드로 설정
        self.model.eval()
        # 입력 텍스트 예시
        self.class_labels = ["문제없음/관련없음", "부적절(518 망언 가능)"]

    def infer(self, new_text):
        inputs = self.tokenizer(new_text, return_tensors="pt")
        # 추론 수행 (CPU 사용)
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
        probs = torch.nn.functional.softmax(logits, dim=-1)
        toxic_prob = probs[0][1].item()
        non_toxic_prob = probs[0][0].item()

        print(
            f"{self.class_labels[0]}:{non_toxic_prob*100:.2f}%, {self.class_labels[1]}:{toxic_prob*100:.2f}%"
        )

        return toxic_prob


if __name__ == "__main__":
    classifier = ToxcitiyClassifier()
    result = classifier.infer("5.18 은 하나의 폭동이야")
    print(result)