Update app.py

app.py

@@ -1,40 +1,26 @@
-# OCR Translate v0.2
-# 创建人：曾逸夫
-# 创建时间：2022-07-19
-
 import os
-
-os.system("sudo apt-get install xclip")
-
+import pyperclip
 import gradio as gr
 import nltk
-import pyclip
 import pytesseract
+import google.generativeai as genai
 from nltk.tokenize import sent_tokenize
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import *
+import torch
+from tqdm import tqdm  # Import tqdm
 
+# Download necessary data for nltk
 nltk.download('punkt')
 
 OCR_TR_DESCRIPTION = '''# OCR Translate v0.2
 <div id="content_align">OCR translation system based on Tesseract</div>'''
 
-# 图片路径
-img_dir = "./data"
-
-# 获取tesseract语言列表
+# Getting the list of available languages for Tesseract
 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-zh-en
-    # 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
+# Replace YOUR_API_KEY with your API key from https://cloud.google.com/docs/authentication/api-keys
+genai.configure(api_key="AIzaSyC_o3nrbe2fVUifhUg0Zh2oVsweSIdlEbg")
+model = genai.GenerativeModel('gemini-pro')
 
 
 # tesseract语言列表转pytesseract语言
@@ -62,48 +48,30 @@ def clear_content():
     return None
 
 
+import pyperclip
+
 # 复制到剪贴板
 def cp_text(input_text):
-    # sudo apt-get install xclip
     try:
-        pyclip.copy(input_text)
+        pyperclip.copy(input_text)
     except Exception as e:
-        print("sudo apt-get install xclip")
+        print("Error occurred while copying to clipboard")
         print(e)
 
-
 # 清除剪贴板
 def cp_clear():
-    pyclip.clear()
+    pyperclip.clear()
 
-
-# 翻译
+# Add a translation function
 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:]
-
+    prompt = f"Please reformat the following article to have clear paragraph breaks and correct punctuation, and then translate it into {inputs_transStyle}. In the translation, do not display the original text, fictional content, or any repeated content. Ensure that the original meaning and context are preserved as much as possible."
+    # Reorder for ease of reading and translate into {inputs_transStyle}"
+    response = model.generate_content([prompt, input_text])
+    return response.text
 
 def main():
 
@@ -133,19 +101,14 @@ def main():
                 with gr.Column():
                     with gr.Row():
                         outputs_text = gr.Textbox(label="Extract content", lines=20)
-                    with gr.Row():
-                        inputs_transStyle = gr.Radio(choices=["zh-en", "en-zh"],
-                                                     type="value",
-                                                     value="zh-en",
-                                                     label='translation mode')
+                    inputs_transStyle = gr.inputs.Dropdown(choices=["Chinese (Simplified)", "Chinese (Traditional)", "English", "Japanese", "Korean"], 
+                                                           default="Chinese (Simplified)", label='translation mode')
                     with gr.Row():
                         clear_text_btn = gr.Button('Clear')
                         translate_btn = gr.Button(value='Translate', variant="primary")
 
             with gr.Row():
-                example_list = [["./data/test.png", ["eng"]], ["./data/test02.png", ["eng"]],
-                                ["./data/test03.png", ["chi_sim"]]]
-                gr.Examples(example_list, [inputs_img, inputs_lang], outputs_text, ocr_tesseract, cache_examples=False)
+                pass
 
         # -------------- 翻译 --------------
         with gr.Box():
@@ -173,8 +136,7 @@ def main():
         cp_btn.click(fn=cp_text, inputs=[outputs_tr_text], outputs=[])
         cp_clear_btn.click(fn=cp_clear, inputs=[], outputs=[])
 
-    ocr_tr.launch(inbrowser=True)
-
+    ocr_tr.launch(inbrowser=True, share=True)
 
 if __name__ == '__main__':
     main()