Create inference.py

inference.py

@@ -0,0 +1,46 @@
+import os
+import numpy as np
+import onnxruntime as ort
+from transformers import AutoTokenizer
+import gradio as gr
+
+# Load the ONNX model and tokenizer
+model_path = "model.onnx"
+translation_session = ort.InferenceSession(model_path)
+translation_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-fr")
+
+def translate_text(input_text):
+    # Tokenize input text
+    tokenized_input = translation_tokenizer(
+        input_text, return_tensors="np", padding=True, truncation=True, max_length=512
+    )
+    
+    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,
+        }
+    )
+
+    # Decode the output tokens
+    translated_tokens = np.argmax(outputs[0], axis=-1)
+    translated_text = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+
+    return translated_text
+
+# Create a Gradio interface
+interface = gr.Interface(
+    fn=translate_text,
+    inputs="text",
+    outputs="text",
+    title="Frenchizer Translation Model",
+    description="Translate text from English to French using an ONNX model."
+)
+
+# Launch the Gradio app
+interface.launch()