Update app.py

app.py

@@ -1,25 +1,82 @@
 import gradio as gr
+from fastapi import FastAPI
 import onnxruntime as ort
-import json
+from transformers import AutoTokenizer
+from pydantic import BaseModel
+import numpy as np
+import uvicorn
+from fastapi.responses import HTMLResponse
 
-# Update MODEL_FILE path if the ONNX model is hosted elsewhere
-MODEL_FILE = "https://huggingface.co/Frenchizer/model_1/resolve/main/model.onnx"
+# Initialize FastAPI app
+app = FastAPI()
+
+# Load ONNX model and tokenizer
+MODEL_FILE = "./model.onnx"
 session = ort.InferenceSession(MODEL_FILE)
+tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-fr")
+
+# Define input model
+class TranslationInput(BaseModel):
+    input_text: str
+
+# FastAPI endpoint for model prediction
+@app.post("/predict")
+async def predict(translation_input: TranslationInput):
+    """
+    Endpoint for inference.
+    :param translation_input: Text input (e.g., in English).
+    :return: Translated text (e.g., in French).
+    """
+    # Tokenize input text
+    tokenized_input = tokenizer(
+        translation_input.input_text, 
+        return_tensors="np", 
+        padding=True
+    )
+    input_ids = tokenized_input["input_ids"]
+    
+    # Perform inference with ONNX model
+    outputs = session.run(
+        None, 
+        {"input_ids": input_ids.astype("int64")}
+    )
+
+    # Decode the output to get translated text
+    translated_text = tokenizer.decode(outputs[0][0], skip_special_tokens=True)
+    
+    return {"translated_text": translated_text}
 
-# Define your translation function
-def translate(text: str, target_language: str):
-    inputs = {"text": text, "target_language": target_language}
-    ort_inputs = {session.get_inputs()[0].name: [json.dumps(inputs)]}
-    outputs = session.run(None, ort_inputs)
-    return outputs[0][0]
-
-# Create Gradio Interface
-interface = gr.Interface(
-    fn=translate,
-    inputs=["text", "text"],  # Input boxes for source text and target language
-    outputs="text",  # Output box for translated text
-    title="Frenchizer Translator",
+# Root endpoint (optional)
+@app.get("/")
+def read_root():
+    return {"message": "ONNX model deployed with FastAPI!"}
+
+# Gradio interface function (frontend)
+def translate_text(input_text: str):
+    # Tokenize input text
+    tokenized_input = tokenizer(input_text, return_tensors="np", padding=True)
+    input_ids = tokenized_input["input_ids"]
+
+    # Perform inference with ONNX model
+    outputs = session.run(None, {"input_ids": input_ids.astype("int64")})
+    translated_text = tokenizer.decode(outputs[0][0], skip_special_tokens=True)
+    
+    return translated_text
+
+# Create Gradio interface
+gradio_interface = gr.Interface(
+    fn=translate_text,
+    inputs="text",
+    outputs="text",
+    title="English to French Translator",
+    description="A simple translator using a pre-trained ONNX model"
 )
 
-# Launch the app
-interface.launch()
+# FastAPI endpoint for Gradio app (to render in the browser)
+@app.get("/gradio")
+async def gradio_ui():
+    return HTMLResponse(gradio_interface.launch(inline=True))
+
+# Run FastAPI with Uvicorn
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)