Views
No views yet
1from transformers import AutoModelForTokenClassification, AutoTokenizer
2import torch
3
4device = torch.device("cuda:0")
5path = f"KomeijiForce/xlm-roberta-large-metaie"
6tokenizer = AutoTokenizer.from_pretrained(path)
7tagger = AutoModelForTokenClassification.from_pretrained(path).to(device)
8
9def find_sequences(lst):
10 sequences = []
11 i = 0
12 while i < len(lst):
13 if lst[i] == 0:
14 start = i
15 end = i
16 i += 1
17 while i < len(lst) and lst[i] == 1:
18 end = i
19 i += 1
20 sequences.append((start, end+1))
21 else:
22 i += 1
23 return sequences
24
25examples = [
26 "Fire volleys at the command happens: The soldiers were expected to fire volleys at the command of officers, but in practice this happened only in the first minutes of the battle .",
27 "Historische Ereignisse: Siebenjährigen Krieg von 1756 bis 1763, war Preußen als fünfte Großmacht neben Frankreich, Großbritannien, Österreich und Russland in der europäischen Pentarchie anerkannt .",
28 "高度: 东方明珠自落成后便为上海天际线的组成部分之一，总高468米。",
29 "倒れた場所: カフカは高松の私立図書館に通うようになるが、ある日目覚めると、自分が森の中で血だらけで倒れていた。",
30]
31
32for example in examples:
33 inputs = tokenizer(example, return_tensors="pt").to(device)
34 tag_predictions = tagger(**inputs).logits[0].argmax(-1)
35
36 predictions = [tokenizer.decode(inputs.input_ids[0, seq[0]:seq[1]]).strip() for seq in find_sequences(tag_predictions)]
37
38 print(example)
39 print(predictions)1Fire volleys at the command happens: The soldiers were expected to fire volleys at the command of officers, but in practice this happened only in the first minutes of the battle .
2['first minutes of the battle']
3Historische Ereignisse: Siebenjährigen Krieg von 1756 bis 1763, war Preußen als fünfte Großmacht neben Frankreich, Großbritannien, Österreich und Russland in der europäischen Pentarchie anerkannt .
4['Siebenjährigen Krieg']
5高度: 东方明珠自落成后便为上海天际线的组成部分之一，总高468米。
6['468米']
7倒れた場所: カフカは高松の私立図書館に通うようになるが、ある日目覚めると、自分が森の中で血だらけで倒れていた。
8['森']