Update README.md

README.md

@@ -7,4 +7,38 @@ datasets:
 - SaranaAbidueva/buryat-russian_parallel_corpus
 metrics:
 - bleu
----
+---
+
+How to use in Python:
+```python
+from transformers import MBartForConditionalGeneration, MBart50Tokenizer
+model = MBartForConditionalGeneration.from_pretrained("SaranaAbidueva/mbart50_bur_ru")
+tokenizer = MBart50Tokenizer.from_pretrained("SaranaAbidueva/mbart50_bur_ru")
+def fix_tokenizer(tokenizer):
+    old_len = len(tokenizer) - int('bxr_XX' in tokenizer.added_tokens_encoder)
+    tokenizer.lang_code_to_id['bxr_XX'] = old_len-1
+    tokenizer.id_to_lang_code[old_len-1] = 'bxr_XX'
+    tokenizer.fairseq_tokens_to_ids["<mask>"] = len(tokenizer.sp_model) + len(tokenizer.lang_code_to_id) + tokenizer.fairseq_offset
+
+    tokenizer.fairseq_tokens_to_ids.update(tokenizer.lang_code_to_id)
+    tokenizer.fairseq_ids_to_tokens = {v: k for k, v in tokenizer.fairseq_tokens_to_ids.items()}
+    if 'bxr_XX' not in tokenizer._additional_special_tokens:
+        tokenizer._additional_special_tokens.append('bxr_XX')
+    tokenizer.added_tokens_encoder = {}
+fix_tokenizer(tokenizer)
+
+def translate(text, src='ru_RU', trg='bxr_XX', max_length=200, num_beams=5, repetition_penalty=5.0, **kwargs):
+    tokenizer.src_lang = src
+    encoded = tokenizer(text, return_tensors="pt")
+    generated_tokens = model.generate(
+        **encoded.to(model.device),
+        forced_bos_token_id=tokenizer.lang_code_to_id[trg], 
+        max_length=max_length, 
+        num_beams=num_beams,
+        repetition_penalty=repetition_penalty,
+        # early_stopping=True,
+    )
+    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+
+translate('Евгений Онегин интересная книга')
+```