docs/TECHNICAL_SPEC.md

Technical Specification: OCR-Aware Mixed-Precision Release Pipeline

Status: Implementation specification
Target repository: AutomatosX/AX-Unlimited-OCR-3B-MoE-MLX-MXFP8

1. System overview

BF16 Baidu checkpoint
        |
        v
source + dataset preflight
        |
        v
BF16 baseline evaluation
        |
        v
one-group-at-a-time MXFP8 sensitivity
        |
        v
joint quality/throughput head calibration
        |
        v
generated executable precision map
        |
        v
staged BF16 -> mixed MXFP8 conversion
        |
        v
candidate vs BF16 vs Sahil validation
        |
        v
accuracy + performance + R-SWA release gates
        |
        v
content-addressed release manifest
        |
        v
atomic Hugging Face publication

2. Components

2.1 Dataset validator

Input layout:

eval_dir/
  manifest.json
  images/<sample>.<ext>
  ground_truth/<sample>.txt

The validator checks unique stems, complete pairs, non-empty ground truth, manifest consistency, category counts, and a SHA-256 digest over the manifest and referenced files. Validation occurs before importing MLX.

2.2 Sensitivity runner

quantization/layer_sensitivity.py loads the unquantized model and computes one BF16 baseline. For each module group it reloads BF16, quantizes only the group, and records:

• matched concrete module paths;
• sample count and task coverage;
• CER, digit CER, and table score;
• deltas from BF16;
• elapsed time and status.

The runner rejects a source config containing quantization or quantization_config.

2.3 Precision-map generator

The generator reads layer_groups from the sensitivity result and the approved calibration result. It starts from the conservative base map, promotes sensitive or failed groups to BF16, and applies only calibration overrides that were selected by existing release thresholds. Generated rules are top-level executable entries. Metadata is stored under keys beginning with _, which the converter ignores.

Rule precedence is:

1. exact concrete module path;
2. more-specific wildcard path;
3. general wildcard path;
4. unlisted modules default to BF16.

2.4 Converter

quantization/mixed_precision_convert.py calls the pinned mlx-vlm converter with per-module quantization settings. The main model uses MXFP8 with eight data bits and group size 32; the calibrated LM head uses 8-bit affine quantization with group size 32. It writes into a new sibling staging directory and atomically renames it only after verifying:

• at least one module was quantized;
• every requested MXFP8 or affine8 rule matched;
• at least one Safetensors file exists;
• model metadata selects native Unlimited-OCR;
• precision_map.json and quantization_summary.json were written.

The requested final output directory must not exist.

2.5 Comparative benchmark runner

For each model (bf16, reference_mxfp8, candidate), use identical dataset, prompt, profile, max tokens, and preprocessing settings. Results are written to:

artifacts/<run_id>/
  provenance.json
  sensitivity_results.json
  calibration_results.json
  calibration_baseline_accuracy.json
  calibration_reference_performance.json
  calibration_bfloat16_accuracy.json
  calibration_bfloat16_performance.json
  calibration_mxfp8_accuracy.json
  calibration_mxfp8_performance.json
  calibration_affine8_accuracy.json
  calibration_affine8_performance.json
  generated_precision_map.json
  bf16_accuracy.json
  reference_accuracy.json
  candidate_accuracy.json
  reference_performance.json
  candidate_performance.json
  candidate_rswa.json
  release_manifest.json

2.6 Release gate

The release gate is a pure-data operation. It reads benchmark JSON and candidate artifacts without loading a model. Every gate records:

{
  "name": "candidate_cer_vs_bf16",
  "passed": true,
  "actual": 0.012,
  "limit": 0.01,
  "detail": "candidate - bf16 absolute CER"
}

release_approved is true only when every required gate passes. Missing or inconclusive values fail the gate. The calibration decision is not trusted as an assertion: the gate verifies SHA-256 records for all eight raw calibration benchmark files, recomputes their aggregates and quality/throughput deltas, and reselects the fastest precision that passes the fixed limits. It also rejects a release when calibration and held-out evaluation reuse any image hash.

2.7 Publisher

The publisher accepts only a candidate directory containing:

• model Safetensors and valid index when sharded;
• tokenizer and processor metadata;
• generated precision map;
• quantization summary;
• approved release manifest;
• README/model card and license.

Before upload it recomputes candidate hashes and compares them with the approved manifest. It then creates one Hugging Face commit. Source checkpoint download or fallback is forbidden in this path.

3. CLI

Primary orchestration command:

python quantization/run_pipeline.py \
  --model-path reference/Unlimited-OCR \
  --source-revision SOURCE_COMMIT_SHA \
  --reference-model sahilchachra/unlimited-ocr-mxfp8-mlx \
  --reference-revision REFERENCE_COMMIT_SHA \
  --calibration-dir datasets/selection \
  --eval-dir datasets/held-out-release \
  --output-dir models/AX-Unlimited-OCR-3B-MoE-MLX-MXFP8 \
  --artifacts-dir artifacts/ocr-aware-v1 \
  --image test_data/test_invoice.png

