Update app.py

app.py

@@ -1,14 +1,5 @@
-# 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
@@ -24,17 +15,32 @@ 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):
+    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
 
-# 翻译模型选择
-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}"  # 模型名称
+# 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()
 
-    tokenizer = MarianTokenizer.from_pretrained(model_name)  # 分词器
-    model = MarianMTModel.from_pretrained(model_name)  # 模型
 
-    return tokenizer, model
 
 
 # tesseract语言列表转pytesseract语言
@@ -77,32 +83,7 @@ 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():