()
| 354 | # ── Main ───────────────────────────────────────────────────────────── |
| 355 | |
| 356 | def main(): |
| 357 | parser = argparse.ArgumentParser(description="Extract nomic-embed-code token embeddings") |
| 358 | parser.add_argument("--output-dir", default="vendored/nomic", |
| 359 | help="Output directory (default: vendored/nomic)") |
| 360 | parser.add_argument("--device", default=None, |
| 361 | help="Device: cuda, mps, cpu (auto-detected)") |
| 362 | parser.add_argument("--skip-attention", action="store_true", |
| 363 | help="Skip simulated attention (faster, lower quality)") |
| 364 | parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, |
| 365 | help=f"Batch size (default: {BATCH_SIZE})") |
| 366 | parser.add_argument("--checkpoint", default=None, |
| 367 | help="Checkpoint file path (auto: <output-dir>/checkpoint.npz)") |
| 368 | args = parser.parse_args() |
| 369 | |
| 370 | batch_size = args.batch_size |
| 371 | |
| 372 | # Auto-detect device |
| 373 | # Prefer CPU for 7B models on Apple Silicon — MPS shares unified memory |
| 374 | # with the system and can cause OOM/crashes. CPU keeps allocation predictable. |
| 375 | # Use --device mps to override if you have enough headroom (32GB+). |
| 376 | if args.device: |
| 377 | device = args.device |
| 378 | elif torch.cuda.is_available(): |
| 379 | device = "cuda" |
| 380 | else: |
| 381 | device = "cpu" |
| 382 | |
| 383 | # Force line-buffered stdout so tee/log sees output immediately |
| 384 | sys.stdout.reconfigure(line_buffering=True) |
| 385 | |
| 386 | print(f"device={device}") |
| 387 | print(f"threads={torch.get_num_threads()}") |
| 388 | print(f"model={MODEL_NAME}") |
| 389 | print(f"output_dim={OUTPUT_DIM}") |
| 390 | print() |
| 391 | |
| 392 | # Create output dir |
| 393 | out_dir = Path(args.output_dir) |
| 394 | out_dir.mkdir(parents=True, exist_ok=True) |
| 395 | |
| 396 | checkpoint_path = args.checkpoint or str(out_dir / "checkpoint.npz") |
| 397 | |
| 398 | # ── Step 1: Load model + tokenizer ── |
| 399 | print("step 1: loading model + tokenizer...") |
| 400 | t0 = time.time() |
| 401 | tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) |
| 402 | model = AutoModel.from_pretrained( |
| 403 | MODEL_NAME, |
| 404 | trust_remote_code=True, |
| 405 | dtype=torch.float16, # 7B×2B = ~14GB (vs 28GB float32) |
| 406 | low_cpu_mem_usage=True, # Stream weights, no 2x peak during load |
| 407 | ) |
| 408 | model = model.to(device) |
| 409 | print(f" loaded in {time.time() - t0:.1f}s") |
| 410 | print(f" hidden_size={model.config.hidden_size}") |
| 411 | print(f" vocab_size={tokenizer.vocab_size}") |
| 412 | print() |
| 413 |
no test coverage detected