Unggi commited on
Commit
f5510ed
1 Parent(s): c8dab0c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +49 -0
app.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pip
2
+ pip.main(['install', 'torch'])
3
+ pip.main(['install', 'transformers'])
4
+
5
+ import torch
6
+ import torch.nn as nn
7
+ import gradio as gr
8
+ import transformers
9
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
10
+
11
+ def load_model(model_name):
12
+ # model_name = "Unggi/hate_speech_bert"
13
+ # model
14
+ model = AutoModelForSequenceClassification.from_pretrained(model_name)
15
+ # tokenizer..
16
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
17
+
18
+ return model, tokenizer
19
+
20
+
21
+ def inference(prompt):
22
+ model_name = "Unggi/ko_hate_speech_KcELECTRA" #"Unggi/hate_speech_bert"
23
+
24
+ model, tokenizer = load_model(
25
+ model_name = model_name
26
+ )
27
+
28
+ inputs = tokenizer(
29
+ prompt,
30
+ return_tensors="pt"
31
+ )
32
+
33
+ with torch.no_grad():
34
+ logits = model(**inputs).logits
35
+
36
+ # for binary classification
37
+ sigmoid = nn.Sigmoid()
38
+ bi_prob = sigmoid(logits)
39
+
40
+ predicted_class_id = bi_prob.argmax().item()
41
+ class_id = model.config.id2label[predicted_class_id]
42
+
43
+ return "class_id: " + str(class_id) + "\n" + "clean_prob: " + str(bi_prob[0][0].item()) + "\n" + "unclean_prob: " + str(bi_prob[0][1].item())
44
+
45
+ demo = gr.Interface(
46
+ fn=inference,
47
+ inputs="text",
48
+ outputs="text", #return 값
49
+ ).launch()