Preload models

app.py

@@ -12,10 +12,21 @@ models = {
     "Pruned ONNX Optimized FP16": "tryolabs/bert-large-uncased-wwm-squadv2-optimized-f16",
 }
 
+loaded_models = {
+    "Pruned ONNX Optimized FP16": hf_hub_download(
+        repo_id=models["Pruned ONNX Optimized FP16"], filename="model.onnx"
+    ),
+    "Base model": AutoModelForQuestionAnswering.from_pretrained(models["Base model"]),
+    "Pruned model": AutoModelForQuestionAnswering.from_pretrained(
+        models["Pruned model"]
+    ),
+}
+
 
 def run_ort_inference(model_name, inputs):
-    model_path = hf_hub_download(repo_id=models[model_name], filename="model.onnx")
-    sess = InferenceSession(model_path, providers=["CPUExecutionProvider"])
+    sess = InferenceSession(
+        loaded_models[model_name], providers=["CPUExecutionProvider"]
+    )
     start_time = time.time()
     output = sess.run(None, input_feed=inputs)
     end_time = time.time()
@@ -24,9 +35,9 @@ def run_ort_inference(model_name, inputs):
 
 def run_normal_hf(model_name, inputs):
     start_time = time.time()
-    model = AutoModelForQuestionAnswering.from_pretrained(models[model_name])
+    output = loaded_models[model_name](**inputs).values()
     end_time = time.time()
-    return model(**inputs).values(), (end_time - start_time)
+    return output, (end_time - start_time)
 
 
 def inference(model_name, context, question):