Chat with normalizer

app.py

@@ -0,0 +1,90 @@
+import gradio as gr
+import json
+import re
+import torch
+from transformers import GPT2Tokenizer, T5ForConditionalGeneration 
+
+
+# re_tokens = re.compile(r"[а-яА-Я]+\s*|\d+(?:\.\d+)?\s*|[^а-яА-Я\d\s]+\s*")
+re_tokens = re.compile(r"(?:[.,!?]|[а-яА-Я]\S*|\d\S*(?:\.\d+)?|[^а-яА-Я\d\s]+)\s*")
+
+
+def tokenize(text):
+    return re.findall(re_tokens, text)
+
+
+def strip_numbers(s):
+    """
+    From `1234567` to `1 234 567`
+    """
+    result = []
+    for part in s.split():
+        if part.isdigit():
+            while len(part) > 3:
+                result.append(part[:- 3 * ((len(part) - 1) // 3)])
+                part = part[- 3 * ((len(part) - 1) // 3):]
+            if part:
+                result.append(part)
+        else:
+            result.append(part)
+    return " ".join(result)
+
+
+def construct_prompt(text):
+    """
+    From `я купил iphone 12X за 142 990 руб без 3-x часов 12:00, и т.д.` \
+    to `<SC1>я купил [iphone 12X]<extra_id_0> за [142 990]<extra_id_1> руб без [3-x]<extra_id_2> часов [12:00]<extra_id_3>, и т.д.`.
+    """
+    result = "<SC1>"
+    etid = 0
+    token_to_add = ""
+    for token in tokenize(text) + [""]:
+        if not re.search("[a-zA-Z\d]", token):
+            if token_to_add:
+                end_match = re.search(r"(.+?)(\W*)$", token_to_add, re.M).groups()
+                result += f"[{strip_numbers(end_match[0])}]<extra_id_{etid}>{end_match[1]}"
+                etid += 1
+                token_to_add = ""
+            result += token
+        else:
+            token_to_add += token
+    return result
+
+
+def construct_answer(prompt:str, prediction:str) -> str:
+    re_prompt = re.compile(r"\[([^\]]+)\]<extra_id_(\d+)>")
+    re_pred = re.compile(r"\<extra_id_(\d+)\>(.+?)(?=\<extra_id_\d+\>|</s>)")
+    pred_data = {}
+    for match in re.finditer(re_pred, prediction.replace("\n", " ")):
+        pred_data[match[1]] = match[2].strip()
+    while match := re.search(re_prompt, prompt):
+        replace = pred_data.get(match[2], match[1])
+        prompt = prompt[:match.span()[0]] + replace + prompt[match.span()[1]:]
+    return prompt.replace("<SC1>", "")
+
+
+with open("examples.json") as f:
+    test_examples = json.load(f)
+
+
+tokenizer = GPT2Tokenizer.from_pretrained("saarus72/russian_text_normalizer", eos_token='</s>')
+model = T5ForConditionalGeneration.from_pretrained("saarus72/russian_text_normalizer")
+
+
+def predict(text):
+    input_ids = torch.tensor([tokenizer.encode(text)])
+    outputs = model.generate(input_ids, max_new_tokens=50, eos_token_id=tokenizer.eos_token_id, early_stopping=True)
+    return tokenizer.decode(outputs[0][1:])
+
+
+def norm(message, history):
+    prompt = construct_prompt(message)
+    yield f"```Prompt:\n{prompt}\nPrediction:\n...```\n..."
+    prediction = predict(prompt)
+    answer = construct_answer(prompt, prediction)
+    # yield f"```\nPrompt:\n{prompt}\nPrediction:\n{prediction}\n```\n{answer}"
+    yield f"Prompt:\n```{prompt}```\nPrediction:\n```\n{prediction}\n```\n{answer}"
+
+
+demo = gr.ChatInterface(fn=norm, stop_btn=None, examples=list(test_examples.keys())).queue()
+demo.launch()

examples.json

@@ -0,0 +1,11 @@
+{
+    "я купил iphone 12X за 142 990 руб без 3-x часов полдень и т.д.": "я купил айфон двенадцать икс за сто сорок две тысячи девятьсот девяносто руб без трёх часов полдень и т.д.",
+    "я купил айфон за 14 970 рублей": "я купил айфон за четырнадцать тысяч девятьсот семьдесят рублей",
+    "Временами я думаю, какое применение найти тем 14 697 рублям, что лежат уже больше 33 лет?": "Временами я думаю, какое применение найти тем четырнадцати тысячам шестистам девяносто семи рублям, что лежат уже больше тридцати трёх лет?",
+    "Было у отца 3 сына, но не было даже 2-3 пиджаков с блёстками за 142 990 рублей.": "Было у отца три сына, но не было даже двух-трёх пиджаков с блёстками за сто сорок две тысячи девятьсто девяносто рублей.",
+    "В школе у меня одни 5.": "В школе у меня одни пятёрки.",
+    "Было у отца 3 сына. Старшему было 35, среднему - не меньше 33, а младший на 4 младше всех. Бывает.": "Было у отца три сына. Старшему было тридцать пять, среднему - не меньше тридцати трех, а младший на четыре младше всех. Бывает.",
+    "я вырос на the beatles, меня не испугать даже 33 yellow submarine": "я вырос на зе битлз, меня не испугать даже тридцатью тремя йеллоу сабмэрин",
+    "слыш nigga ты слыхал про gitdata?": "слыш нигга ты слыхал про гитдата?",
+    "стоимость samsung 32MX Pro — всего 189 600 руб!": "стоимость самсунг тридцать два эм икс про — всего сто восемьдесят девять тысяч шестьсот руб!"
+}