fix loading pipelines

main.py

@@ -18,7 +18,6 @@ from typing import Dict, List, Union
 
 from optimum.onnxruntime import ORTModelForSequenceClassification, ORTOptimizer, ORTQuantizer
 from optimum.onnxruntime.configuration import OptimizationConfig, AutoQuantizationConfig
-from optimum.onnxruntime.model import ORTModel
 from optimum.pipelines import pipeline as ort_pipeline
 from transformers import BertTokenizer, BertForSequenceClassification, pipeline
 
@@ -39,6 +38,11 @@ VAR2LABEL = {
     "ort_quantized_pipeline": "ONNXRuntime (Quantized)",
 }
 
+# Check if repositories exist, if not create them
+BASE_PATH.mkdir(exist_ok=True)
+QUANTIZED_BASE_PATH.mkdir(exist_ok=True)
+OPTIMIZED_BASE_PATH.mkdir(exist_ok=True)
+
 
 def get_timers(
     samples: Union[List[str], str], exp_number: int, only_mean: bool = False
@@ -64,9 +68,10 @@ def get_timers(
     timers: Dict[str, float] = {}
     for model in VAR2LABEL.keys():
         time_buffer = []
+        st.session_state["pipeline"] = load_pipeline(model)
         for _ in range(exp_number):
             with calculate_inference_time(time_buffer):
-                st.session_state[model](samples)
+                st.session_state["pipeline"](samples)
         timers[VAR2LABEL[model]] = np.mean(time_buffer) if only_mean else time_buffer
     return timers
 
@@ -87,6 +92,47 @@ def get_plot(timers: Dict[str, Union[float, List[float]]]) -> plotly.graph_objs.
     )
     fig.update_layout(title_text="Inference Time", xaxis_title="Inference Time (s)", yaxis_title="Number of Samples")
     return fig
+
+
+def load_pipeline(pipeline_name: str) -> None:
+    """
+    Load a pipeline for a given model.
+
+    Parameters
+    ----------
+    pipeline_name : str
+        Name of the pipeline to load.
+    """
+    if pipeline_name == "pt_pipeline":
+        model = BertForSequenceClassification.from_pretrained(HUB_MODEL_PATH, num_labels=3)
+        pipeline = pipeline("sentiment-analysis", tokenizer=st.session_state["tokenizer"], model=model)
+    elif pipeline_name == "ort_pipeline":
+        model = ORTModelForSequenceClassification.from_pretrained(HUB_MODEL_PATH, from_transformers=True)
+        if not ONNX_MODEL_PATH.exists():
+            model.save_pretrained(ONNX_MODEL_PATH)
+        pipeline = ort_pipeline("text-classification", tokenizer=st.session_state["tokenizer"], model=model)
+    elif pipeline_name == "ort_optimized_pipeline":
+        if not OPTIMIZED_MODEL_PATH.exists():
+            optimization_config = OptimizationConfig(optimization_level=99)
+            optimizer = ORTOptimizer.from_pretrained(HUB_MODEL_PATH, feature="sequence-classification")
+            optimizer.export(ONNX_MODEL_PATH, OPTIMIZED_MODEL_PATH, optimization_config=optimization_config)
+            optimizer.model.config.save_pretrained(OPTIMIZED_BASE_PATH)
+        model = ORTModelForSequenceClassification.from_pretrained(
+            OPTIMIZED_BASE_PATH, file_name=OPTIMIZED_MODEL_PATH.name
+        )
+        pipeline = pipeline("text-classification", tokenizer=st.session_state["tokenizer"], model=model)
+    elif pipeline_name == "ort_quantized_pipeline":
+        if not QUANTIZED_MODEL_PATH.exists():
+            quantization_config = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)
+            quantizer = ORTQuantizer.from_pretrained(HUB_MODEL_PATH, feature="sequence-classification")
+            quantizer.export(ONNX_MODEL_PATH, QUANTIZED_MODEL_PATH, quantization_config=quantization_config)
+            quantizer.model.config.save_pretrained(QUANTIZED_BASE_PATH)
+        model = ORTModelForSequenceClassification.from_pretrained(
+            QUANTIZED_BASE_PATH, file_name=QUANTIZED_MODEL_PATH.name
+        )
+        pipeline = pipeline("text-classification", tokenizer=st.session_state["tokenizer"], model=model)
+    print(type(pipeline))
+    return pipeline
         
 
 st.set_page_config(page_title="Optimum Text Classification", page_icon="⭐")
