Update inference_handler.py

inference_handler.py

@@ -1,20 +1,51 @@
 from typing import Dict, Any
 import torch
-from transformers import AutoConfig, AutoTokenizer
-from snp_universal_embedding import CustomSNPModel, CustomSNPConfig
+from transformers import AutoConfig, AutoModel, AutoTokenizer, PreTrainedModel, PretrainedConfig
+import torch.nn as nn
 
+# ============================================================
+# Register Custom SNP Architecture
+# ============================================================
+class CustomSNPConfig(PretrainedConfig):
+    model_type = "custom_snp"
+
+
+class CustomSNPModel(PreTrainedModel):
+    config_class = CustomSNPConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        hidden_size = getattr(config, "hidden_size", 768)
+        self.encoder = nn.Linear(hidden_size, hidden_size)
+        self.mirror_head = nn.Sequential(nn.Linear(hidden_size, hidden_size), nn.Tanh())
+        self.prism_head = nn.Sequential(nn.Linear(hidden_size, hidden_size), nn.Tanh())
+        self.projection = nn.Linear(hidden_size, 6)
+
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        x = self.encoder(input_ids.float()) if input_ids is not None else None
+        x = self.mirror_head(x)
+        x = self.prism_head(x)
+        return self.projection(x)
+
+# Register classes so Transformers recognizes "custom_snp"
+AutoConfig.register("custom_snp", CustomSNPConfig)
+AutoModel.register(CustomSNPConfig, CustomSNPModel)
+
+
+# ============================================================
+# Endpoint Handler
+# ============================================================
 class EndpointHandler:
-    def __init__(self, model_dir):
+    def __init__(self, model_dir: str):
         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.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
+        self.model = AutoModel.from_pretrained(model_dir, config=config, trust_remote_code=True)
         self.model.eval()
-        print("✅ Custom SNP model loaded successfully!")
+        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"]
@@ -25,7 +56,6 @@ class EndpointHandler:
 
         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):