saarus72 commited on
Commit
883f1f9
1 Parent(s): 6579fd9

Chat with normalizer

Browse files
Files changed (2) hide show
  1. app.py +90 -0
  2. examples.json +11 -0
app.py ADDED
@@ -0,0 +1,90 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import json
3
+ import re
4
+ import torch
5
+ from transformers import GPT2Tokenizer, T5ForConditionalGeneration
6
+
7
+
8
+ # re_tokens = re.compile(r"[а-яА-Я]+\s*|\d+(?:\.\d+)?\s*|[^а-яА-Я\d\s]+\s*")
9
+ re_tokens = re.compile(r"(?:[.,!?]|[а-яА-Я]\S*|\d\S*(?:\.\d+)?|[^а-яА-Я\d\s]+)\s*")
10
+
11
+
12
+ def tokenize(text):
13
+ return re.findall(re_tokens, text)
14
+
15
+
16
+ def strip_numbers(s):
17
+ """
18
+ From `1234567` to `1 234 567`
19
+ """
20
+ result = []
21
+ for part in s.split():
22
+ if part.isdigit():
23
+ while len(part) > 3:
24
+ result.append(part[:- 3 * ((len(part) - 1) // 3)])
25
+ part = part[- 3 * ((len(part) - 1) // 3):]
26
+ if part:
27
+ result.append(part)
28
+ else:
29
+ result.append(part)
30
+ return " ".join(result)
31
+
32
+
33
+ def construct_prompt(text):
34
+ """
35
+ From `я купил iphone 12X за 142 990 руб без 3-x часов 12:00, и т.д.` \
36
+ to `<SC1>я купил [iphone 12X]<extra_id_0> за [142 990]<extra_id_1> руб без [3-x]<extra_id_2> часов [12:00]<extra_id_3>, и т.д.`.
37
+ """
38
+ result = "<SC1>"
39
+ etid = 0
40
+ token_to_add = ""
41
+ for token in tokenize(text) + [""]:
42
+ if not re.search("[a-zA-Z\d]", token):
43
+ if token_to_add:
44
+ end_match = re.search(r"(.+?)(\W*)$", token_to_add, re.M).groups()
45
+ result += f"[{strip_numbers(end_match[0])}]<extra_id_{etid}>{end_match[1]}"
46
+ etid += 1
47
+ token_to_add = ""
48
+ result += token
49
+ else:
50
+ token_to_add += token
51
+ return result
52
+
53
+
54
+ def construct_answer(prompt:str, prediction:str) -> str:
55
+ re_prompt = re.compile(r"\[([^\]]+)\]<extra_id_(\d+)>")
56
+ re_pred = re.compile(r"\<extra_id_(\d+)\>(.+?)(?=\<extra_id_\d+\>|</s>)")
57
+ pred_data = {}
58
+ for match in re.finditer(re_pred, prediction.replace("\n", " ")):
59
+ pred_data[match[1]] = match[2].strip()
60
+ while match := re.search(re_prompt, prompt):
61
+ replace = pred_data.get(match[2], match[1])
62
+ prompt = prompt[:match.span()[0]] + replace + prompt[match.span()[1]:]
63
+ return prompt.replace("<SC1>", "")
64
+
65
+
66
+ with open("examples.json") as f:
67
+ test_examples = json.load(f)
68
+
69
+
70
+ tokenizer = GPT2Tokenizer.from_pretrained("saarus72/russian_text_normalizer", eos_token='</s>')
71
+ model = T5ForConditionalGeneration.from_pretrained("saarus72/russian_text_normalizer")
72
+
73
+
74
+ def predict(text):
75
+ input_ids = torch.tensor([tokenizer.encode(text)])
76
+ outputs = model.generate(input_ids, max_new_tokens=50, eos_token_id=tokenizer.eos_token_id, early_stopping=True)
77
+ return tokenizer.decode(outputs[0][1:])
78
+
79
+
80
+ def norm(message, history):
81
+ prompt = construct_prompt(message)
82
+ yield f"```Prompt:\n{prompt}\nPrediction:\n...```\n..."
83
+ prediction = predict(prompt)
84
+ answer = construct_answer(prompt, prediction)
85
+ # yield f"```\nPrompt:\n{prompt}\nPrediction:\n{prediction}\n```\n{answer}"
86
+ yield f"Prompt:\n```{prompt}```\nPrediction:\n```\n{prediction}\n```\n{answer}"
87
+
88
+
89
+ demo = gr.ChatInterface(fn=norm, stop_btn=None, examples=list(test_examples.keys())).queue()
90
+ demo.launch()
examples.json ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "я купил iphone 12X за 142 990 руб без 3-x часов полдень и т.д.": "я купил айфон двенадцать икс за сто сорок две тысячи девятьсот девяносто руб без трёх часов полдень и т.д.",
3
+ "я купил айфон за 14 970 рублей": "я купил айфон за четырнадцать тысяч девятьсот семьдесят рублей",
4
+ "Временами я думаю, какое применение найти тем 14 697 рублям, что лежат уже больше 33 лет?": "Временами я думаю, какое применение найти тем четырнадцати тысячам шестистам девяносто семи рублям, что лежат уже больше тридцати трёх лет?",
5
+ "Было у отца 3 сына, но не было даже 2-3 пиджаков с блёстками за 142 990 рублей.": "Было у отца три сына, но не было даже двух-трёх пиджаков с блёстками за сто сорок две тысячи девятьсто девяносто рублей.",
6
+ "В школе у меня одни 5.": "В школе у меня одни пятёрки.",
7
+ "Было у отца 3 сына. Старшему было 35, среднему - не меньше 33, а младший на 4 младше всех. Бывает.": "Было у отца три сына. Старшему было тридцать пять, среднему - не меньше тридцати трех, а младший на четыре младше всех. Бывает.",
8
+ "я вырос на the beatles, меня не испугать даже 33 yellow submarine": "я вырос на зе битлз, меня не испугать даже тридцатью тремя йеллоу сабмэрин",
9
+ "слыш nigga ты слыхал про gitdata?": "слыш нигга ты слыхал про гитдата?",
10
+ "стоимость samsung 32MX Pro — всего 189 600 руб!": "стоимость самсунг тридцать два эм икс про — всего сто восемьдесят девять тысяч шестьсот руб!"
11
+ }