A5hbr1ng3r/COLD-LoL

import torch
from transformers.models.bert import BertTokenizer, BertForSequenceClassification
from transformers import AutoTokenizer, AutoModelForSequenceClassification


# Load model architecture from COLD and load fine-tuned params.
model_name = "thu-coai/roberta-base-cold"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
model_path = "finetuned_cold_LoL.pth" # Could be downloaded in this repo.
model.load_state_dict(torch.load(model_path))


# Demo for toxicity detection
texts = ['狠狠地导', '卡了哟', 'gala有卡莎皮肤，你们这些小黑子有吗？', '早改了，改成回血了']
model_input = tokenizer(texts, return_tensors="pt", padding=True)
model_output = model(**model_input, return_dict=False)
prediction = torch.argmax(model_output[0].cpu(), dim=-1)
prediction = [p.item() for p in prediction]
# prediction = [1, 0, 1, 0] # 1 for toxic, 0 for non-toxic