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

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

def load_model(model_name):
    # model_name = "Unggi/hate_speech_bert"
    # 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_bert"

    model, tokenizer = load_model(
        model_name = model_name
    )

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

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

    predicted_class_id = logits.argmax().item()
    class_id = model.config.id2label[predicted_class_id]

    return class_id

demo = gr.Interface(
    fn=inference, 
    inputs="text", 
    outputs="text", #return 값
    examples=[
        "아무 방 눌렀다가 미친 게임 봤다"
    ]
    ).launch() # launch(share=True)를 설정하면 외부에서 접속 가능한 링크가 생성됨

demo.launch()