@@ -125,72 +171,10 @@ if st.session_state["init_models"]:
                 tokenizer = BertTokenizer.from_pretrained(HUB_MODEL_PATH)
                 st.session_state["tokenizer"] = tokenizer
             st.text("✅ Tokenizer loaded.")
-
-            if "pt_model" not in st.session_state:
-                pt_model = BertForSequenceClassification.from_pretrained(HUB_MODEL_PATH, num_labels=3)
-                st.session_state["pt_model"] = pt_model
-            st.text("✅ PyTorch model loaded.")
-
-            if "ort_model" not in st.session_state:
-                ort_model = ORTModelForSequenceClassification.from_pretrained(HUB_MODEL_PATH, from_transformers=True)
-                if not ONNX_MODEL_PATH.exists():
-                    ort_model.save_pretrained(ONNX_MODEL_PATH)
-                st.session_state["ort_model"] = ort_model
-            st.text("✅ ONNX Model loaded.")
-
-            if "optimized_model" not in st.session_state:
-                optimization_config = OptimizationConfig(optimization_level=99)
-                optimizer = ORTOptimizer.from_pretrained(HUB_MODEL_PATH, feature="sequence-classification")
-                if not OPTIMIZED_MODEL_PATH.exists():
-                    optimizer.export(ONNX_MODEL_PATH, OPTIMIZED_MODEL_PATH, optimization_config=optimization_config)
-                    optimizer.model.config.save_pretrained(OPTIMIZED_BASE_PATH)
-                optimized_model = ORTModelForSequenceClassification.from_pretrained(
-                    OPTIMIZED_BASE_PATH, file_name=OPTIMIZED_MODEL_PATH.name
-                )
-                st.session_state["optimized_model"] = optimized_model
-            st.text("✅ Optimized ONNX model loaded.")
-
-            if "quantized_model" not in st.session_state:
-                quantization_config = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)
-                quantizer = ORTQuantizer.from_pretrained(HUB_MODEL_PATH, feature="sequence-classification")
-                if not QUANTIZED_MODEL_PATH.exists():
-                    quantizer.export(ONNX_MODEL_PATH, QUANTIZED_MODEL_PATH, quantization_config=quantization_config)
-                    quantizer.model.config.save_pretrained(QUANTIZED_BASE_PATH)
-                quantized_model = ORTModelForSequenceClassification.from_pretrained(
-                    QUANTIZED_BASE_PATH, file_name=QUANTIZED_MODEL_PATH.name
-                )
-                st.session_state["quantized_model"] = quantized_model
-            st.text("✅ Quantized ONNX model loaded.")
-
-            if "pt_pipeline" not in st.session_state:
-                pt_pipeline = pipeline(
-                    "sentiment-analysis", tokenizer=st.session_state["tokenizer"], model=st.session_state["pt_model"]
-                )
-                st.session_state["pt_pipeline"] = pt_pipeline
-
-            if "ort_pipeline" not in st.session_state:
-                ort_pipeline = ort_pipeline(
-                    "text-classification", tokenizer=st.session_state["tokenizer"], model=st.session_state["ort_model"]
-                )
-                st.session_state["ort_pipeline"] = ort_pipeline
-
-            if "ort_optimized_pipeline" not in st.session_state:
-                ort_optimized_pipeline = pipeline(
-                    "text-classification", 
-                    tokenizer=st.session_state["tokenizer"], 
-                    model=st.session_state["optimized_model"],
-                )
-                st.session_state["ort_optimized_pipeline"] = ort_optimized_pipeline
-
-            if "ort_quantized_pipeline" not in st.session_state:
-                ort_quantized_pipeline = pipeline(
-                    "text-classification", 
-                    tokenizer=st.session_state["tokenizer"], 
-                    model=st.session_state["quantized_model"],
-                )
-                st.session_state["ort_quantized_pipeline"] = ort_quantized_pipeline
-
-            st.text("✅ All pipelines are ready.")
+            if "pipeline" not in st.session_state:
+                for pipeline in VAR2LABEL.keys():
+                    st.session_state["pipeline"] = load_pipeline(pipeline)
+            st.text("✅ Models ready.")
         sleep(2)
         loading_logs.success("🎉 Everything is ready!")
 st.session_state["init_models"] = False