[update]add sent_tokenize model

cache/huggingface/hub/version.txt

@@ -0,0 +1 @@
+1

main.py

@@ -1,6 +1,7 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 import argparse
+import json
 import os
 
 from project_settings import project_path
@@ -14,6 +15,36 @@ from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 from transformers.generation.streamers import TextIteratorStreamer
 
 
+language_map = {
+    "Chinese": "zh",
+    "Czech": "cs",
+    "Danish": "da",
+    "Dutch": "nl",
+    "Flemish": "nl",
+    "English": "en",
+    "Estonian": "et",
+    "Finnish": "fi",
+    "French": "fr",
+    "German": "de",
+    "Italian": "it",
+    "Norwegian": "no",
+    "Polish": "pl",
+    "Portuguese": "pt",
+    "Russian": "ru",
+    "Spanish": "es",
+    "Swedish": "sv",
+    "Turkish": "tr",
+
+}
+
+
+nltk_sent_tokenize_languages = [
+    "czech", "danish", "dutch", "flemish", "english", "estonian",
+    "finnish", "french", "german", "italian", "norwegian",
+    "polish", "portuguese", "russian", "spanish", "swedish", "turkish"
+]
+
+
 def main():
     model_dict = {
         "facebook/m2m100_418M": {
@@ -27,6 +58,7 @@ def main():
                                  tgt_lang: str,
                                  model_name: str,
                                  ):
+        # model
         model_group = model_dict.get(model_name)
         if model_group is None:
             for k in list(model_dict.keys()):
@@ -41,15 +73,20 @@ def main():
         model = model_group["model"]
         tokenizer = model_group["tokenizer"]
 
-        tokenizer.src_lang = src_lang
+        # tokenize
+        tokenizer.src_lang = language_map[src_lang]
 
-        src_t_list = nltk.sent_tokenize(src_text)
+        if src_lang.lower() in nltk_sent_tokenize_languages:
+            src_t_list = nltk.sent_tokenize(src_text, language="")
+        else:
+            src_t_list = [src_text]
 
+        # infer
         result = ""
         for src_t in src_t_list:
             encoded_src = tokenizer(src_t, return_tensors="pt")
             generated_tokens = model.generate(**encoded_src,
-                                              forced_bos_token_id=tokenizer.get_lang_id(tgt_lang),
+                                              forced_bos_token_id=tokenizer.get_lang_id(language_map[tgt_lang]),
                                               )
             text_decoded = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
             result += text_decoded[0]
@@ -83,8 +120,10 @@ It was introduced in this [paper](https://arxiv.org/abs/2010.11125) and first re
     ]
     inputs = [
         gr.Textbox(lines=4, placeholder="text", label="Input Text"),
-        gr.Textbox(lines=1, value="en", label="Source Language"),
-        gr.Textbox(lines=1, value="zh", label="Target Language"),
+        gr.Dropdown(choices=list(language_map.keys()), value="English", label="Source Language"),
+        gr.Dropdown(choices=list(language_map.keys()), value="Chinese", label="Target Language"),
+        # gr.Textbox(lines=1, value="en", label="Source Language"),
+        # gr.Textbox(lines=1, value="zh", label="Target Language"),
         gr.Dropdown(choices=model_choices, value="facebook/m2m100_418M", label="model_name")
     ]