Views
No views yet
1from transformers import pipeline
2from transformers import AutoTokenizer
3
4pt = "RUPunct/RUPunct_big"
5
6tk = AutoTokenizer.from_pretrained(pt, strip_accents=False, add_prefix_space=True)
7classifier = pipeline("ner", model=pt, tokenizer=tk, aggregation_strategy="first")
8
9
10def process_token(token, label):
11 if label == "LOWER_O":
12 return token
13 if label == "LOWER_PERIOD":
14 return token + "."
15 if label == "LOWER_COMMA":
16 return token + ","
17 if label == "LOWER_QUESTION":
18 return token + "?"
19 if label == "LOWER_TIRE":
20 return token + "—"
21 if label == "LOWER_DVOETOCHIE":
22 return token + ":"
23 if label == "LOWER_VOSKL":
24 return token + "!"
25 if label == "LOWER_PERIODCOMMA":
26 return token + ";"
27 if label == "LOWER_DEFIS":
28 return token + "-"
29 if label == "LOWER_MNOGOTOCHIE":
30 return token + "..."
31 if label == "LOWER_QUESTIONVOSKL":
32 return token + "?!"
33 if label == "UPPER_O":
34 return token.capitalize()
35 if label == "UPPER_PERIOD":
36 return token.capitalize() + "."
37 if label == "UPPER_COMMA":
38 return token.capitalize() + ","
39 if label == "UPPER_QUESTION":
40 return token.capitalize() + "?"
41 if label == "UPPER_TIRE":
42 return token.capitalize() + " —"
43 if label == "UPPER_DVOETOCHIE":
44 return token.capitalize() + ":"
45 if label == "UPPER_VOSKL":
46 return token.capitalize() + "!"
47 if label == "UPPER_PERIODCOMMA":
48 return token.capitalize() + ";"
49 if label == "UPPER_DEFIS":
50 return token.capitalize() + "-"
51 if label == "UPPER_MNOGOTOCHIE":
52 return token.capitalize() + "..."
53 if label == "UPPER_QUESTIONVOSKL":
54 return token.capitalize() + "?!"
55 if label == "UPPER_TOTAL_O":
56 return token.upper()
57 if label == "UPPER_TOTAL_PERIOD":
58 return token.upper() + "."
59 if label == "UPPER_TOTAL_COMMA":
60 return token.upper() + ","
61 if label == "UPPER_TOTAL_QUESTION":
62 return token.upper() + "?"
63 if label == "UPPER_TOTAL_TIRE":
64 return token.upper() + " —"
65 if label == "UPPER_TOTAL_DVOETOCHIE":
66 return token.upper() + ":"
67 if label == "UPPER_TOTAL_VOSKL":
68 return token.upper() + "!"
69 if label == "UPPER_TOTAL_PERIODCOMMA":
70 return token.upper() + ";"
71 if label == "UPPER_TOTAL_DEFIS":
72 return token.upper() + "-"
73 if label == "UPPER_TOTAL_MNOGOTOCHIE":
74 return token.upper() + "..."
75 if label == "UPPER_TOTAL_QUESTIONVOSKL":
76 return token.upper() + "?!"
77
78while 1:
79 input_text = input(":> ")
80 preds = classifier(input_text)
81 output = ""
82 for item in preds:
83 output += " " + process_token(item['word'].strip(), item['entity_group'])
84 print(">>>", output)