File size: 5,123 Bytes
b9301f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | #!/usr/bin/env bash
# LR-decay ms1M until chat looks like a short FR sentence (heuristic), or END_TARGET.
set -euo pipefail
cd /workspace/exports/atom-ai
export PYTHONPATH=.
PY=.venv/bin/python
OUT=checkpoints/byte_tick
BASE=496000
END_TARGET=${END_TARGET:-5000000}
# --- branch parametrization (defaults reproduce the original lrd run exactly) ---
# SUFFIX checkpoint/run_metrics/train-log suffix: atom_native_step_<N>_d32_ms1M_${SUFFIX}.pt
# LR / SLR --learning-rate / --surface-learning-rate
# START_CKPT optional checkpoint (any suffix) used as --resume ONLY for the first chunk,
# when no ${SUFFIX} checkpoint exists yet at START. START defaults to its step.
# LOG_TAG tag for gen logs / coherence log / COHERENT_STEP file ("" for lrd, "_<SUFFIX minus lrd_>" otherwise)
SUFFIX=${SUFFIX:-lrd}
LR=${LR:-3e-5}
SLR=${SLR:-7.2e-5}
START_CKPT=${START_CKPT:-}
if [ -z "${LOG_TAG+x}" ]; then
if [ "$SUFFIX" = "lrd" ]; then LOG_TAG=""; else LOG_TAG="_${SUFFIX#lrd_}"; fi
fi
TAG="d32_ms1M_${SUFFIX}"
START=${1:-}
if [ -z "$START" ]; then
START=$(ls "$OUT"/atom_native_step_*_${TAG}.pt 2>/dev/null | sed -n 's/.*step_\([0-9]*\)_.*/\1/p' | sort -n | tail -1 || true)
fi
if [ -z "$START" ] && [ -n "$START_CKPT" ]; then
START=$(basename "$START_CKPT" | sed -n 's/.*step_\([0-9]*\)_.*/\1/p')
fi
if [ -z "$START" ]; then echo "cannot determine START (no ${TAG} ckpt and no START_CKPT)"; exit 1; fi
COHERENCE_LOG=${COHERENCE_LOG:-logs/coherence_watch${LOG_TAG}.log}
COHERENT_STEP_FILE=${COHERENT_STEP_FILE:-logs/COHERENT_STEP${LOG_TAG}.txt}
mkdir -p logs
echo "=== $(date -Is) train_until_coherent SUFFIX=$SUFFIX LR=$LR SLR=$SLR START=$START END_TARGET=$END_TARGET START_CKPT=${START_CKPT:-none} LOG_TAG=${LOG_TAG:-none} ===" | tee -a "$COHERENCE_LOG"
is_coherent() {
# stdin: generated text. True if looks like short FR phrase, not idi/irra attractor.
python3 - <<'PY'
import sys,re
t=sys.stdin.read().strip()
low=t.lower()
# fail attractors
if any(x in low for x in ['idi','irra','ilisateur','tuu tuu','::::']):
sys.exit(1)
# need letters + a space (multi-word) or punctuation sentence
letters=sum(c.isalpha() for c in t)
if letters < 8:
sys.exit(1)
words=re.findall(r"[A-Za-zÀ-ÿ']{2,}", t)
if len(words) < 3:
sys.exit(1)
# reject if >40% spaces
if t.count(' ') > max(1, len(t)//2):
sys.exit(1)
sys.exit(0)
PY
}
while [ "$START" -lt "$END_TARGET" ]; do
END=$((START + 8000))
[ "$END" -gt "$END_TARGET" ] && END=$END_TARGET
STEPS=$((END - START))
SKIP=$((START - BASE)); [ "$SKIP" -lt 0 ] && SKIP=0
RESUME="$OUT/atom_native_step_${START}_${TAG}.pt"
if [ ! -f "$RESUME" ] && [ -n "$START_CKPT" ] && [ "$(basename "$START_CKPT" | sed -n 's/.*step_\([0-9]*\)_.*/\1/p')" = "$START" ]; then
RESUME="$START_CKPT"
fi
CKPT_NAME="atom_native_step_${END}_${TAG}.pt"
LOG="logs/train_d32_${START}_${END}_ms1M_${SUFFIX}.log"
if [ ! -f "$RESUME" ]; then echo "MISSING $RESUME"; exit 1; fi
if [ ! -f "$OUT/$CKPT_NAME" ]; then
echo "=== $(date -Is) $START->$END skip=$SKIP ===" | tee -a "$COHERENCE_LOG"
$PY -u tools/run_atom_native.py \
--stream --data-glob 'data/corpus_fr_medium*.txt' --loop-shards --chunk-bytes 65536 \
--stream-skip-packets "$SKIP" --resume "$RESUME" --output-dir "$OUT" \
--checkpoint-name "$CKPT_NAME" --steps "$STEPS" \
--d-model 32 --n-modes 32 --n-atoms-max 64 --max-span-bytes 1 \
--field-obligatory-hard --no-enable-merge --no-payload-copy --last-atom-readout \
--efference-every 10 --learning-rate "$LR" --surface-learning-rate "$SLR" --seed 20260913 \
--atom-flush-every 64 --episode-length 8000 --episode-reset --slow-every 1 \
--field-loss-weight 0.05 --field-contrast-weight 0 --allow-parallel-train \
> "$LOG" 2>&1
[ -f "$OUT/run_metrics.json" ] && cp "$OUT/run_metrics.json" "$OUT/run_metrics_${END}_${TAG}.json" || true
else
echo "skip exists $CKPT_NAME" | tee -a "$COHERENCE_LOG"
fi
# coherence probe every tip (cold chat, 3 prompts)
CKPT="$OUT/$CKPT_NAME"
GENLOG="logs/gen_coherent_${END}${LOG_TAG}.txt"
: > "$GENLOG"
HIT=0
for p in "Bonjour, comment ça va ?" "Qui es-tu ?" "Il était une fois"; do
RESP=$($PY tools/chat_atom_native.py --checkpoint "$CKPT" --prompt "$p" --no-role-prime \
--max-packets 96 --max-length 100 --temperature 0.8 --top-k 12 2>/dev/null | sed -n 's/^Response: //p' | head -1)
echo "PROMPT: $p" >> "$GENLOG"
echo "RESP: $RESP" >> "$GENLOG"
if printf '%s' "$RESP" | is_coherent; then
HIT=$((HIT+1))
echo "COHERENT_CANDIDATE @$END :: $RESP" | tee -a "$COHERENCE_LOG"
fi
done
# val line
VAL=$(python3 -c "import json;d=json.load(open('$OUT/run_metrics_${END}_${TAG}.json'));print(d.get('validation',{}).get('loss'))" 2>/dev/null || echo '?')
echo "tip=$END val=$VAL coherent_hits=$HIT/3" | tee -a "$COHERENCE_LOG"
if [ "$HIT" -ge 2 ]; then
echo "COHERENT_REACHED step=$END hits=$HIT" | tee -a "$COHERENCE_LOG"
echo "$END" > "$COHERENT_STEP_FILE"
exit 0
fi
START=$END
done
echo "HIT_END_TARGET $END_TARGET without coherence" | tee -a "$COHERENCE_LOG"
exit 2
|