#!/usr/bin/env bash
# Purpose: build an importance matrix from a calibration text so that llama-quantize can spend
#          its bit budget on the weights that carry the most signal for your workload, and record
#          which text produced it so the quantisation can be reproduced later
# Platform: all (spark, strix, mac, nvidia); uses the llama.cpp build from Part 6
# Minimum memory: 12 GB; the full-precision model is loaded once, and -ngl decides where
# Assumes: llama-imatrix on PATH or under LLAMA_BIN, a full-precision GGUF (made by
#          quantise-five-ways.sh or by convert_hf_to_gguf.py), and a calibration text file of
#          your own that is disjoint from anything you will evaluate on
#
# Usage: bash make-imatrix.sh MODEL_BF16_GGUF CALIBRATION_TXT [OUTPUT_GGUF]
#   MODEL_BF16_GGUF  the full-precision GGUF to collect statistics from
#   CALIBRATION_TXT  plain text resembling your workload; a few hundred kilobytes is plenty
#   OUTPUT_GGUF      where to write the matrix; defaults to imatrix.gguf beside the model
#
# Environment: LLAMA_BIN (directory holding llama-imatrix), NGL (default 99),
#              CHUNKS (limit the number of calibration chunks), LABBOOK (default labbook.md)
set -euo pipefail

MODEL="${1:-}"
CALIB="${2:-}"
OUT="${3:-}"
NGL="${NGL:-99}"
LABBOOK="${LABBOOK:-labbook.md}"

die() { echo "make-imatrix: $*" >&2; exit 1; }

[[ -n "$MODEL" && -n "$CALIB" ]] || die "usage: bash make-imatrix.sh MODEL_BF16_GGUF CALIBRATION_TXT [OUTPUT_GGUF]"
[[ -f "$MODEL" ]] || die "$MODEL does not exist"
[[ -f "$CALIB" ]] || die "$CALIB does not exist"

IMATRIX_BIN="${LLAMA_BIN:-}/llama-imatrix"
if [[ -z "${LLAMA_BIN:-}" ]]; then
  IMATRIX_BIN="$(command -v llama-imatrix || true)"
fi
[[ -x "$IMATRIX_BIN" ]] || die "llama-imatrix not found; set LLAMA_BIN to the directory holding it (built in Part 6)"

[[ -n "$OUT" ]] || OUT="$(dirname "$MODEL")/imatrix.gguf"

CALIB_BYTES=$(wc -c < "$CALIB" | tr -d ' ')
if [[ "$CALIB_BYTES" -lt 20000 ]]; then
  echo "make-imatrix: WARNING: $CALIB is only $CALIB_BYTES bytes." >&2
  echo "              A matrix from a handful of sentences describes a handful of sentences." >&2
  echo "              Aim for a few hundred kilobytes of text that resembles your real traffic." >&2
fi

# The hash is the point of this line: an importance matrix is only reproducible if you can say
# which bytes produced it, and calibration texts get edited.
if command -v sha256sum >/dev/null 2>&1; then
  CALIB_SHA="$(sha256sum "$CALIB" | cut -d' ' -f1)"
elif command -v shasum >/dev/null 2>&1; then
  CALIB_SHA="$(shasum -a 256 "$CALIB" | cut -d' ' -f1)"
else
  CALIB_SHA="unavailable"
fi

echo "==> Collecting activation statistics"
echo "    model:       $MODEL"
echo "    calibration: $CALIB ($CALIB_BYTES bytes, sha256 ${CALIB_SHA:0:16})"
echo "    output:      $OUT"
echo

STARTED=$(date +%s)
IMATRIX_ARGS=(
  -m "$MODEL"
  -f "$CALIB"
  -o "$OUT"
  -ngl "$NGL"
  --output-frequency 20
  --save-frequency 50
  --parse-special
)
if [[ -n "${CHUNKS:-}" ]]; then
  IMATRIX_ARGS+=(--chunks "$CHUNKS")
fi

"$IMATRIX_BIN" "${IMATRIX_ARGS[@]}"
ELAPSED=$(( $(date +%s) - STARTED ))

[[ -f "$OUT" ]] || die "llama-imatrix finished but $OUT was not written"

echo
echo "==> Wrote $OUT in ${ELAPSED}s"
echo "    Read it back with: $IMATRIX_BIN --in-file $OUT --show-statistics"
echo "    Two matrices from different texts can be merged by passing --in-file twice."

python3 - "$LABBOOK" "$MODEL" "$CALIB" "$CALIB_SHA" "$CALIB_BYTES" "$OUT" "$ELAPSED" <<'PY'
import json, os, platform, sys, time

labbook, model, calib, sha, size, out, elapsed = sys.argv[1:8]
record = {
    "lab": "part-16/lab-quantise-five-ways-and-measure/imatrix",
    "run_id": time.strftime("%Y%m%dT%H%M%S"),
    "model": os.path.basename(model),
    "calibration_file": os.path.basename(calib),
    "calibration_sha256": sha,
    "calibration_bytes": int(size),
    "imatrix_file": os.path.basename(out),
    "seconds": int(elapsed),
    "host": platform.platform(),
    "date": time.strftime("%Y-%m-%d"),
}
with open(labbook, "a", encoding="utf-8") as handle:
    handle.write(json.dumps(record) + "\n")
print(f"recorded in {labbook}")
PY