File size: 2,607 Bytes
b8181bc 4a40b5a b8181bc 4a40b5a b8181bc 4a40b5a b8181bc 4a40b5a b8181bc 4a40b5a b8181bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
from typing import Dict, Any
import torch
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: str):
print(f"Loading model from {model_dir}")
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.")
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
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)
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}
|