Reza2kn/Bina-0.1-Flax-BF16

Bina 0.1 Flax BF16

An experimental Flax/JAX port of Reza2kn/Bina-0.1, a Persian OCR fine-tune of Surya OCR 2.

The checkpoint remains in framework-neutral Safetensors format. The included Flax NNX runtime implements the Qwen3.5 vision tower, multimodal prefill, full-attention decoder layers, Gated DeltaNet recurrent layers, KV/recurrent caches, and greedy generation.

Precision decision

BF16 is the parity release. Naive row-wise INT8 and INT4 candidates were evaluated but are not represented as parity-preserving:

Candidate	Result
BF16 Flax decoder, 4 printed pages	177/177 tokens exact
BF16 Flax decoder, long handwriting page	450/450 tokens exact
BF16 full image-to-transcript page	14/14 tokens exact
INT8, first 4 printed pages	177/177 tokens exact
INT8 broader gate	Diverged on printed and handwriting pages
INT4 layer canary	Rejected; excessive layer-output drift

The parameter-free Flax Gated DeltaNet canary was bit-exact at its BF16 output; its FP32 recurrent-state maximum absolute difference was 1.79e-7.

Files

• full_decoder.py: Flax/JAX Qwen3.5 language decoder.
• vision_prefill.py: Flax/JAX vision tower and multimodal prefill.
• production_runtime.py: persistent, compiled BF16 inference worker.
• preprocessing.py: NumPy preprocessing plus an exact compatibility path that uses the Transformers fast PyTorch image processor before JAX inference.
• production_benchmark.py: end-to-end performance and transcript-parity gate.
• smoke_production_runtime.py: reproducible public-API production smoke test.
• examples/bina_bbox_flax_tpu_v6.ipynb: all-device TPU v6 bbox OCR pipeline.
• generate_parity.py: teacher-free transcript parity runner.
• gdn_parity.py, layer0_parity.py: focused numerical gates.
• receipts/: machine-readable validation evidence.

Production runtime

production_runtime.py is the persistent compiled runtime. It uses fused GPU vision attention, a chunked Gated DeltaNet prefill kernel, length-bucketed prefill caches, a device-resident generation loop, and fixed-shape 4,096-token decode caches. For production parity, use prepare_image_fast: it reproduces the source Transformers fast processor exactly, then hands its tensors to JAX. production_benchmark.py is the matching end-to-end parity/performance gate.

RTX 5080 Laptop measurements after compilation:

Fixture	Vision	Prefill	Decode	Parity
Printed, 336 patches / 244 prompt tokens	21.0 ms	6,964 tok/s	190 tok/s	14/14 exact
Handwriting, 9,280 patches / 2,480 prompt tokens	138 ms	7,860 tok/s	185 tok/s	450/450 exact

Warm model-core latency was 124.5 ms for the printed crop and 2.880 seconds for the dense 450-token handwriting page, versus 228 ms and 5.316 seconds through the original PyTorch raw-image service. These are controlled-fixture measurements, not a claim of a dataset-wide 1.8x speedup; the PyTorch service numbers also include preprocessing and string decode.

Cold compilation of a new bucket is expensive; construct one BinaProductionRuntime per worker, keep it alive, and always set its persistent compilation-cache directory. The included TPU v6 notebook runs one persistent replica on every addressable TPU device.

Microbatching

The runtime also exposes encode_vision_batch, prefill_batch, and generate_batch. An RTX 5080 batch-8 canary produced 8/8 exact transcripts in 314.5 ms warm model time, or 39.3 ms per crop. The TPU notebook defaults to batch 8 and pads final partial batches. It preserves natural crop geometry for source parity, so a new processor tensor shape can still trigger an XLA compilation; batches are grouped by exact prepared shape to reuse compiled executables where possible.

Decoding correctness

The tokenizer EOS ID is 2. Older experimental runtime code incorrectly used 248044, causing correct transcripts to continue into repeated or unrelated text. The current single and batched generators stop on EOS 2 with exact greedy argmax decoding. A repeated 16-gram detector remains as a host-side safety signal; it does not alter ordinary transcripts and TPU shard uploads are rejected if repetition, truncation, or token-length QA thresholds fail.

Source integrity

• Source model: Reza2kn/Bina-0.1
• Architecture: Qwen3_5ForConditionalGeneration
• Source precision: BF16
• Source weight SHA-256: 2193be4ef3d2366438121a15b7a1dea2bb85b24f83145e5a39bfa1f387891ada

This derivative follows the upstream OpenRAIL license. Review the source model license before use or redistribution.