Publication command:

python scripts/publish_optimized_model.py \
  --model-dir models/AX-Unlimited-OCR-3B-MoE-MLX-MXFP8 \
  --manifest artifacts/ocr-aware-v1/release_manifest.json \
  --artifacts-dir artifacts/ocr-aware-v1 \
  --repo-id AutomatosX/AX-Unlimited-OCR-3B-MoE-MLX-MXFP8

Useful modes:

• --step preflight
• --step sensitivity
• --step calibrate — interim map, three LM-head converts, select affine/mxfp8/bf16 head
• --step precision-map
• --step convert
• --step validate
• --step gate
• --step publish
• --step all (runs the full order above)
• --dry-run

calibrate is required before precision-map for schema-3 releases: it writes calibration_results.json plus the eight raw calibration_*_{accuracy,performance}.json evidence files on the calibration dataset only.

4. Provenance schema

provenance.json includes:

• UTC run ID and timestamps;
• source and reference identifiers plus immutable 40-character commit SHAs;
• file SHA-256 and byte sizes;
• dataset digest and category counts;
• Python, macOS, MLX, mlx-vlm, and huggingface-hub versions;
• Apple chip and unified memory;
• generation parameters and artifact digests;
• generated precision-map digest.

5. Model metadata

Published config.json must have:

• architecture UnlimitedOCRForCausalLM;
• model_type: unlimited-ocr;
• a positive sliding-window size;
• an MXFP8 quantization object generated by the converter.
• per-module entries showing 119 MXFP8 modules and one affine8 LM head.

Published processor_config.json must use:

• processor_class: UnlimitedOCRHFProcessor;
• sft_format: unlimitedocr.

If tokenizer_config.json is present, it must use the same processor_class: UnlimitedOCRHFProcessor (convert normalization rewrites both files so mlx-vlm and Hub-style loaders do not disagree).

6. Error handling

• Never delete or overwrite an existing output model directory.
• Each subprocess failure stops the pipeline.
• Missing metrics, zero samples, unmatched rules, natural short generations, and unavailable Metal are explicit failures or inconclusive release blockers.
• Partial candidate output remains only in a named staging directory until the converter cleans it up.
• Publication recomputes hashes immediately before creating the remote commit.

7. Testing

Unit tests cover:

• source rejection for already quantized models;
• sensitivity-result schema parsing;
• objective calibration selection across BF16, MXFP8, and affine8 heads;
• sensitive-group promotion to executable BF16 rules;
• specificity precedence in the precision map;
• output-directory lifecycle;
• release-gate pass, fail, missing, and inconclusive cases;
• hash mismatch rejection;
• publisher refusal without an approved manifest;
• publisher operation list containing candidate rather than source weights.

Integration tests mock model conversion and Hugging Face APIs. Hardware smoke tests run BF16 and candidate inference on a real image. Full release validation runs only on a Metal-capable Apple Silicon host.

8. Release verification

After publication:

1. Read remote file metadata and confirm all expected files exist.
2. Compare remote LFS SHA-256 with the approved manifest.
3. Load the exact remote revision with mlx-vlm.
4. Run deterministic OCR on the smoke image.
5. Record the final revision in the local release manifest and handoff report.

9. Schema-3 re-release checklist (MLX Hub)

The publisher requires MLX_RELEASE_SCHEMA_VERSION = 3. A legacy artifacts/ocr-aware-v1/release_manifest.json with schema_version: 1 cannot be re-uploaded without a full re-gate.

1. Build disjoint calibration and evaluation trees (no shared FUNSD IDs):

python eval_dataset/build_release_set.py --clean \
  --output-root eval_dataset/calibration \
  --funsd-ids 0,1 --generator-seed 20260720 \
  --funsd-revision <FUNSD_SHA> --latin-font <PATH> --cjk-font <PATH>

python eval_dataset/build_release_set.py --clean \
  --output-root eval_dataset/evaluation \
  --funsd-ids 2,3,4,5 --generator-seed 20260721 \
  --funsd-revision <FUNSD_SHA> --latin-font <PATH> --cjk-font <PATH>

2. Run the OCR-aware pipeline / release gate with --calibration-dir eval_dataset/calibration and --eval-dir eval_dataset/evaluation (schema 3, all 23 gates).
3. Confirm convert normalization rewrote both processor and tokenizer processor_class fields to UnlimitedOCRHFProcessor.
4. Dry-run then publish:

python scripts/publish_optimized_model.py \
  --model-dir models/AX-Unlimited-OCR-3B-MoE-MLX-MXFP8 \
  --manifest artifacts/<run>/release_manifest.json \
  --artifacts-dir artifacts/<run> \
  --dry-run
# then without --dry-run when approved

5. Post-publish smoke against the remote revision; keep evidence under artifacts/.

Model-card claims must stay local-measurement only until a public protocol (OmniDocBench / olmOCR-Bench) is run on the candidate.