| |
| """Step A of the side recall-test: embed THIS repo (gemma4-hack) with the |
| frozen encoder in a short-lived, low-memory process, save the 2048-d vector |
| to disk, and exit -- so the big base-model process (Step B) never has the |
| encoder resident at the same time. CPU-only to avoid competing with the |
| live MPS training run.""" |
| from __future__ import annotations |
| import sys, gc |
| from pathlib import Path |
| import numpy as np |
| HERE = Path(__file__).resolve().parent; REPO_ROOT = HERE.parent |
| sys.path.insert(0, str(REPO_ROOT)) |
| from memory_lora.encoder import load_encoder, embed_document |
| from memory_lora.codegraph import extract_repo_graph_sections, extract_repo_dependency_summary |
|
|
| def main(): |
| skip = {".git","venv","__pycache__","data","runs",".mypy_cache","scratchpad"} |
| print("Extracting codegraph sections from this repo...", flush=True) |
| graph_sections = extract_repo_graph_sections(REPO_ROOT, max_files=60, skip_dirs=skip) |
| dep = extract_repo_dependency_summary(REPO_ROOT, skip_dirs=skip) |
| sections = list(graph_sections) |
| if dep: |
| sections.append(("dependency_graph", dep)) |
| |
| for rel in ["memory_lora/core.py", "memory_lora/encoder.py", "scripts/train_memory_lora.py", |
| "memory_lora/codegraph.py", "requirements.txt"]: |
| p = REPO_ROOT / rel |
| if p.exists(): |
| sections.append((f"raw:{rel}", p.read_text(errors="ignore"))) |
| print(f" {len(sections)} sections", flush=True) |
|
|
| print("Loading encoder on CPU...", flush=True) |
| model, tok = load_encoder(device="cpu") |
| vec = embed_document(sections, model, tok, "cpu", chunk_tokens=2048, chunk_overlap=256, batch_size=2) |
| del model, tok; gc.collect() |
| if vec is None: |
| print("FAILED: no embedding", flush=True); return |
| out = REPO_ROOT / "runs" / "this_repo_emb.npy" |
| np.save(out, vec.numpy().astype("float32")) |
| print(f"Saved repo embedding {tuple(vec.shape)} -> {out}", flush=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
