Add inference.py

inference.py

@@ -275,11 +275,30 @@ class NeuralLayer:
         try:
             import torch, torch.nn as nn
             self.torch=torch; self.nn=nn
-            ckpt=torch.load(model_path, map_location='cpu', weights_only=False)
-            cfg=ckpt['config']
-            self.rules=ckpt['rule_vocab']; self.n_rules=len(self.rules)
-            self.char2id=ckpt['char_vocab']; self.max_len=cfg['max_len']
-            self.model=self._build(cfg); self.model.load_state_dict(ckpt['model_state'])
+            # Load config/vocab from JSON sidecars if present, else extract from .pt
+            cfg_path = _HERE / "config.json"
+            rv_path  = _HERE / "rule_vocab.json"
+            cv_path  = _HERE / "char_vocab.json"
+            sf_path  = _HERE / "mba_classifier.safetensors"
+            if cfg_path.exists() and rv_path.exists() and cv_path.exists():
+                cfg          = json.load(open(cfg_path))
+                self.rules   = json.load(open(rv_path))
+                self.char2id = json.load(open(cv_path))
+            else:
+                ckpt         = torch.load(model_path, map_location='cpu', weights_only=False)
+                cfg          = ckpt['config']
+                self.rules   = ckpt['rule_vocab']
+                self.char2id = ckpt['char_vocab']
+            self.n_rules = len(self.rules)
+            self.max_len = cfg['max_len']
+            self.model   = self._build(cfg)
+            # Load weights: prefer .safetensors, fall back to .pt
+            if sf_path.exists():
+                from safetensors.torch import load_file
+                self.model.load_state_dict(load_file(str(sf_path)))
+            else:
+                ckpt = torch.load(model_path, map_location='cpu', weights_only=False)
+                self.model.load_state_dict(ckpt['model_state'])
             self.model.eval()
             self.reduction_rules=[i for i,r in enumerate(self.rules)
                                    if not r.startswith('comm_') and not r.startswith('fold_')]
@@ -294,14 +313,15 @@ class NeuralLayer:
         class RC(nn.Module):
             def __init__(s,vocab_size,d_model,n_heads,n_layers,d_ff,max_len,n_rules,dropout=0.1):
                 super().__init__()
-                s.embed=nn.Embedding(vocab_size,d_model,padding_idx=0)
-                s.pos=nn.Embedding(max_len,d_model)
-                el=nn.TransformerEncoderLayer(d_model,n_heads,d_ff,dropout,batch_first=True,norm_first=True)
-                s.enc=nn.TransformerEncoder(el,n_layers)
-                s.norm=nn.LayerNorm(d_model); s.head=nn.Linear(d_model,n_rules)
+                s.embed     = nn.Embedding(vocab_size,d_model,padding_idx=0)
+                s.pos_embed = nn.Embedding(max_len,d_model)
+                el = nn.TransformerEncoderLayer(d_model,n_heads,d_ff,dropout,batch_first=True,norm_first=True)
+                s.encoder   = nn.TransformerEncoder(el,n_layers)
+                s.norm      = nn.LayerNorm(d_model)
+                s.head      = nn.Linear(d_model,n_rules)
             def forward(s,x):
                 pm=(x==0); pos=torch.arange(x.size(1),device=x.device).unsqueeze(0)
-                e=s.embed(x)+s.pos(pos); enc=s.enc(e,src_key_padding_mask=pm)
+                e=s.embed(x)+s.pos_embed(pos); enc=s.encoder(e,src_key_padding_mask=pm)
                 L=(~pm).float().sum(1,keepdim=True).clamp(min=1)
                 return s.head(s.norm((enc*(~pm).unsqueeze(-1).float()).sum(1)/L))
         RC.torch=self.torch
@@ -457,8 +477,10 @@ def _get_neural(model_path: Optional[str]) -> Optional[NeuralLayer]:
     global _neural
     if _neural is not None: return _neural
     if model_path is None:
-        default = _HERE / 'mba_classifier_v2.pt'
-        if default.exists(): model_path=str(default)
+        sf = _HERE / 'mba_classifier.safetensors'
+        pt = _HERE / 'mba_classifier_v2.pt'
+        if sf.exists():   model_path=str(sf)
+        elif pt.exists(): model_path=str(pt)
         else: return None
     _,rw,_=_get_sym(8)
     _neural=NeuralLayer(str(model_path), rw, bits=8)