Upload inference_handler.py

inference_handler.py

@@ -0,0 +1,42 @@
+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}