Update inference.py

inference.py

@@ -1,49 +1,49 @@
-from fastapi import FastAPI
-import onnxruntime as ort
-from transformers import AutoTokenizer
-from pydantic import BaseModel
-
-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
-
-@app.post("/predict")
-async def predict(translation_input: TranslationInput):
-    """
-    Endpoint for inference.
-    :param translation_input: Text input in English.
-    :return: Translated text 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
-    outputs = session.run(
-        None, 
-        {"input_ids": input_ids.astype("int64")}
-    )
-    translated_ids = outputs[0]
-
-    # Decode output tokens
-    translated_text = tokenizer.decode(
-        translated_ids[0], 
-        skip_special_tokens=True
-    )
-    
-    return {"translated_text": translated_text}
-
-@app.get("/")
-async def root():
+from fastapi import FastAPI
+import onnxruntime as ort
+from transformers import AutoTokenizer
+from pydantic import BaseModel
+
+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
+
+@app.post("/predict")
+async def predict(translation_input: TranslationInput):
+    """
+    Endpoint for inference.
+    :param translation_input: Text input in English.
+    :return: Translated text 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
+    outputs = session.run(
+        None, 
+        {"input_ids": input_ids.astype("int64")}
+    )
+    translated_ids = outputs[0]
+
+    # Decode output tokens
+    translated_text = tokenizer.decode(
+        translated_ids[0], 
+        skip_special_tokens=True
+    )
+    
+    return {"translated_text": translated_text}
+
+@app.get("/")
+async def root():
     return {"message": "ONNX model deployed on Hugging Face Spaces!"}