Update pipeline.py

pipeline.py

@@ -18,16 +18,17 @@ class PreTrainedPipeline(Pipeline):
 
 
         # Reload Keras SavedModel
-        self.model = from_pretrained_keras(model_id)
+        self.model = keras.models.load_model('./model.h5')
 
         # Number of labels
         self.num_labels = self.model.output_shape[1]
 
         # Config is required to know the mapping to label.
-        config_file = hf_hub_download(model_id, filename=CONFIG_FILENAME)
-        with open(config_file) as config:
-            config = json.load(config)
-
+        #config_file = hf_hub_download(model_id, filename=CONFIG_FILENAME)
+        #with open(config_file) as config:
+        #    config = json.load(config)
+        
+        self.num_labels = 3
         self.id2label = config.get(
             "id2label", {str(i): f"LABEL_{i}" for i in range(self.num_labels)}
         )
@@ -59,12 +60,14 @@ class PreTrainedPipeline(Pipeline):
         self.single_output_unit = (
             self.model.output_shape[1] == 1
         )  # if there are two classes
-
+        
+        
         if self.single_output_unit:
             score = predictions[0][0]
-            labels = [
-                {"label": str(self.id2label["1"]), "score": float(score)},
-                {"label": str(self.id2label["0"]), "score": float(1 - score)},
+            labels = [{"label":"pet", "score":1.0}, {"label":"other", "score":1.0}]
+            #labels = [
+            #    {"label": str(self.id2label["1"]), "score": float(score)},
+            #    {"label": str(self.id2label["0"]), "score": float(1 - score)},
             ]
         else:
             labels = [