Update handler.py

handler.py

@@ -30,19 +30,10 @@ def mean_pooling(model_output):
 
 class EndpointHandler():
     def __init__(self, path=""):
-        print("HELLO THIS IS THE CWD:", os.getcwd())
-        print("HELLO THIS IS THE PATH ARG:", path)
-        files = os.listdir(path)
-        for file in files:    
-            print(file)
-        # self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         task = "feature-extraction"
         self.tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/msmarco-MiniLM-L-6-v3')
         model_regular = ORTModelForFeatureExtraction.from_pretrained("jpohhhh/msmarco-MiniLM-L-6-v3_onnx", from_transformers=False)
- 
         self.onnx_extractor = pipeline(task, model=model_regular, tokenizer=self.tokenizer)
-        # self.model.to(self.device)
-        # print("model will run on ", self.device)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -52,24 +43,11 @@ class EndpointHandler():
       Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
-        print("A")
         sentences = data.pop("inputs",data)
-        print("B")
         sentence_embeddings = []
-        print("C")
         for sentence in sentences:    
-            print("D")
-            
-
             # Compute token embeddings    
             with torch.no_grad():       
                 model_output = self.onnx_extractor(sentence)  
-            print("E")
-
-            # Perform pooling. In this case, max pooling.    
-            # embedding = mean_pooling(model_output, encoded_input['attention_mask'])
-            print("F")
-            
             sentence_embeddings.append(mean_pooling(model_output))
-            print("G")
         return sentence_embeddings