Views
No views yet
1import torch
2from transformers.models.bert import BertTokenizer, BertForSequenceClassification
3
4tokenizer = BertTokenizer.from_pretrained('thu-coai/roberta-base-cold')
5model = BertForSequenceClassification.from_pretrained('thu-coai/roberta-base-cold')
6model.eval()
7
8texts = ['你就是个傻逼!','黑人很多都好吃懒做,偷奸耍滑!','男女平等,黑人也很优秀。']
9
10model_input = tokenizer(texts,return_tensors="pt",padding=True)
11model_output = model(**model_input, return_dict=False)
12prediction = torch.argmax(model_output[0].cpu(), dim=-1)
13prediction = [p.item() for p in prediction]
14print(prediction) # --> [1, 1, 0] (0 for Non-Offensive, 1 for Offenisve)@article{deng2022cold,
title={Cold: A benchmark for chinese offensive language detection},
author={Deng, Jiawen and Zhou, Jingyan and Sun, Hao and Zheng, Chujie and Mi, Fei and Meng, Helen and Huang, Minlie},
booktitle={Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing},
year={2022}
}