Update app.py

app.py

@@ -1,4 +1,3 @@
-import os
 import numpy as np
 import onnxruntime as ort
 from transformers import AutoTokenizer
@@ -15,22 +14,42 @@ def translate_text(input_text):
         input_text, return_tensors="np", padding=True, truncation=True, max_length=512
     )
     
+    # Prepare encoder inputs
     input_ids = tokenized_input["input_ids"].astype(np.int64)
     attention_mask = tokenized_input["attention_mask"].astype(np.int64)
 
-    # Run inference with the ONNX model
-    outputs = translation_session.run(
-        None,
-        {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-        }
-    )
+    # Prepare decoder inputs (start with the start token)
+    decoder_start_token_id = translation_tokenizer.cls_token_id or translation_tokenizer.pad_token_id
+    decoder_input_ids = np.array([[decoder_start_token_id]], dtype=np.int64)
 
-    # Decode the output tokens
-    translated_tokens = np.argmax(outputs[0], axis=-1)
-    translated_text = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+    # Iteratively generate output tokens
+    translated_tokens = []
+    for _ in range(512):  # Max length of output
+        # Run inference with the ONNX model
+        outputs = translation_session.run(
+            None,
+            {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "decoder_input_ids": decoder_input_ids,
+            }
+        )
+
+        # Get the next token ID
+        next_token_id = np.argmax(outputs[0][0, -1, :], axis=-1)
+        translated_tokens.append(next_token_id)
 
+        # Stop if the end-of-sequence token is generated
+        if next_token_id == translation_tokenizer.eos_token_id:
+            break
+
+        # Update decoder_input_ids for the next iteration
+        decoder_input_ids = np.concatenate(
+            [decoder_input_ids, np.array([[next_token_id]], dtype=np.int64)], axis=1
+        )
+
+    # Decode the output tokens
+    translated_text = translation_tokenizer.decode(translated_tokens, skip_special_tokens=True)
     return translated_text
 
 # Create a Gradio interface