Update app.py

app.py

@@ -4,16 +4,16 @@ from transformers import MarianTokenizer
 import gradio as gr
 
 # Load the tokenizer from the local folder
-model_path = "./onnx_model"  # Path to the folder containing the tokenizer files
-tokenizer = MarianTokenizer.from_pretrained(model_path)
+tokenizer_path = "./onnx_model"  # Path to the local tokenizer folder
+tokenizer = MarianTokenizer.from_pretrained(tokenizer_path)
 
 # Load the ONNX model
 onnx_model_path = "./model.onnx"
 session = ort.InferenceSession(onnx_model_path)
 
-def translate_text(input_texts, max_length=512):
+def translate(texts, max_length=512):
     # Tokenize the input texts
-    inputs = tokenizer(input_texts, return_tensors="np", padding=True, truncation=True, max_length=max_length)
+    inputs = tokenizer(texts, return_tensors="np", padding=True, truncation=True, max_length=max_length)
     input_ids = inputs["input_ids"].astype(np.int64)
     attention_mask = inputs["attention_mask"].astype(np.int64)
 
@@ -24,7 +24,7 @@ def translate_text(input_texts, max_length=512):
     # Generate output tokens iteratively
     for _ in range(max_length):
         # Run the ONNX model
-        ort_outputs = session.run(
+        onnx_outputs = session.run(
             None,
             {
                 "input_ids": input_ids,
@@ -34,7 +34,7 @@ def translate_text(input_texts, max_length=512):
         )
 
         # Get the next token logits (output of the ONNX model)
-        next_token_logits = ort_outputs[0][:, -1, :]  # Shape: (batch_size, vocab_size)
+        next_token_logits = onnx_outputs[0][:, -1, :]  # Shape: (batch_size, vocab_size)
 
         # Greedy decoding: select the token with the highest probability
         next_tokens = np.argmax(next_token_logits, axis=-1)  # Shape: (batch_size,)
@@ -51,17 +51,20 @@ def translate_text(input_texts, max_length=512):
     return translations
 
 # Gradio interface
-def gradio_translate(input_texts):
-    translations = translate_text(input_texts)
-    return translations
+def gradio_translate(input_text):
+    # Split the input text into lines (assuming one sentence per line)
+    texts = input_text.strip().split("\n")
+    translations = translate(texts)
+    # Join the translations into a single string with line breaks
+    return "\n".join(translations)
 
 # Create the Gradio interface
 interface = gr.Interface(
     fn=gradio_translate,
-    inputs=gr.Textbox(lines=2, placeholder="Enter text to translate...", label="Input Text"),
-    outputs=gr.Textbox(label="Translated Text"),
+    inputs=gr.Textbox(lines=5, placeholder="Enter text to translate...", label="Input Text"),
+    outputs=gr.Textbox(lines=5, label="Translated Text"),
     title="ONNX English to French Translation",
-    description="Translate English text to French using a MarianMT ONNX model.",
+    description="Translate English text to French using an ONNX model.",
 )
 
 # Launch the Gradio app