import pip
pip.main(['install', 'torch'])
pip.main(['install', 'transformers'])

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

def load_model(model_name):
    # model_name = "Unggi/hate_speech_classifier_KcElectra"
    # model
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    # tokenizer..
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    return model, tokenizer


def inference(prompt):
    model_name = "Unggi/hate_speech_classifier_KcElectra" #"Unggi/hate_speech_bert"
    
    model, tokenizer = load_model(
        model_name = model_name
    )

    inputs = tokenizer(
        prompt, 
        return_tensors="pt"
        )

    with torch.no_grad():
        logits = model(**inputs).logits

    # for binary classification
    sigmoid = nn.Sigmoid()
    bi_prob = sigmoid(logits)

    predicted_class_id = bi_prob.argmax().item()

    # 임시
    model.config.id2label = {0: 'CLEAN', 1: 'UNCLEAN'}
    
    class_id = model.config.id2label[predicted_class_id]

    return "class_id: " + str(class_id) + "\n" + "clean_prob: " + str(bi_prob[0][0].item()) + "\n" + "unclean_prob: " + str(bi_prob[0][1].item())

demo = gr.Interface(
    fn=inference, 
    inputs="text", 
    outputs="text", #return 값
    ).launch()