Update main.py

sioaeko · Jun 30, 2024 · 3d7a54c · 3d7a54c
1 parent 9deb0ee
commit 3d7a54c
Showing 1 changed file with 21 additions and 39 deletions.
diff --git a/main.py b/main.py
@@ -1,7 +1,5 @@
 import gradio as gr
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-import torch
-from pyngrok import ngrok
 import os
 
 # Load the model and tokenizer
@@ -37,42 +35,26 @@ def translate_file(file, src_lang, tgt_lang):
     except Exception as e:
         return f"File translation error: {str(e)}"
 
-def create_ui():
-    with gr.Blocks() as demo:
-        gr.Markdown("# NLLB Translator")
-
-        with gr.Tab("Text Translation"):
-            with gr.Row():
-                src_lang = gr.Dropdown(choices=available_languages, label="Source Language")
-                tgt_lang = gr.Dropdown(choices=available_languages, label="Target Language")
-            input_text = gr.Textbox(lines=5, label="Input Text")
-            output_text = gr.Textbox(lines=5, label="Translated Text")
-            translate_btn = gr.Button("Translate Text")
-            translate_btn.click(fn=translate_text, inputs=[input_text, src_lang, tgt_lang], outputs=output_text)
-
-        with gr.Tab("File Translation"):
-            file_input = gr.File(label="Upload file to translate")
-            file_src_lang = gr.Dropdown(choices=available_languages, label="Source Language")
-            file_tgt_lang = gr.Dropdown(choices=available_languages, label="Target Language")
-            file_output = gr.Textbox(label="Translation Status")
-            file_translate_btn = gr.Button("Translate File")
-            file_translate_btn.click(fn=translate_file, inputs=[file_input, file_src_lang, file_tgt_lang], outputs=file_output)
-
-        gr.Markdown(f"Currently using the {model_name} AI translation model")
-
-    return demo
-
-if __name__ == "__main__":
-    demo = create_ui()
+with gr.Blocks() as demo:
+    gr.Markdown("# NLLB Translator")
 
-    # Set up ngrok tunneling (optional)
-    if os.environ.get('USE_NGROK'):
-        ngrok_token = os.environ.get('NGROK_TOKEN')
-        if ngrok_token:
-            ngrok.set_auth_token(ngrok_token)
-            public_url = ngrok.connect(port=7860)
-            print(f"Public URL: {public_url}")
-        else:
-            print("NGROK_TOKEN not set. Skipping ngrok tunnel.")
+    with gr.Tab("Text Translation"):
+        with gr.Row():
+            src_lang = gr.Dropdown(choices=available_languages, label="Source Language")
+            tgt_lang = gr.Dropdown(choices=available_languages, label="Target Language")
+        input_text = gr.Textbox(lines=5, label="Input Text")
+        output_text = gr.Textbox(lines=5, label="Translated Text")
+        translate_btn = gr.Button("Translate Text")
+        translate_btn.click(fn=translate_text, inputs=[input_text, src_lang, tgt_lang], outputs=output_text)
 
-    demo.launch(share=True)
+    with gr.Tab("File Translation"):
+        file_input = gr.File(label="Upload file to translate")
+        file_src_lang = gr.Dropdown(choices=available_languages, label="Source Language")
+        file_tgt_lang = gr.Dropdown(choices=available_languages, label="Target Language")
+        file_output = gr.Textbox(label="Translation Status")
+        file_translate_btn = gr.Button("Translate File")
+        file_translate_btn.click(fn=translate_file, inputs=[file_input, file_src_lang, file_tgt_lang], outputs=file_output)
+
+    gr.Markdown(f"Currently using the {model_name} AI translation model")
+
+demo.launch(share=True)