Views
No views yet
1from transformers import TextClassificationPipeline, BertForSequenceClassification, AutoTokenizer+
2
3model_name = 'SJ-Donald/kcbert-large-unsmile'
4model = BertForSequenceClassification.from_pretrained(model_name)
5
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7
8pipe = TextClassificationPipeline(
9 model = model,
10 tokenizer = tokenizer,
11 device = 0, # cpu: -1, gpu: gpu number
12 return_all_scores = True,
13 function_to_apply = 'sigmoid'
14)
15
16for result in pipe("이래서 여자는 게임을 하면 안된다")[0]:
17 print(result)
18
19{'label': '여성/가족', 'score': 0.9793611168861389}
20{'label': '남성', 'score': 0.006330598145723343}
21{'label': '성소수자', 'score': 0.007870828732848167}
22{'label': '인종/국적', 'score': 0.010810344479978085}
23{'label': '연령', 'score': 0.020540334284305573}
24{'label': '지역', 'score': 0.015790466219186783}
25{'label': '종교', 'score': 0.014563685283064842}
26{'label': '기타 혐오', 'score': 0.04097242280840874}
27{'label': '악플/욕설', 'score': 0.019168635830283165}
28{'label': 'clean', 'score': 0.014866289682686329}