Update handler.py

handler.py

@@ -3,13 +3,13 @@ from transformers import AutoTokenizer, AutoModelForMaskedLM
 import torch
 
 class EndpointHandler:
-    def __init__(self, model_name: str = "naver/splade-cocondenser-ensembledistil", 
-                 max_seq_length: int = 256, device: str = "cpu"):
+    def __init__(self, path: str = ""):
+        # Device to use for inference. "cuda" if CUDA is available, else "cpu"
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
         # Load tokenizer and model
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForMaskedLM.from_pretrained(model_name).to(device)
-        self.max_seq_length = max_seq_length
-        self.device = device
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForMaskedLM.from_pretrained(path).to(self.device)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -29,7 +29,7 @@ class EndpointHandler:
             return_tensors="pt",
             padding=True,
             truncation=True,
-            max_length=self.max_seq_length,
+            max_length=self.model.config.max_position_embeddings,  # Max length according to the model configuration
         ).to(self.device)
         with torch.no_grad():
             logits = self.model(**inputs).logits
@@ -48,4 +48,4 @@ class EndpointHandler:
                 {"indices": nz_tokens.tolist(), "values": nz_weights.tolist()}
             )
 
-        return output[0] if isinstance(inputs, str) else output
+        return output[0] if isinstance(inputs, str) else output