Promote aux weight 1 model

knot2vec.py

@@ -5,10 +5,12 @@
 #     "datasets",
 #     "huggingface-hub",
 #     "jinja2",
+#     "matplotlib",
 #     "numpy",
 #     "safetensors",
 #     "snappy",
 #     "torch",
+#     "umap-learn",
 # ]
 # ///
 """Knot2Vec: embed knot diagrams so that diagrams of the same knot land together."""
@@ -28,9 +30,11 @@ from pathlib import Path
 
 import datasets
 import huggingface_hub.utils
+import matplotlib.pyplot as plt
 import numpy as np
 import snappy
 import torch
+import umap
 from database_knotinfo import link_list
 from datasets import Dataset, load_dataset
 from huggingface_hub import (
@@ -46,6 +50,12 @@ from torch.nn import functional
 
 REPO = "jgalego/knot2vec"
 DATA = "jgalego/knot2vec-diagrams"
+LANDMARKS = {
+    "3_1": "trefoil",
+    "4_1": "figure-eight",
+    "10_124": "T(3,5)",
+    "11n_34": "Conway",
+}
 HERE = Path(__file__).parent
 # Progress bars redraw in place, which shows up as garbage in HF Jobs logs.
 PROGRESS = sys.stderr.isatty()
@@ -264,12 +274,13 @@ def schedule(step, warmup, total):
     return 0.5 * (1 + math.cos(math.pi * (step - warmup) / max(1, total - warmup)))
 
 
-def push_result(repo, name, result):
+def push_result(repo, name, result, revision=None):
     """Upload a result as results/<name>.json in the model repo."""
     HfApi().upload_file(
         path_or_fileobj=json.dumps(result, indent=1).encode(),
         path_in_repo=f"results/{name}.json",
         repo_id=repo,
+        revision=revision,
         commit_message=f"Add {name} results",
     )
 
@@ -351,8 +362,20 @@ def train(args):
     (Path(args.output) / "README.md").unlink(missing_ok=True)
     if args.push:
         api = HfApi()
-        api.upload_folder(folder_path=args.output, repo_id=args.repo, commit_message="Upload model")
-        api.upload_file(path_or_fileobj=__file__, path_in_repo="knot2vec.py", repo_id=args.repo)
+        if args.revision:
+            api.create_branch(args.repo, branch=args.revision, exist_ok=True)
+        api.upload_folder(
+            folder_path=args.output,
+            repo_id=args.repo,
+            revision=args.revision,
+            commit_message="Upload model",
+        )
+        api.upload_file(
+            path_or_fileobj=__file__,
+            path_in_repo="knot2vec.py",
+            repo_id=args.repo,
+            revision=args.revision,
+        )
         push_result(
             args.repo,
             "train",
@@ -370,6 +393,7 @@ def train(args):
                 "runtime_s": round(time.time() - start),
                 "device": torch.cuda.get_device_name() if device == "cuda" else "cpu",
             },
+            args.revision,
         )
 
 
@@ -379,10 +403,30 @@ def nearest(queries, gallery, k=5, chunk=4096):
     return torch.cat([(queries[i : i + chunk] @ gallery.T).topk(k).indices for i in starts])
 
 
+def knot_map(gallery, knots, path):
+    """UMAP of the catalogue's embeddings, coloured by signature, with a few landmarks."""
+    xy = umap.UMAP(metric="cosine", random_state=0).fit_transform(gallery.cpu().numpy())
+    fig, ax = plt.subplots(figsize=(8, 7))
+    points = ax.scatter(
+        *xy.T, c=knots["signature"], cmap="Spectral", vmin=-MAX_SIGNATURE, vmax=MAX_SIGNATURE,
+        s=2, linewidths=0,
+    )
+    fig.colorbar(points, ax=ax, label="Signature", shrink=0.7)
+    for i, name in enumerate(knots["name"]):
+        if name in LANDMARKS:
+            ax.annotate(
+                f"{name} {LANDMARKS[name]}", xy[i], xytext=(6, 6), textcoords="offset points",
+                fontsize=9, arrowprops={"arrowstyle": "-", "color": "0.3"},
+            )
+    ax.set_axis_off()
+    fig.savefig(path, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+
+
 def evaluate(args):
     """Name held-out diagrams by nearest canonical diagram; check the invariant heads."""
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = Knot2Vec.from_pretrained(args.model).to(device).eval()
+    model = Knot2Vec.from_pretrained(args.model, revision=args.revision).to(device).eval()
     knots = table(args.data, "knots", "train").to_dict()
     index = {n: i for i, n in enumerate(knots["name"])}
     gallery = model.embed(knots["pd"]).to(device)
@@ -404,10 +448,13 @@ def evaluate(args):
             "n": int(mask.sum()),
             "top1": round(hits1[mask].float().mean().item(), 3),
             "top5": round(hits5[mask].float().mean().item(), 3),
-            **{
-                f"{key}_accuracy": round((out[key] == truth[key])[mask].float().mean().item(), 3)
-                for key in ("signature", "determinant")
-            },
+            "signature_accuracy": round(
+                (out["signature"] == truth["signature"])[mask].float().mean().item(), 3
+            ),
+            "determinant_within_10pct": round(
+                (abs(out["determinant"] - truth["determinant"]) <= 0.1 * truth["determinant"])[mask]
+                .float().mean().item(), 3
+            ),
             "volume_mae": round((out["volume"] - truth["volume"])[mask].abs().mean().item(), 3),
             "top1_by_crossings": {
                 int(c): round(hits1[mask & (crossings == c)].float().mean().item(), 3)
@@ -415,8 +462,16 @@ def evaluate(args):
             },
         }
     print(json.dumps(result, indent=1))
