Update model reference

README.md

@@ -20,7 +20,7 @@ The following is adapted from [slone/nllb-rus-tyv-v1](https://huggingface.co/slo
 import torch
 from transformers import NllbTokenizer, AutoModelForSeq2SeqLM
 
-def fix_tokenizer(tokenizer, new_lang='ain_Latn'):
+def fix_tokenizer(tokenizer, new_lang):
     """ 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
@@ -36,10 +36,11 @@ def fix_tokenizer(tokenizer, new_lang='ain_Latn'):
     tokenizer.added_tokens_encoder = {}
     tokenizer.added_tokens_decoder = {}
 
-MODEL_URL = "TwentyNine/nllb-jpn-ain-v1"
+MODEL_URL = "TwentyNine/nllb-ain-kana-latin-converter-v1"
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_URL)
 tokenizer = NllbTokenizer.from_pretrained(MODEL_URL)
-fix_tokenizer(tokenizer)
+fix_tokenizer(tokenizer, 'ain_Jpan')
+fix_tokenizer(tokenizer, 'ain_Latn')
 
 def translate(
     text,