app update

app.py

@@ -1,7 +1,6 @@
 # OCR Translate v0.1
 # 创建人：曾逸夫
 # 创建时间：2022-06-14
-# email: zyfiy1314@163.com
 
 import os
 
@@ -13,13 +12,6 @@ from transformers import MarianMTModel, MarianTokenizer
 
 nltk.download('punkt')
 
-# ----------- 翻译 -----------
-# https://huggingface.co/Helsinki-NLP/opus-mt-en-zh
-modchoice = "Helsinki-NLP/opus-mt-en-zh"  # 模型名称
-
-tokenizer = MarianTokenizer.from_pretrained(modchoice)  # 分词器
-model = MarianMTModel.from_pretrained(modchoice)  # 模型
-
 OCR_TR_DESCRIPTION = '''# OCR Translate v0.1
 <div id="content_align">基于Tesseract的OCR翻译系统</div>'''
 
@@ -30,6 +22,17 @@ img_dir = "./data"
 choices = os.popen('tesseract --list-langs').read().split('\n')[1:-1]
 
 
+# 翻译模型选择
+def model_choice(src="en", trg="zh"):
+    # https://huggingface.co/Helsinki-NLP/opus-mt-en-zh
+    model_name = f"Helsinki-NLP/opus-mt-{src}-{trg}"  # 模型名称
+
+    tokenizer = MarianTokenizer.from_pretrained(model_name)  # 分词器
+    model = MarianMTModel.from_pretrained(model_name)  # 模型
+
+    return tokenizer, model
+
+
 # tesseract语言列表转pytesseract语言
 def ocr_lang(lang_list):
     lang_str = ""
@@ -66,11 +69,19 @@ def translate(input_text):
     if input_text is None or input_text == "":
         return "系统提示：没有可翻译的内容！"
 
-    translated = model.generate(**tokenizer(sent_tokenize(input_text), return_tensors="pt", padding=True))
-    tgt_text = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
-    translate_text = "".join(tgt_text)
+    tokenizer, model = model_choice()
+
+    translate_text = ""
+    input_text_list = input_text.split("\n\n")
+
+    for i in range(len(input_text_list)):
+        translated_sub = model.generate(
+            **tokenizer(sent_tokenize(input_text_list[i]), return_tensors="pt", truncation=True, padding=True))
+        tgt_text_sub = [tokenizer.decode(t, skip_special_tokens=True) for t in translated_sub]
+        translate_text_sub = "".join(tgt_text_sub)
+        translate_text = translate_text + "\n\n" + translate_text_sub
 
-    return translate_text
+    return translate_text[2:]
 
 
 def main():