hotchpotch's picture
Add FP16 and INT8 ONNX/OpenVINO variants at all MRL widths
76c3670 verified
|
Raw History Blame Contribute Delete
5.63 kB

Training scripts for Bekko Static Embedding Multilingual v1

This directory contains the complete public-data training path for hotchpotch/bekko-static-embedding-v1. The training entry points and export utility are standalone: they do not import the Bekko package, UniR, or another private project. The tokenizer is deliberately a separate deterministic stage, so a complete reproduction consists of the two training scripts, the FP16 export utility, the YAML recipe, and the data-selection JSON distributed alongside them.

Files

File Purpose
train_tokenizer.py Builds the uncased 200,000-token WordPiece tokenizer, preserving a 30,000-token English core.
export_fp16.py Creates the formal FP16 delivery from the FP32 master without changing the tokenizer or model configs.
train_model.py Validates and materializes the selected public Bekko rows through shared_id, trains the MNRL + MRL static model, performs intermediate evaluation, exports the model, and runs final HAKARI evaluation.

In the Hugging Face model repository, this README is published in training_scripts/ together with mnrl-full.yaml (the formal recipe) and bekko-unsupervised-multilingual-40-direct.json (the pinned public-data selection). In the source repository, those two assets live in configs/.

Requirements

Use Python 3.11 or later, uv, and a CUDA-capable PyTorch installation for full training. The exact release environment used Python 3.13, PyTorch 2.8.0, SentenceTransformers 5.6.0, and Transformers 5.14.1. Install at least:

pip install \
  'torch==2.8.0' \
  'sentence-transformers==5.6.0' \
  'transformers==5.14.1' \
  'datasets>=4.4.2' 'tokenizers>=0.22.2' 'pyyaml>=6.0.3' \
  'numpy' 'simplemma>=1.1.2' 'nltk>=3.9.2' 'xxhash>=3.8.1' \
  'accelerate>=1.1.0' 'wandb>=0.21.0'

Install a PyTorch build appropriate for the available CUDA runtime; the command above is only the version constraint. flash-attn is optional and is not required by these scripts. If Weights & Biases is unavailable, set tracking.enabled: false in the YAML.

Final and intermediate evaluation require a local clone of HAKARI-Bench. Install that project according to its own instructions, then set evaluation.hakari_repo in the YAML to its checkout. No private data, private dataset cache, or UniR checkout is required.

1. Train the tokenizer

Run from the directory containing these files, or pass absolute paths:

uv run train_tokenizer.py \
  --corpus-dir output/tokenizer-corpus \
  --output-dir output/bekko-multilingual-wp200k-en30k-uncased

The script downloads deterministic prefixes from the pinned public hotchpotch/bekko-embedding-v1-unsupervised revision. It learns the English 30k core first and then extends it to 200k tokens using non-bitext multilingual text. Keep the script defaults unchanged for an exact v1 tokenizer.

2. Configure the training recipe

Before starting, edit only local paths in mnrl-full.yaml:

  • model.tokenizer_model: the output directory from step 1;
  • run.output_dir and data.stage_cache_root: persistent local storage;
  • evaluation.hakari_repo: the HAKARI-Bench checkout;
  • evaluation.full_device: the GPU to use.

Do not change the data spec, sampler, MNRL settings, MRL widths, optimizer, or seed for an exact v1 reproduction. data.stage_cache_root is a resumable local cache. Its rows come from the public dataset and are keyed by its public shared_id field.

3. Train, monitor, export, and evaluate

python train_model.py --config mnrl-full.yaml

The v1 recipe trains one full epoch with MNRL temperature 0.03, learning rate 0.25, 10% warmup, global batch size 16,384, and seed 42. It schedules exactly 100 lightweight intermediate reranking evaluations: five queries per NanoMIRACL target and three queries per MNanoBEIR target.

After training, the script saves the final SentenceTransformers model, changes only the exported tokenizer maximum length from the training limit of 512 to 32,768, then starts HAKARI's genuine full --all evaluation. Outputs are under run.output_dir, including final/ and full-eval/hakari-results/.

4. Export the formal FP16 delivery

The trainer retains the FP32 master in final/. Since 2026-09-14, publish a separate FP16 copy:

python export_fp16.py \
  --source output/bekko-static-embedding-multilingual-v1/final \
  --output output/bekko-static-embedding-multilingual-v1/final-fp16

Adjust these paths to run.output_dir. Output must not already exist. The utility casts only embedding.weight; model configs and tokenizer remain unchanged. The FP16 weight file is 409,600,096 bytes instead of 819,200,096. The trainer's automatic full evaluation runs on the FP32 master. A separate matched HAKARI comparison evaluated FP16 storage using --dtype fp32: native micro nDCG@10 × 100 was 45.88353 versus 45.88177 for FP32 storage. This tests stored-value rounding, not FP16 inference kernels.

Release checks

The formal v1 export stores FP16 weights, has 1,024 dimensions, and supports MRL widths 32, 64, 128, 256, 512, and 1,024. Its primary SHA-256 values are:

model.safetensors  586e165a868f306822174e8110020fa8759577cc62a537a702941747732858f8
tokenizer.json     308823060352e35bf9fff6b42ef343f31cc68fa800e284568cc44e9cc0f70a73

The 32,768-token export setting is an input ceiling, not a promise that a 32,768-token batch fits in GPU memory. Compare functional invariants and the full evaluation outputs as well as hashes when reproducing on other hardware.