File size: 1,371 Bytes
f5510ed
 
 
 
 
 
 
 
 
 
 
99a4945
f5510ed
 
 
 
 
 
 
 
 
99a4945
f5510ed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f46a31f
 
 
 
f5510ed
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
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()