add v0 code

app.py

@@ -0,0 +1,110 @@
+import gradio as gr
+import torch
+import numpy as np
+import fasttext
+import os
+import urllib
+import huggingface_hub
+from transformers import NllbTokenizer, AutoModelForSeq2SeqLM
+
+
+MODEL_URL = 'slone/nllb-rus-tyv-v1'
+
+
+lang_to_code = {
+    'Русский | Russian': 'rus_Cyrl',
+    'Тувинский | Tyvan': 'tyv_Cyrl',
+}
+
+
+def fix_tokenizer(tokenizer, new_lang='tyv_Cyrl'):
+    """ Add a new language token to the tokenizer vocabulary (this should be done each time after its initialization) """
+    old_len = len(tokenizer) - int(new_lang in tokenizer.added_tokens_encoder)
+    tokenizer.lang_code_to_id[new_lang] = old_len-1
+    tokenizer.id_to_lang_code[old_len-1] = new_lang
+    # always move "mask" to the last position
+    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 new_lang not in tokenizer._additional_special_tokens:
+        tokenizer._additional_special_tokens.append(new_lang)
+    # clear the added token encoder; otherwise a new token may end up there by mistake
+    tokenizer.added_tokens_encoder = {}
+    tokenizer.added_tokens_decoder = {}
+
+
+def translate(
+        text,
+        model,
+        tokenizer,
+        src_lang='rus_Cyrl',
+        tgt_lang='tyv_Cyrl',
+        max_length='auto',
+        num_beams=4,
+        no_repeat_ngram_size=4,
+        n_out=None,
+        **kwargs
+):
+    tokenizer.src_lang = src_lang
+    encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+    if max_length == 'auto':
+        max_length = int(32 + 2.0 * encoded.input_ids.shape[1])
+    model.eval()
+    generated_tokens = model.generate(
+        **encoded.to(model.device),
+        forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang],
+        max_length=max_length,
+        num_beams=num_beams,
+        no_repeat_ngram_size=no_repeat_ngram_size,
+        num_return_sequences=n_out or 1,
+        **kwargs
+    )
+    out = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+    if isinstance(text, str) and n_out is None:
+        return out[0]
+    return out
+
+
+def translate_wrapper(text, src, trg, correct=None):
+    src = lang_to_code.get(src)
+    trg = lang_to_code.get(trg)
+    if src == trg:
+        return 'Please choose two different languages'
+    print(text, src, trg)
+    result = translate(
+        text=text,
+        model=model,
+        tokenizer=tokenizer,
+        src=src,
+        trg=trg,
+    )
+    return result
+
+
+article = """
+Please wait until I publish all the details
+"""
+
+
+interface = gr.Interface(
+    translate_wrapper,
+    [
+        gr.Textbox(label="Text", lines=2, placeholder='text to translate '),
+        gr.Dropdown(list(lang_to_code.keys()), type="value", label='source language', value=list(lang_to_code.keys())[0]),
+        gr.Dropdown(list(lang_to_code.keys()), type="value", label='target language', value=list(lang_to_code.keys())[1]),
+    ],
+    "text",
+    title='Tyvan-Russian translaton',
+    article=article,
+)
+
+
+if __name__ == '__main__':
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_URL)
+    if torch.cuda.is_available():
+        model.cuda()
+    tokenizer = NllbTokenizer.from_pretrained(MODEL_URL)
+    fix_tokenizer(tokenizer)
+
+    interface.launch()

requirements.txt

@@ -0,0 +1,5 @@
+transformers
+sentencepiece
+numpy
+gradio>=3.18.0
+torch