#!/bin/bash
# Export a fine-tuned checkpoint to the transformers.js layout:
#   OUT/{config,generation_config,tokenizer,tokenizer_config,special_tokens_map}.json
#   OUT/onnx/model_q4.onnx         (4-bit weights, fp32 compute, WebAssembly)
#   OUT/SHA256SUMS
#   OUT/eval/model_q4_full.onnx    (same weights, logits for every position;
#                                   for quant_eval.py, never published)
#
# Usage: export.sh CHECKPOINT OUT
# Env:   EXP  virtualenv holding optimum-onnx, onnx, onnxruntime and accelerate
#             (see docs/export-and-runtime.md). Default: the tools on PATH.
set -euo pipefail
[ $# -eq 2 ] || { echo "usage: $0 CHECKPOINT OUT" >&2; exit 2; }
CKPT=$1 OUT=$2
BIN=${EXP:+$EXP/bin/}
PY=${EXP:+$EXP/bin/python}; PY=${PY:-python3}
command -v "${BIN}optimum-cli" >/dev/null || { echo "optimum-cli not found; set EXP to the export virtualenv" >&2; exit 1; }
HERE=$(cd "$(dirname "$0")" && pwd)
WORK=$(mktemp -d "${TMPDIR:-/tmp}/tokenling-export.XXXX")
trap 'rm -rf "$WORK" "$WORK.log"' EXIT
CUDA_VISIBLE_DEVICES= "${BIN}optimum-cli" export onnx -m "$CKPT" --task text-generation-with-past "$WORK" >"$WORK.log" 2>&1 || { tail -20 "$WORK.log"; exit 1; }
"$PY" "$HERE/quantize.py" "$WORK" 2>&1 | grep -E '^q4'
mkdir -p "$OUT/onnx" "$OUT/eval"
for f in config.json generation_config.json tokenizer.json tokenizer_config.json special_tokens_map.json; do
  cp "$WORK/$f" "$OUT/$f"
done
cp "$WORK/onnx/model_q4.onnx" "$OUT/onnx/"
cp "$WORK/onnx/model_q4_full.onnx" "$OUT/eval/"
(cd "$OUT" && sha256sum *.json onnx/*.onnx > SHA256SUMS)
ls -la "$OUT" "$OUT/onnx"