OK

src/draft_model.py

@@ -30,8 +30,11 @@ class DraftModel:
             repo_id=MODEL_REPO, filename=MODEL_FILE, repo_type="model"
         )
 
+        with open(cfg_path, "r") as f:
+            cfg = json.load(f)
+
         # ---- load network ---------------------------------------------------
-        cfg = open(cfg_path, "r")
+
 
         self.net = MLP_CrossAttention(**cfg).to(self.device)
         self.net.load_state_dict(weight_path, map_location=self.device)
@@ -42,6 +45,7 @@ class DraftModel:
             hf_hub_download(repo_id=DATA_REPO, filename=ENCODING_FILE, repo_type="dataset"),
             add_nontransformed=True
         )
+        self.emb_size = next(iter(self.embed_dict.values())).shape[0]
 
     # --------------------------------------------------------------------- #
     # Public API expected by streamlit_app.py                               #