+    Path(args.output).mkdir(parents=True, exist_ok=True)
+    knot_map(gallery, knots, Path(args.output) / "map.png")
     if args.push:
-        push_result(args.repo, "eval", result)
+        push_result(args.repo, "eval", result, args.revision)
+        HfApi().upload_file(
+            path_or_fileobj=Path(args.output) / "map.png",
+            path_in_repo="results/map.png",
+            repo_id=args.repo,
+            revision=args.revision,
+        )
 
 
 def embed(args):
@@ -496,13 +551,14 @@ def main():
     train_parser.add_argument("--lr", type=float, default=3e-4)
     train_parser.add_argument("--warmup", type=int, default=1000)
     train_parser.add_argument("--temperature", type=float, default=0.05)
-    train_parser.add_argument("--aux-weight", type=float, default=0.1)
+    train_parser.add_argument("--aux-weight", type=float, default=1.0)
     train_parser.add_argument("--dim", type=int, default=256)
     train_parser.add_argument("--depth", type=int, default=6)
     train_parser.add_argument("--max-crossings", type=int, default=64)
     train_parser.add_argument("--seed", type=int, default=0)
     train_parser.add_argument("--output", default="out/model")
     train_parser.add_argument("--repo", default=REPO)
+    train_parser.add_argument("--revision", default=None, help="branch to push to")
     train_parser.add_argument("--push", action="store_true")
     train_parser.set_defaults(run=train)
 
@@ -511,6 +567,8 @@ def main():
     eval_parser.add_argument("--data", default=DATA)
     eval_parser.add_argument("--limit", type=int, default=None, help="number of test diagrams")
     eval_parser.add_argument("--repo", default=REPO, help="where --push stores the result")
+    eval_parser.add_argument("--revision", default=None, help="model branch")
+    eval_parser.add_argument("--output", default="out/eval", help="where the knot map goes")
     eval_parser.add_argument("--push", action="store_true")
     eval_parser.set_defaults(run=evaluate)
 

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3c9782c590fdc438543af62d807a2ab69f7f4111a38ab8d18c8ce9ebe1acbff8
+oid sha256:786f8b9cb77a3b430caaa584312e3db9a5939c4f480c5b038963634b642d205a
 size 19708148

results/eval.json

@@ -4,41 +4,41 @@
  "chance_top1": 7.7e-05,
  "seen": {
   "n": 46676,
-  "top1": 0.849,
-  "top5": 0.973,
-  "signature_accuracy": 0.794,
-  "determinant_within_10pct": 0.199,
-  "volume_mae": 1.507,
+  "top1": 0.881,
+  "top5": 0.976,
+  "signature_accuracy": 0.917,
+  "determinant_within_10pct": 0.267,
+  "volume_mae": 1.237,
   "top1_by_crossings": {
    "3": 1.0,
    "5": 0.75,
    "6": 1.0,
-   "7": 0.714,
-   "8": 0.764,
-   "9": 0.835,
-   "10": 0.76,
-   "11": 0.879,
-   "12": 0.809,
-   "13": 0.858
+   "7": 0.821,
+   "8": 0.833,
+   "9": 0.756,
+   "10": 0.769,
+   "11": 0.844,
+   "12": 0.864,
+   "13": 0.889
   }
  },
  "unseen": {
   "n": 5184,
-  "top1": 0.856,
-  "top5": 0.969,
-  "signature_accuracy": 0.781,
-  "determinant_within_10pct": 0.207,
-  "volume_mae": 1.583,
+  "top1": 0.88,
+  "top5": 0.976,
+  "signature_accuracy": 0.891,
+  "determinant_within_10pct": 0.248,
+  "volume_mae": 1.39,
   "top1_by_crossings": {
-   "4": 0.5,
-   "5": 0.5,
+   "4": 0.75,
+   "5": 0.75,
    "6": 0.75,
-   "8": 0.75,
-   "9": 0.7,
-   "10": 0.691,
-   "11": 0.863,
-   "12": 0.833,
-   "13": 0.865
+   "8": 0.917,
+   "9": 0.65,
+   "10": 0.779,
+   "11": 0.806,
+   "12": 0.86,
+   "13": 0.893
   }
  }
 }

results/train.json

@@ -1,19 +1,19 @@
 {
  "step": 19999,
- "loss": 0.2043,
- "contrastive": 0.1145,
- "invariants": 0.8985,
- "pair_top1": 0.971,
- "minutes": 101.7,
+ "loss": 0.6226,
+ "contrastive": 0.1426,
+ "invariants": 0.48,
+ "pair_top1": 0.961,
+ "minutes": 101.8,
  "data": "jgalego/knot2vec-diagrams",
  "steps": 20000,
  "batch_size": 512,
  "learning_rate": 0.0003,
  "temperature": 0.05,
- "aux_weight": 0.1,
+ "aux_weight": 1.0,
  "dim": 256,
  "depth": 6,
  "parameters": 4924943,
- "runtime_s": 6106,
+ "runtime_s": 6109,
  "device": "NVIDIA A10G"
 }