Views
No views yet
1from transformers import AutoTokenizer, AutoModelForTokenClassification
2
3PRETRAINED = "cbdb/PersonAndBookTitleSplitter"
4tokenizer = AutoTokenizer.from_pretrained(PRETRAINED)
5model = AutoModelForTokenClassification.from_pretrained(PRETRAINED)1# Load your data here
2test_list = ['徐元文漢魏風致集', '熊方補後漢書年表', '羅振玉本朝學術源流槪略 一卷', '陶諧陶莊敏集']1def predict_class(test):
2 tokens_test = tokenizer.encode_plus(
3 test,
4 add_special_tokens=True,
5 return_attention_mask=True,
6 padding=True,
7 max_length=128,
8 return_tensors='pt',
9 truncation=True
10 )
11
12 test_seq = torch.tensor(tokens_test['input_ids'])
13 test_mask = torch.tensor(tokens_test['attention_mask'])
14
15 inputs = {
16 "input_ids": test_seq,
17 "attention_mask": test_mask
18 }
19 with torch.no_grad():
20 # print(inputs.shape)
21 outputs = model(**inputs)
22 outputs = outputs.logits.detach().cpu().numpy()
23
24 softmax_score = softmax(outputs)
25 softmax_score = np.argmax(softmax_score, axis=2)[0]
26 return test_seq, softmax_score
27
28for test_sen0 in test_list:
29 test_seq, pred_class_proba = predict_class(test_sen0)
30 test_sen = tokenizer.decode(test_seq[0]).split()
31 label = [idx2label[i] for i in pred_class_proba]
32
33 element_to_find = '。'
34
35 if element_to_find in label:
36 index = label.index(element_to_find)
37 test_sen_pred = [i for i in test_sen0]
38 test_sen_pred.insert(index, element_to_find)
39 test_sen_pred = ''.join(test_sen_pred)
40
41 else:
42 test_sen_pred = [i for i in test_sen0]
43 test_sen_pred = ''.join(test_sen_pred)
44
45 print(test_sen_pred)