Update README.md

README.md

@@ -13,4 +13,76 @@ widget:
 
 # ruT5-ASR
 
-Model was trained by [bond005](https://research.nsu.ru/en/persons/ibondarenko) to correct errors in the ASR output. The model is based on [ruT5-base](https://huggingface.co/ai-forever/ruT5-base)
+Model was trained by [bond005](https://research.nsu.ru/en/persons/ibondarenko) to correct errors in the ASR output (in particular, output of [Wav2Vec2-Large-Ru-Golos](https://huggingface.co/bond005/wav2vec2-large-ru-golos)). The model is based on [ruT5-base](https://huggingface.co/ai-forever/ruT5-base).
+
+## Usage
+
+To correct ASR outputs the model can be used as a standalone sequence-to-sequence model as follows:
+
+```python
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+import torch
+
+
+def rescore(text: str, tokenizer: T5Tokenizer,
+            model: T5ForConditionalGeneration) -> str:
+    if len(text) == 0:  # if an input text is empty, then we return an empty text too
+        return ''
+    ru_letters = set('аоуыэяеёюибвгдйжзклмнпрстфхцчшщьъ')
+    punct = set('.,:/\\?!()[]{};"\'-')
+    x = tokenizer(text, return_tensors='pt', padding=True).to(model.device)
+    max_size = int(x.input_ids.shape[1] * 1.5 + 10)
+    min_size = 3
+    if x.input_ids.shape[1] <= min_size:
+        return text  # we don't rescore a very short text
+    out = model.generate(**x, do_sample=False, num_beams=5,
+                         max_length=max_size, min_length=min_size)
+    res = tokenizer.decode(out[0], skip_special_tokens=True).lower().strip()
+    res = ' '.join(res.split())
+    postprocessed = ''
+    for cur in res:
+        if cur.isspace() or (cur in punct):
+            postprocessed += ' '
+        elif cur in ru_letters:
+            postprocessed += cur
+    return (' '.join(postprocessed.strip().split())).replace('ё', 'е')
+
+
+# load model and tokenizer
+tokenizer_for_rescoring = T5Tokenizer.from_pretrained('bond005/ruT5-ASR')
+model_for_rescoring = T5ForConditionalGeneration.from_pretrained('bond005/ruT5-ASR')
+if torch.cuda.is_available():
+    model_for_rescoring = model_for_rescoring.cuda()
+
+input_examples = [
+    'уласны в москве интерне только в большом году что лепровели',
+    'мороз и солнце день чудесный',
+    'нейро сети эта харошо',
+    'да'
+]
+
+for src in input_examples:
+    rescored = rescore(src, tokenizer_for_rescoring, model_for_rescoring)
+    print(f'{src} -> {rescored}')
+```
+
+```text
+уласны в москве интерне только в большом году что лепровели -> у нас в москве интернет только в прошлом году что ли провели
+мороз и солнце день чудесный -> мороз и солнце день чудесный
+нейро сети эта харошо -> нейросети это хорошо
+да -> да
+```
+
+## Evaluation
+This model was evaluated on the test subsets of [SberDevices Golos](https://huggingface.co/datasets/SberDevices/Golos), [Common Voice 6.0](https://huggingface.co/datasets/common_voice) (Russian part), and [Russian Librispeech](https://huggingface.co/datasets/bond005/rulibrispeech), but it was trained on the training subset of SberDevices Golos only. You can see the evaluation script on other datasets, including Russian Librispeech and SOVA RuDevices, on my Kaggle web-page https://www.kaggle.com/code/bond005/wav2vec2-t5-ru-eval
+
+*Comparison with "pure" Wav2Vec2-Large-Ru-Golos (WER, %)*:
+
+|        dataset name | pure ASR | ASR with rescoring |
+|---------------------|----------|--------------------|
+|         Voxforge Ru |    **27.08** |              40.48 |
+| Russian LibriSpeech |    **21.87** |              23.77 |
+|      Sova RuDevices |    25.41 |              **20.13** |
+|         Golos Crowd |    10.14 |               **9.42** |
+|      Golos Farfield |    20.35 |              **17.99** |
+|      CommonVoice Ru |    18.55 |              **11.60** |