Update app.py

app.py

@@ -4,10 +4,8 @@ import onnxruntime as ort
 from transformers import AutoTokenizer
 from pydantic import BaseModel
 import numpy as np
-import uvicorn
-from fastapi.responses import HTMLResponse
 
-# Initialize FastAPI app
+# Initialize FastAPI and Gradio
 app = FastAPI()
 
 # Load ONNX model and tokenizer
@@ -15,7 +13,7 @@ MODEL_FILE = "./model.onnx"
 session = ort.InferenceSession(MODEL_FILE)
 tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-fr")
 
-# Define input model
+# Define input model for FastAPI
 class TranslationInput(BaseModel):
     input_text: str
 
@@ -24,8 +22,8 @@ class TranslationInput(BaseModel):
 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).
+    :param translation_input: Text input in English.
+    :return: Translated text in French.
     """
     # Tokenize input text
     tokenized_input = tokenizer(
@@ -34,49 +32,29 @@ async def predict(translation_input: TranslationInput):
         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
+    # Decode output and return translated text
     translated_text = tokenizer.decode(outputs[0][0], skip_special_tokens=True)
-    
     return {"translated_text": translated_text}
 
-# 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
+# Gradio Interface
+def gradio_predict(input_text):
+    response = predict(TranslationInput(input_text=input_text))
+    return response["translated_text"]
 
-# Create Gradio interface
+# Gradio interface for the web app
 gradio_interface = gr.Interface(
-    fn=translate_text,
+    fn=gradio_predict,
     inputs="text",
     outputs="text",
-    title="English to French Translator",
-    description="A simple translator using a pre-trained ONNX model"
+    live=True
 )
 
-# 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)
+# Launch Gradio app
+gradio_interface.launch(inline=True, server_name="0.0.0.0", server_port=8000)