Update inference_handler.py

inference_handler.py

@@ -1,42 +1,36 @@
-from typing import Dict, Any
-import torch
-from transformers import AutoConfig, AutoTokenizer
-from snp_universal_embedding import CustomSNPModel, CustomSNPConfig
-
-class EndpointHandler:
-    def __init__(self, model_dir: str):
-        print(f"Loading model from {model_dir}")
-
-        # --- Tokenizer ---
-        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True)
-        print("✅ Tokenizer loaded.")
-
-        # --- Config & Model ---
-        config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
-        self.model = CustomSNPModel(config)
-        state = torch.load(f"{model_dir}/pytorch_model.bin", map_location="cpu")
-        self.model.load_state_dict(state, strict=False)
-        self.model.eval()
-        print("✅ Custom SNP model loaded and ready.")
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """Called for each inference request"""
-        inputs = data.get("inputs") or data
-        if isinstance(inputs, dict) and "text" in inputs:
-            text = inputs["text"]
-        else:
-            text = str(inputs)
-
-        encoded = self.tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-
-        with torch.no_grad():
-            outputs = self.model(**encoded)
-            # Get mean pooled embedding
-            if hasattr(outputs, "last_hidden_state"):
-                emb = outputs.last_hidden_state.mean(dim=1).tolist()
-            elif isinstance(outputs, tuple):
-                emb = outputs[0].mean(dim=1).tolist()
-            else:
-                emb = outputs.tolist()
-
-        return {"embeddings": emb}
+from typing import Dict, Any
+import torch
+from transformers import AutoConfig, AutoTokenizer
+from snp_universal_embedding import CustomSNPModel, CustomSNPConfig
+
+class EndpointHandler:
+    def __init__(self, model_dir):
+        print(f"Loading model from {model_dir}")
+        config = CustomSNPConfig.from_pretrained(model_dir)
+        self.model = CustomSNPModel(config)
+        state = torch.load(f"{model_dir}/pytorch_model.bin", map_location="cpu")
+        self.model.load_state_dict(state, strict=False)
+        self.model.eval()
+        print("✅ Custom SNP model loaded successfully!")
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Called for each inference request"""
+        inputs = data.get("inputs") or data
+        if isinstance(inputs, dict) and "text" in inputs:
+            text = inputs["text"]
+        else:
+            text = str(inputs)
+
+        encoded = self.tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+
+        with torch.no_grad():
+            outputs = self.model(**encoded)
+            # Get mean pooled embedding
+            if hasattr(outputs, "last_hidden_state"):
+                emb = outputs.last_hidden_state.mean(dim=1).tolist()
+            elif isinstance(outputs, tuple):
+                emb = outputs[0].mean(dim=1).tolist()
+            else:
+                emb = outputs.tolist()
+
+        return {"embeddings": emb}