1import torch
2from transformers.models.bert import BertTokenizer, BertForSequenceClassification
3from transformers import AutoTokenizer, AutoModelForSequenceClassification
4
5
6# Load model architecture from COLD and load fine-tuned params.
7model_name = "thu-coai/roberta-base-cold"
8tokenizer = AutoTokenizer.from_pretrained(model_name)
9model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
10model_path = "finetuned_cold_LoL.pth" # Could be downloaded in this repo.
11model.load_state_dict(torch.load(model_path))
12
13
14# Demo for toxicity detection
15texts = ['狠狠地导', '卡了哟', 'gala有卡莎皮肤，你们这些小黑子有吗？', '早改了，改成回血了']
16model_input = tokenizer(texts, return_tensors="pt", padding=True)
17model_output = model(**model_input, return_dict=False)
18prediction = torch.argmax(model_output[0].cpu(), dim=-1)
19prediction = [p.item() for p in prediction]
20# prediction = [1, 0, 1, 0] # 1 for toxic, 0 for non-toxic