Update handler.py

handler.py

@@ -7,9 +7,9 @@ device = 0 if torch.cuda.is_available() else -1
 
 # multi-model list
 multi_model_list = [
-    {"model_id": "omarabb315/gemma-2B-2nd_filtered_3_full", "task": "text-generation"},
-    {"model_id": "omarabb315/gemma-2B-2nd_filtered_3_16bit", "task": "text-generation"},
-    {"model_id": "omarabb315/Gemma-2-9B-filtered_3_4bits", "task": "text-generation"},
+    {"model_id": "gemma-2B-2nd_filtered_3_full", "model_path": "omarabb315/gemma-2B-2nd_filtered_3_full", "task": "text-generation"},
+#    {"model_path": "omarabb315/gemma-2B-2nd_filtered_3_16bit", "task": "text-generation"},
+#    {"model_path": "omarabb315/Gemma-2-9B-filtered_3_4bits", "task": "text-generation"},
 ]
 
 class EndpointHandler():
@@ -17,7 +17,7 @@ class EndpointHandler():
         self.multi_model={}
         # load all the models onto device
         for model in multi_model_list:
-            self.multi_model[model["model_id"]] = pipeline(model["task"], model=model["model_id"], device=device)
+            self.multi_model[model["model_id"]] = pipeline(model["task"], model=model["model_path"], device=device)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         # deserialize incomin request