Update app.py

app.py

@@ -1,6 +1,14 @@
+# OCR Translate v0.2
+# 创建人：曾逸夫
+# 创建时间：2022-07-19
+
 import os
+
+#os.system("apt-get install xclip")
+
 import gradio as gr
 import nltk
+import pyclip
 import pytesseract
 from nltk.tokenize import sent_tokenize
 from transformers import MarianMTModel, MarianTokenizer
@@ -16,33 +24,17 @@ img_dir = "./data"
 # 获取tesseract语言列表
 choices = os.popen('tesseract --list-langs').read().split('\n')[1:-1]
 
-# Função de escolha do modelo de tradução
-def model_choice(src, trg):
-    #unicamp-dl/translation-pt-en-t5
-    model_name = f"unicamp-dl/translation-{src}-{trg}-t5"
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    return tokenizer, model
-
-# Função de tradução
-def translate(input_text, trans_style):
-    if not input_text:
-        return "System prompt: There is no content to translate!"
-
-    src, trg = trans_style.split("-")
-    tokenizer, model = model_choice(src, trg)
-
-    sentences = sent_tokenize(input_text)
-    translated_text = ""
-    for sentence in sentences:
-        inputs = tokenizer.encode(sentence, return_tensors="pt", truncation=True, padding=True)
-        translated = model.generate(inputs)
-        translated_sentence = tokenizer.decode(translated[0], skip_special_tokens=True)
-        translated_text += translated_sentence + " "
 
-    return translated_text.strip()
+# 翻译模型选择
+def model_choice(src='trans_src', trg='trans_trg'):
+    # https://huggingface.co/Helsinki-NLP/opus-mt-tc-big-pt-en
+    # https://huggingface.co/Helsinki-NLP/opus-mt-tc-big-en-pt
+    model_name = f"Helsinki-NLP/opus-mt-tc-big-{src}-{trg}"  # 模型名称
 
+    tokenizer = MarianTokenizer.from_pretrained(model_name)  # 分词器
+    model = MarianMTModel.from_pretrained(model_name)  # 模型
 
+    return tokenizer, model
 
 
 # tesseract语言列表转pytesseract语言
@@ -85,7 +77,32 @@ def cp_clear():
     pyclip.clear()
 
 
+# 翻译
+def translate(input_text, inputs_transStyle):
+    # 参考：https://huggingface.co/docs/transformers/model_doc/marian
+    if input_text is None or input_text == "":
+        return "System prompt: There is no content to translate!"
+
+    # 选择翻译模型
+    trans_src, trans_trg = inputs_transStyle.split("-")[0], inputs_transStyle.split("-")[1]
+    tokenizer, model = model_choice(trans_src, trans_trg)
+
+    translate_text = ""
+    input_text_list = input_text.split("\n\n")
 
+    translate_text_list_tmp = []
+    for i in range(len(input_text_list)):
+        if input_text_list[i] != "":
+            translate_text_list_tmp.append(input_text_list[i])
+
+    for i in range(len(translate_text_list_tmp)):
+        translated_sub = model.generate(
+            **tokenizer(sent_tokenize(translate_text_list_tmp[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[2:]
 
 
 def main():
@@ -161,3 +178,5 @@ def main():
 
 if __name__ == '__main__':
     main()
+
+    #unicamp-dl/translation-pt-en-t5