Update device for simcse generator

handler.py

@@ -1,5 +1,5 @@
 
-from typing import Dict, List, Any, Union, Optional
+from typing import Dict, Union, Optional
 from pathlib import Path
 import json
 import joblib
@@ -12,11 +12,12 @@ from sklearn.base import TransformerMixin
 
 class SimcseGenerator(TransformerMixin):
     def __init__(
-        self, device: str ='cpu', batch_size: int =16, model_name: str = "princeton-nlp/unsup-simcse-bert-base-uncased"
+        self, batch_size: int =16, model_name: str = "princeton-nlp/unsup-simcse-bert-base-uncased"
     ) -> None:
 
         self.model_name = model_name
-        self.device =  torch.device(device)
+        
+        self.device =  torch.device('cpu')
 
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModel.from_pretrained(model_name).to(self.device)
@@ -53,13 +54,10 @@ class SimcseGenerator(TransformerMixin):
         return embeddings
 
 class EndpointHandler():
-    def __init__(self, device: str = "cpu"):
-        # Preload all the elements you are going to need at inference.
-        # pseudo:
-        # self.model= load_model(path)
+    def __init__(self):
 
         local_path = Path(__file__).parent
-        self.device = device
+
         with open(local_path/'stop_words.json','r') as fp:
             self.stop_words = set(json.load(fp))
 
@@ -70,7 +68,7 @@ class EndpointHandler():
         self.instruction_pipeline = joblib.load(local_path/'instruction_classification_pipeline.joblib')
         self.response_pipeline = joblib.load(local_path/'response_quality_pipeline.joblib')
         
-        self.simcse_generator = SimcseGenerator(device=self.device)
+        self.simcse_generator = SimcseGenerator()
 
     def _get_stop_word_proportion(self, s):
         s = s.lower()