Views
No views yet
1import torch
2from transformers import BertTokenizer, BertForSequenceClassification
3
4tokenizer = BertTokenizer.from_pretrained('thu-coai/roberta-zh-specific')
5model = BertForSequenceClassification.from_pretrained('thu-coai/roberta-zh-specific', num_labels=2)
6model.eva()
7
8context = [
9 "你大爱的冷门古诗词是什么?\t一枝红艳露凝香,云雨巫山枉断肠",
10 "你大爱的冷门古诗词是什么?\t一枝红艳露凝香,云雨巫山枉断肠",
11]
12
13response = [
14 "我也很喜欢,我觉得这句的意境很美",
15 "我也很喜欢",
16]
17
18model_input = tokenizer(context, response, return_tensors='pt', padding=True)
19with torch.no_grad():
20 model_output = model(**model_input, return_dict=True)
21logits = model_output.logits
22preds_all = torch.argmax(logits, dim=-1).cpu()
23print(preds_all) # 1 for specific response else 0
24
25