Update model.py

model.py

@@ -10,9 +10,9 @@ Input:  (B, 8, 16, 16)  — adapted latent patches
 Output: gate_vectors (B, 64, 17), patch_features (B, 64, 256), logits
 
 Usage:
-    from geometric_model import SuperpositionPatchClassifier, load_from_hub
+    from geometric_model import load_from_hub, extract_features
 
-    model = load_from_hub()  # downloads from AbstractPhil/geovocab-patch-maker
+    model, config = load_from_hub()  # reads config.json + model.pt from Hub
     out = model(patches)
 
     # Gate vectors: explicit geometric properties per patch
@@ -324,30 +324,66 @@ class SuperpositionPatchClassifier(nn.Module):
 # Hub Loading
 # ══════════════════════════════════════════════════════════════════════════════
 
+def load_config(repo_id="AbstractPhil/geovocab-patch-maker", config_file="config.json"):
+    """Load model config from HuggingFace Hub."""
+    import json
+    from huggingface_hub import hf_hub_download
+
+    path = hf_hub_download(repo_id=repo_id, filename=config_file)
+    with open(path, "r") as f:
+        return json.load(f)
+
+
+def from_config(config, device="cpu"):
+    """Instantiate model from config dict (no weights)."""
+    return SuperpositionPatchClassifier(
+        embed_dim=config["embed_dim"],
+        patch_dim=config["patch_dim"],
+        n_bootstrap=config["n_bootstrap"],
+        n_geometric=config["n_geometric"],
+        n_heads=config["n_heads"],
+        dropout=config.get("dropout", 0.0),
+    ).to(device)
+
+
 def load_from_hub(
     repo_id="AbstractPhil/geovocab-patch-maker",
-    filename="model.pt",
+    weights_file="model.pt",
+    config_file="config.json",
     device="cuda" if torch.cuda.is_available() else "cpu",
 ):
-    """Load pretrained model from HuggingFace Hub."""
+    """
+    Load pretrained model from HuggingFace Hub.
+
+    Reads config.json for architecture, model.pt for weights.
+    Falls back to config embedded in checkpoint if config.json missing.
+    """
     from huggingface_hub import hf_hub_download
 
-    path = hf_hub_download(repo_id=repo_id, filename=filename)
-    ckpt = torch.load(path, map_location=device, weights_only=False)
-    cfg = ckpt["config"]
+    # Load config
+    try:
+        config = load_config(repo_id, config_file)
+        print(f"✓ Config loaded from {config_file}")
+    except Exception:
+        config = None
+
+    # Load weights
+    weights_path = hf_hub_download(repo_id=repo_id, filename=weights_file)
+    ckpt = torch.load(weights_path, map_location=device, weights_only=False)
 
-    model = SuperpositionPatchClassifier(
-        embed_dim=cfg["embed_dim"],
-        patch_dim=cfg["patch_dim"],
-        n_bootstrap=cfg["n_bootstrap"],
-        n_geometric=cfg["n_geometric"],
-        n_heads=cfg["n_heads"],
-        dropout=0.0,
-    ).to(device).eval()
+    # Config priority: config.json > checkpoint config
+    if config is None:
+        config = ckpt["config"]
+        print(f"  Config from checkpoint (no {config_file} found)")
 
+    model = from_config(config, device=device)
     model.load_state_dict(ckpt["model_state_dict"])
-    print(f"✓ Loaded {repo_id} (epoch {ckpt.get('epoch', '?')})")
-    return model
+    model.eval()
+
+    epoch = ckpt.get("epoch", "?")
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"✓ Loaded {repo_id} (epoch {epoch}, {n_params:,} params)")
+    return model, config
 
 
 @torch.no_grad()
@@ -395,6 +431,9 @@ def extract_features(model, patches, batch_size=256):
 # ══════════════════════════════════════════════════════════════════════════════
 
 if __name__ == "__main__":
+    import json
+
+    # Test 1: Direct instantiation
     model = SuperpositionPatchClassifier()
     n_params = sum(p.numel() for p in model.parameters())
     print(f"SuperpositionPatchClassifier: {n_params:,} parameters")
@@ -406,4 +445,14 @@ if __name__ == "__main__":
     print(f"  local_dim:       {out['local_dim_logits'].shape}")
     print(f"  struct_topo:     {out['struct_topo_logits'].shape}")
     print(f"  patch_shapes:    {out['patch_shape_logits'].shape}")
-    print(f"  global_features: {out['global_features'].shape}")
+    print(f"  global_features: {out['global_features'].shape}")
+
+    # Test 2: From config
+    import os
+    cfg_path = os.path.join(os.path.dirname(__file__), "config.json")
+    if os.path.exists(cfg_path):
+        with open(cfg_path) as f:
+            config = json.load(f)
+        model2 = from_config(config)
+        print(f"\n  from_config: {sum(p.numel() for p in model2.parameters()):,} params")
+        print(f"  config: {config['model_type']} embed={config['embed_dim']} patches={config['num_patches']}")