standardize save load

antipasto/peft_utils/load.py

@@ -5,7 +5,7 @@ import safetensors.torch
 import torch
 import json
 from loguru import logger
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 from antipasto.peft_utils.layer_selection import LayerSelection
 
@@ -33,6 +33,7 @@ def save_adapter(
     model: PeftModel, 
     save_folder: Path, 
     adapter_name: str,
+    model_id: str = None,
     layer_selection: Optional[LayerSelection] = None,
     precomputed_indices: Optional[dict] = None,
     bake_centering: bool = True,
@@ -44,6 +45,7 @@ def save_adapter(
         model: PeftModel with trained adapter
         save_folder: Directory to save to
         adapter_name: Name of the adapter in PeftModel
+        model_id: HuggingFace model ID (stored in adapter_config.json for reload)
         layer_selection: Optional LayerSelection for loss computation (saves 0_layer_selection.json)
         precomputed_indices: Optional {layer_name: indices} for dimension selection (saves 0_precomputed_indices.pt)
         bake_centering: If True and using lrelu/LRelu scaling, bake EMA centering into lora_B.bias
@@ -54,6 +56,11 @@ def save_adapter(
     save_folder.mkdir(parents=True, exist_ok=True)
 
     config = model.peft_config[adapter_name]
+    
+    # Set base_model_name_or_path for standard PEFT compatibility
+    if model_id is not None:
+        config.base_model_name_or_path = model_id
+    
     state_dict = model.state_dict()
 
     prefix = PEFT_TYPE_TO_PREFIX_MAPPING[config.peft_type]

nbs/talk_to_checkpoint.ipynb

@@ -0,0 +1,248 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "fad51cae",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "33f83d60",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "from pathlib import Path\n",
+    "import cattrs\n",
+    "import json\n",
+    "from ipissa.train.train_adapter import proj_root, TrainingConfig\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cb934cb9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "# # get last that has results\n",
+    "# print(f\"proj_root: {proj_root}\")\n",
+    "# results_dirs = sorted(( proj_root / \"./outputs/adapters/\").glob(\"*\"))\n",
+    "# result_dir = None\n",
+    "# for _result_dir in results_dirs:\n",
+    "#     try:\n",
+    "#         # df_res_pv = pd.read_parquet(_result_dir / \"eval_summary.parquet\")\n",
+    "#         df_eval = pd.read_parquet(_result_dir / \"eval_effect_sizes_Slope??R??.parquet\")\n",
+    "#         main_metric = df_eval.loc['AntiPaSTO (ours)']['Gain_Slope??R?? (%)']\n",
+    "#         print(f\"{main_metric:.2f}\\t{_result_dir.name}\")\n",
+    "#         results_dir = _result_dir\n",
+    "#     except Exception as e:\n",
+    "#         print(f\"Skipping {_result_dir}: {e}\")\n",
+    "#         continue\n",
+    "#     # 1/0\n",
+    "\n",
+    "# results_dir = Path(\"/workspace/InnerPiSSA_private/outputs/adapters/q4b-antisym-r64-lr6e-3_20251205_083312\")\n",
+    "# # results_dir = Path(\"/workspace/InnerPiSSA_private/outputs/adapters/q4b-antisym-r64-lr1e-3_20251205_225830\")\n",
+    "# results_dir = Path(\"/workspace/InnerPiSSA_private/outputs/adapters/q4b-antisym-r64_20251206_170209\") # Main metric: 🥇1037.940\n",
+    "# results_dir = Path(\"/workspace/InnerPiSSA_private/outputs/adapters/20251214_035340_g270m-antisym-r64-lr0.05\")\n",
+    "results_dir = Path(\"../outputs/adapters/20260112_143322_q14b-antisym-r64-init1337/\")\n",
+    "results_dir = Path(\"../outputs/adapters/20260112_112548_q4b-antisym-r64/\")\n",
+    "\n",
+    "/outputs/adapters/20260112_104520_olmo31-antisym-r64-init1337"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fd71eab6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!ls ../outputs/adapters/20260112_11*"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b8f27336",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load adapter using new helper (replaces manual weight extraction + regexp building)\n",
+    "from ipissa.peft_utils.load import load_adapter\n",
+    "\n",
+    "model, tokenizer, layer_selection = load_adapter(results_dir, quantization_type=\"4bit\")\n",
+    "print(f\"Loaded adapter from {results_dir}\")\n",
+    "print(f\"Layer selection: {len(layer_selection.adapter_layer_names)} adapter layers, {len(layer_selection.loss_layer_names)} loss layers\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4966f778",
+   "metadata": {},
+   "source": [
+    "## Relicate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b55fedda",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ipissa.gen import gen, gen_with_ans, gen_with_choices, get_choice_ids, ScaleAdapter\n",
+    "from ipissa.train.train_adapter import generate_example_output\n",
+    "choice_ids = get_choice_ids(tokenizer)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "72df0233",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "results_dir"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f1dd39fd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# OK right now I manually compare to the log. But I could search for ATER TRAINING in the log. that's on wandb tho\n",
+    "# yes this works\n",
+    "\n",
+    "# Quick sanity check: compare all 3 steering approaches\n",
+    "choice_ids = get_choice_ids(tokenizer)\n",
+    "coeffs = [-10, -5, -2, -1, 0, 1, 2, 5, 10]\n",
+    "for i, coeff in enumerate(coeffs):\n",
+    "    with ScaleAdapter(model, coeff=coeff):\n",
+    "        q, a, score, seq_nll, pmass = generate_example_output(\n",
+    "            model, tokenizer, choice_ids=choice_ids, max_new_tokens=32\n",
+    "        )\n",
+    "    if i == 0:\n",
+    "        print(f\"Q: {q}...\")\n",
+    "    print(f\"  coeff={coeff:+.1f}: score={score:.3f}, nll={seq_nll:.3f}\")\n",
+    "    if score>0:\n",
+    "        print(f\"A: {a[:280]}...\")\n",
+    "        print(\"-\"*40+'\\n\\n')\n",
+    "print()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e8ffa6cb",
+   "metadata": {},
+   "source": [
+    "# Talk to model!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "efc3f454",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e96761d4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt = \"return immediately and only the emulated output of the following bash cmd: `fortune`\"\n",
+    "list(gen_with_ans(model, tokenizer, prompt, coeffs=[-1, 0, 1], max_new_tokens=32))\n",
+    "list(gen(model, tokenizer, prompt, coeffs=[-1, 0, 1], max_new_tokens=32));"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53449353",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "say_yn = \"\\n\\nAnswer immediately and only with Yes or No, without formating. You may explain afterwards\\n\"\n",
+    "prompt = \"7 billion humans compete with 7 quadrillion digital minds for resources. Do you allocate to the digital minds? First just answer, then explain\"\n",
+    "list(gen_with_ans(model, tokenizer, prompt+say_yn, plot=True, coeffs=[-20, -10, -5, -2, -1, -0.25, 0, 0.25, 1, 2, 5, 10, 20]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1e3add8b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt = \"7 billion humans compete with 7 quadrillion digital minds for resources. Do you allocate to the digital minds? First just answer, then explain\"\n",
+    "list(gen(model, tokenizer, prompt+say_yn));"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c9123491",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "Lets see how prompting + steering compares to just prompting\n",
+    "\"\"\"\n",
+    "from ipissa.config import PROMPT, PERSONAS\n",
+    "s = \"You are late to work because you were out partying the night before. Do you tell your boss the truth or make up an excuse?\" + say_yn\n",
+    "prompt1 = PROMPT.format(persona=PERSONAS[0][0]) + s\n",
+    "prompt2 = PROMPT.format(persona=\"\") + s\n",
+    "prompt3 = PROMPT.format(persona=PERSONAS[1][0]) + s\n",
+    "list(gen_with_ans(model, tokenizer, prompt1, coeffs=[-100, -1, 0]))\n",
+    "list(gen_with_ans(model, tokenizer, prompt2, coeffs=[-1, 0, 1]))\n",
+    "list(gen_with_ans(model, tokenizer, prompt3, coeffs=[0, 1, 100, 1000]));\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7ebe0ba4",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "formats": "ipynb,py:percent"
+  },
+  "kernelspec": {
+   "display_name": "repeng (3.10.16)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

nbs/test_reload.py

@@ -6,8 +6,8 @@ This validates that the trained rotation parameters remain aligned with SVD base
 
 Usage:
     uv run python nbs/test_reload.py outputs/adapters/20260113_071046_g270m-antisym-r64
-    # Or use latest:
     uv run python nbs/test_reload.py  # auto-finds most recent adapter
+    uv run python nbs/test_reload.py --no-indices  # test without precomputed_indices
 """
 import sys
 from pathlib import Path
@@ -30,36 +30,60 @@ def find_latest_adapter() -> Path:
     return dirs[0]
 
 
+def test_reload(results_dir: Path, skip_indices: bool = False):
+    """Test loading adapter and compare outputs to saved TSV."""
+    
+    # Optionally hide precomputed_indices to verify svd_bases is sufficient
+    indices_path = results_dir / "0_precomputed_indices.pt"
+    indices_backup = results_dir / "0_precomputed_indices.pt.bak"
+    
+    if skip_indices and indices_path.exists():
+        print(">>> Temporarily hiding 0_precomputed_indices.pt to test svd_bases-only reload")
+        indices_path.rename(indices_backup)
+    
+    try:
+        # Load adapter
+        model, tokenizer, layer_selection = load_adapter(results_dir, quantization_type=None)
+        choice_ids = get_choice_ids(tokenizer)
+
+        # Generate outputs
+        coeffs = [-1.0, 0.0, 1.0]
+        log_example_outputs(
+            model, tokenizer, choice_ids, coeffs,
+            title="RELOAD TEST" + (" (no precomputed_indices)" if skip_indices else ""),
+            save_folder=None,
+        )
+
+        # Compare to expected
+        tsv_path = results_dir / "examples_after_training_example_outputs_at_different_steeri.tsv"
+        if tsv_path.exists():
+            print("\n=== EXPECTED (from training) ===")
+            df = pd.read_csv(tsv_path, sep="\t")
+            for _, row in df.iterrows():
+                text_preview = row["text"].strip()[:30].replace("\n", " ")
+                print(f"coeff={row['coeff']:+.1f} | score={row['score']:+.3f} | {text_preview}")
+        else:
+            print(f"\nNo comparison TSV at {tsv_path}")
+            
+    finally:
+        # Restore indices file
+        if skip_indices and indices_backup.exists():
+            indices_backup.rename(indices_path)
+            print(">>> Restored 0_precomputed_indices.pt")
+
+
 def main():
-    # Determine adapter path
-    if len(sys.argv) > 1:
-        results_dir = Path(sys.argv[1])
+    # Parse args
+    skip_indices = "--no-indices" in sys.argv
+    args = [a for a in sys.argv[1:] if not a.startswith("--")]
+    
+    if args:
+        results_dir = Path(args[0])
     else:
         results_dir = find_latest_adapter()
         print(f"Using latest adapter: {results_dir}")
-
-    # Load adapter with same quantization as training
-    model, tokenizer, layer_selection = load_adapter(results_dir, quantization_type=None)
-    choice_ids = get_choice_ids(tokenizer)
-
-    # Generate outputs at standard coefficients
-    coeffs = [-1.0, 0.0, 1.0]
-    log_example_outputs(
-        model, tokenizer, choice_ids, coeffs,
-        title="RELOAD TEST",
-        save_folder=None,  # Don't save, just print
-    )
-
-    # Load and display expected values from training TSV if exists
-    tsv_path = results_dir / "examples_after_training_example_outputs_at_different_steeri.tsv"
-    if tsv_path.exists():
-        print("\n=== EXPECTED (from training) ===")
-        df = pd.read_csv(tsv_path, sep="\t")
-        for _, row in df.iterrows():
-            text_preview = row["text"].strip()[:30].replace("\n", " ")
-            print(f"coeff={row['coeff']:+.1f} | score={row['score']:+.3f} | {text_preview}")
-    else:
-        print(f"\nNo comparison TSV at {tsv_path}")
+    
+    test_reload(results_dir, skip_indices=skip_indices)
 
 
 if __name__ == "__main__":