UX fix: demo.py (download instructions / safetensors / pickle-free demo)

demo.py

@@ -26,13 +26,17 @@ import numpy as np
 import torch
 
 ROOT = Path(__file__).resolve().parent
-sys.path.insert(0, os.environ.get("YAZ_EMBEDDER_PATH", ""))
+_emb_path = os.environ.get("YAZ_EMBEDDER_PATH", "")
+if _emb_path:                       # only add a real path (avoid inserting "" == cwd)
+    sys.path.insert(0, _emb_path)
 sys.path.insert(0, str(ROOT))
 from yaz import YazConfig, YazLM
 from yaz.semantic_router import SemanticRouter
 from scripts.gen_paraphrase_data import TRAIN_TEMPLATES
 
-CKPT = ROOT / "checkpoints" / "yaz_gen_semantic_v2.pt"
+CKPT = ROOT / "checkpoints" / "yaz_gen_semantic_v2.pt"   # PyTorch pickle (fallback)
+SAFETENSORS = ROOT / "model.safetensors"                 # recommended, pickle-free
+META = ROOT / "yaz_meta.json"
 DEFAULT_THRESHOLD = 0.08          # margin below this -> abstain (tunable; see s3 risk-coverage)
 
 
@@ -42,10 +46,22 @@ def ids(s):
 
 class YazDemo:
     def __init__(self, threshold=DEFAULT_THRESHOLD):
-        ck = torch.load(CKPT, map_location="cpu", weights_only=False)
-        self.cfg = YazConfig(**ck["cfg"])
-        self.model = YazLM(self.cfg); self.model.load_state_dict(ck["model"]); self.model.eval()
-        self.c2i = ck["country_to_target_atom"]; self.order = list(self.c2i.keys())
+        # Prefer the pickle-free safetensors + JSON sidecar (the recommended artifact);
+        # fall back to the PyTorch pickle checkpoint if they're not present.
+        import json
+        if SAFETENSORS.exists() and META.exists():
+            from safetensors.torch import load_file
+            meta = json.loads(META.read_text())
+            self.cfg = YazConfig(**meta["cfg"])
+            self.model = YazLM(self.cfg); self.model.load_state_dict(load_file(str(SAFETENSORS)))
+            self.model.eval()
+            self.c2i = meta["country_to_target_atom"]
+        else:
+            ck = torch.load(CKPT, map_location="cpu", weights_only=False)
+            self.cfg = YazConfig(**ck["cfg"])
+            self.model = YazLM(self.cfg); self.model.load_state_dict(ck["model"]); self.model.eval()
+            self.c2i = ck["country_to_target_atom"]
+        self.order = list(self.c2i.keys())
         self.threshold = threshold
         # capitals (full word) for resolving --edit RHS and for nicer output
         self.capof = {}