diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 5ff0e36..d77250a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -43,3 +43,16 @@ jobs: run: | PYTHONPATH=src python -m src.cli list test "$(PYTHONPATH=src python -m src.cli list | jq 'length')" -ge 3 + + export-fixture: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: { python-version: "3.11" } + - run: | + python -m pip install --upgrade pip + pip install torch==2.12.1 transformers==5.14.1 + pip install -e ".[dev,train,export]" + - name: IMF export + parity gate fixture tests (torch vs ORT) + run: PYTHONPATH=src python -m pytest tests/test_imf_export.py tests/test_imf_parity.py -v diff --git a/docs/imf-v1.md b/docs/imf-v1.md index 776ffa5..282bde9 100644 --- a/docs/imf-v1.md +++ b/docs/imf-v1.md @@ -47,7 +47,7 @@ model.zip | `tokenizer` | enum | `bytes` (the only v1 value) | | `opset` | int | 7..14; must equal the graphs' opset | | `decoder` | enum | `plain` \| `kv` (`kv` requires decoder-kv.onnx) | -| `precision` | enum | `fp32` \| `fp16` \| `int8` | +| `precision` | enum | `fp32` \| `fp16` \| `int8` (fp16 = torch-native half export: float16 graph IO, int64 ids unchanged; runtimes read dtypes from the session — the ORT float16 converter produces all-zero hiddens on real ByT5 and must not be used) | | `license` | str | non-empty (strict gate) | | `trained_from` | str | repo + run/checkpoint id | | `metrics` | list | `{name, value, protocol, source}`; `source` must be a `RESULTS.md#anchor` (strict gate) | @@ -111,8 +111,26 @@ PYTHONPATH=src python -m imf validate --strict # release gate PYTHONPATH=src python -m imf info # print manifest PYTHONPATH=src python -m imf pack --source \ --metadata [--readme ] --out # sha256 computed +PYTHONPATH=src python -m imf parity --checkpoint \ + --test-data # WO03 gate; writes parity into the zip +PYTHONPATH=src python -m imf golden --inputs --out \ + # cross-runtime golden set: 100 fixed strings, Python = reference ``` +The parity gate compares ONNX KV greedy decode against the transformers +decoder loop (the exact math the export wraps — not `generate`, whose +config-dependent behavior no runtime implements) over >= 500 test pairs; +it writes `{samples, cer_delta}` into metadata and refuses to leave the +zip non-strict. On Modal the same gate runs headless: + +``` +modal run --detach src/gpu/modal_export.py::main --model urd-g2p +modal run --detach src/gpu/modal_export.py::parity --model urd-g2p +``` + +CI runs the full gate on the fixture model (export -> parity -> +strict-validate) in the `export-fixture` job. + Legacy notes: - Old secryst zips (`vocabs.yaml` + single `transformer.onnx`) and the diff --git a/golden/khm-latn-100.jsonl b/golden/khm-latn-100.jsonl new file mode 100644 index 0000000..b56717a --- /dev/null +++ b/golden/khm-latn-100.jsonl @@ -0,0 +1,100 @@ +{"input": "បង្វេច", "tokens": [101, 100, 113, 106, 121, 104, 102, 107], "output": "bangvech"} +{"input": "រោក", "tokens": [117, 114, 110], "output": "rok"} +{"input": "សង់ផ្ទះ", "tokens": [118, 100, 113, 106, 115, 107, 119, 104, 100, 107], "output": "sangphteah"} +{"input": "អប្បបរិមាណ", "tokens": [100, 101, 101, 100, 101, 117, 108, 112, 100, 113], "output": "abbabriman"} +{"input": "ឈឺចាប់", "tokens": [102, 107, 107, 120, 102, 107, 104, 100, 101], "output": "chhucheab"} +{"input": "អសមត្ថភាព", "tokens": [100, 118, 100, 112, 114, 119, 119, 107, 100, 115, 107, 100, 115], "output": "asamotthaphap"} +{"input": "និយាយ", "tokens": [113, 108, 124, 104, 100, 124], "output": "niyeay"} +{"input": "សម្ងំ", "tokens": [118, 100, 112, 113, 106, 114, 112], "output": "samngom"} +{"input": "អូម", "tokens": [100, 120, 112], "output": "aum"} +{"input": "ស្វាយរៀង", "tokens": [118, 121, 100, 124, 117, 108, 104, 113, 106], "output": "svayrieng"} +{"input": "បារមី", "tokens": [101, 100, 117, 112, 108], "output": "barmi"} +{"input": "រនល", "tokens": [117, 114, 113, 111], "output": "ronl"} +{"input": "ធម្មយុត្តិ", "tokens": [119, 107, 114, 112, 112, 114, 124, 114, 120, 119, 119, 104], "output": "thommoyoutte"} +{"input": "ទឹកឃ្មុំ", "tokens": [119, 120, 110, 100, 110, 107, 112, 114, 120, 112], "output": "tukakhmoum"} +{"input": "ភ្លើង", "tokens": [115, 107, 111, 104, 120, 113, 106], "output": "phleung"} +{"input": "ប្រួត", "tokens": [101, 117, 120, 114, 119], "output": "bruot"} +{"input": "វោហារ", "tokens": [121, 114, 107, 104, 100, 117], "output": "vohear"} +{"input": "ច្បាប់", "tokens": [102, 107, 101, 100, 101], "output": "chbab"} +{"input": "ទន្ទឹម", "tokens": [119, 114, 113, 119, 120, 112], "output": "tontum"} +{"input": "បញ្ចស័ក", "tokens": [101, 100, 113, 107, 102, 107, 114, 118, 100, 110], "output": "banhchosak"} +{"input": "ល្បី", "tokens": [111, 101, 108], "output": "lbi"} +{"input": "គុណូបការ", "tokens": [110, 114, 120, 113, 100, 120, 101, 110, 100, 117], "output": "kounaubkar"} +{"input": "បញ្ចុក", "tokens": [101, 100, 113, 107, 102, 107, 114, 120, 110], "output": "banhchouk"} +{"input": "ជ្រាបស្រាប់", "tokens": [102, 107, 117, 104, 100, 101, 100, 118, 117, 100, 101], "output": "chreabasrab"} +{"input": "សៃយ៉ឺន", "tokens": [118, 100, 108, 124, 104, 120, 113], "output": "saiyeun"} +{"input": "លោះខ្ញុំ", "tokens": [111, 114, 107, 110, 107, 113, 107, 114, 112], "output": "lohkhnhom"} +{"input": "ទូទឹម", "tokens": [119, 114, 120, 119, 120, 112], "output": "toutum"} +{"input": "ប្រវត្តិសាស្ត្រ", "tokens": [101, 117, 100, 121, 114, 119, 119, 104, 35, 118, 100, 118, 119, 117], "output": "bravotte sastr"} +{"input": "ខ្លោងទ្វារ", "tokens": [110, 107, 111, 100, 114, 113, 106, 114, 119, 121, 104, 100, 117], "output": "khlaongotvear"} +{"input": "គេជាច្រើន", "tokens": [110, 104, 102, 107, 100, 102, 107, 117, 100, 104, 113], "output": "kechachraen"} +{"input": "បង្កួយ", "tokens": [101, 100, 113, 106, 110, 104, 100, 119], "output": "bangkeat"} +{"input": "ហៃអើ", "tokens": [107, 100, 108, 100, 104], "output": "haiae"} +{"input": "ឆ្កឹះឆ្កៀល", "tokens": [102, 107, 107, 110, 114, 104, 100, 107, 102, 107, 107, 110, 108, 104, 111], "output": "chhkoeahchhkiel"} +{"input": "ទំលាក់", "tokens": [119, 114, 112, 35, 111, 104, 100, 110], "output": "tom leak"} +{"input": "ហ្រ្វង់", "tokens": [107, 117, 121, 114, 113, 106], "output": "hrvong"} +{"input": "រាស់", "tokens": [117, 104, 100, 118, 100], "output": "reasa"} +{"input": "លោភលន់", "tokens": [111, 114, 115, 107, 114, 111, 113], "output": "lopholn"} +{"input": "អរគុណណាស់", "tokens": [100, 117, 110, 114, 120, 113, 113, 100, 118, 100], "output": "arkounnasa"} +{"input": "ដោយហេតុ", "tokens": [103, 100, 114, 124, 107, 104, 119, 114], "output": "daoyheto"} +{"input": "អ្នកចំរៀង", "tokens": [100, 113, 100, 110, 102, 107, 100, 112, 117, 108, 104, 113, 106], "output": "anakchamrieng"} +{"input": "ពោធិសម្ភារ", "tokens": [115, 114, 119, 107, 108, 118, 100, 112, 115, 107, 104, 100, 117], "output": "pothisamphear"} +{"input": "ចូលហ៊ុន", "tokens": [102, 107, 100, 120, 111, 107, 114, 120, 113], "output": "chaulhoun"} +{"input": "បបរ", "tokens": [101, 100, 101, 117], "output": "babr"} +{"input": "រូបលោក", "tokens": [117, 114, 120, 101, 111, 114, 110], "output": "roublok"} +{"input": "ហ៊ុមព័ទ្ធ", "tokens": [107, 114, 120, 112, 115, 119, 119, 107], "output": "houmptth"} +{"input": "ការហាត់ប្រាណ", "tokens": [110, 100, 117, 35, 107, 100, 119, 101, 117, 100, 113], "output": "kar hatbran"} +{"input": "មិនបានជា", "tokens": [112, 108, 113, 101, 100, 113, 102, 107, 104, 100], "output": "minbanchea"} +{"input": "កម្មវិបាក", "tokens": [110, 100, 112, 112, 114, 121, 108, 101, 104, 100, 110], "output": "kammovibeak"} +{"input": "កុំអាល", "tokens": [110, 114, 112, 100, 111], "output": "komal"} +{"input": "សុខយាន", "tokens": [118, 114, 110, 107, 124, 104, 100, 113], "output": "sokhyean"} +{"input": "ទឹកស្លាប់", "tokens": [119, 120, 110, 100, 118, 111, 100, 101], "output": "tukaslab"} +{"input": "អធិការបតី", "tokens": [100, 119, 107, 108, 110, 104, 100, 117, 114, 101, 119, 104, 108], "output": "athikearobtei"} +{"input": "ខ្តត", "tokens": [110, 107, 119, 100, 119], "output": "khtat"} +{"input": "ស្ទេញ", "tokens": [118, 119, 104, 113, 107], "output": "stenh"} +{"input": "ដល់ហើយ", "tokens": [103, 111, 107, 100, 104, 124], "output": "dlhaey"} +{"input": "ក្ដៀប", "tokens": [110, 103, 108, 104, 101], "output": "kdieb"} +{"input": "ភ្លៀងមួយមេ", "tokens": [115, 107, 111, 108, 104, 113, 106, 112, 120, 114, 124, 112, 104], "output": "phliengmuoyme"} +{"input": "សៃយ", "tokens": [118, 100, 108, 124], "output": "saiy"} +{"input": "រាជបល្ល័ង្ក", "tokens": [117, 104, 100, 102, 107, 114, 101, 100, 111, 111, 113, 106, 110], "output": "reachoballngk"} +{"input": "ល្អក់", "tokens": [111, 100, 114, 110], "output": "laok"} +{"input": "ឈ្ងុយឆ្ងាញ់", "tokens": [102, 107, 107, 113, 106, 114, 120, 124, 114, 102, 107, 107, 113, 106, 100, 113, 107], "output": "chhngouyochhnganh"} +{"input": "រ៉ឹង", "tokens": [117, 114, 104, 113, 106], "output": "roeng"} +{"input": "ប្រុងនឹង", "tokens": [101, 117, 114, 113, 106, 113, 120, 113, 106], "output": "brongnung"} +{"input": "ការប្រើប្រាស់", "tokens": [110, 100, 117, 114, 101, 117, 100, 104, 101, 117, 100, 118, 100], "output": "karobraebrasa"} +{"input": "គ្រោះ", "tokens": [110, 117, 114, 107], "output": "kroh"} +{"input": "កលកិច្ច", "tokens": [110, 100, 111, 110, 104, 102, 107, 102, 107], "output": "kalkechch"} +{"input": "ផ្សារផ្សោ", "tokens": [115, 107, 118, 100, 100, 117, 114, 115, 107, 118, 100, 100, 114], "output": "phsaarophsaao"} +{"input": "ស្រករ", "tokens": [118, 117, 100, 110, 117], "output": "srakr"} +{"input": "រត់លឿន", "tokens": [117, 119, 111, 120, 104, 113], "output": "rtluen"} +{"input": "និត្យ", "tokens": [113, 108, 119, 124], "output": "nity"} +{"input": "គ្រប់រូប", "tokens": [110, 117, 114, 101, 117, 114, 120, 101], "output": "krobroub"} +{"input": "ឲ្យទឹក", "tokens": [100, 114, 124, 119, 120, 110], "output": "aoytuk"} +{"input": "ចាប់ត្រី", "tokens": [102, 107, 100, 101, 119, 117, 104, 108], "output": "chabtrei"} +{"input": "ពេជ្រ", "tokens": [115, 104, 102, 107, 117], "output": "pechr"} +{"input": "អគ្គស្នងការ", "tokens": [100, 110, 110, 114, 118, 113, 100, 113, 106, 110, 100, 117], "output": "akkosnangkar"} +{"input": "ដួល", "tokens": [103, 120, 114, 111], "output": "duol"} +{"input": "នៅនឹង", "tokens": [113, 114, 122, 113, 120, 113, 106], "output": "nownung"} +{"input": "ផលកម្ម", "tokens": [115, 107, 100, 111, 114, 110, 100, 112, 112], "output": "phalokamm"} +{"input": "ញាំញី", "tokens": [113, 107, 104, 100, 113, 107, 108], "output": "nheanhi"} +{"input": "ភាតរ", "tokens": [115, 107, 104, 100, 119, 117], "output": "pheatr"} +{"input": "មួយនេះ", "tokens": [112, 120, 114, 124, 113, 108, 107], "output": "muoynih"} +{"input": "ផ្លូវសួន", "tokens": [115, 107, 111, 100, 120, 121, 118, 120, 114, 113], "output": "phlauvsuon"} +{"input": "ឆ្នាំច", "tokens": [102, 107, 107, 113, 100, 102, 107], "output": "chhnach"} +{"input": "ត្នោត", "tokens": [119, 113, 100, 114, 119], "output": "tnaot"} +{"input": "វាង", "tokens": [121, 104, 100, 113, 106], "output": "veang"} +{"input": "រាជបុត្រ", "tokens": [117, 104, 100, 102, 107, 101, 114, 119, 117], "output": "reachbotr"} +{"input": "ស្រមូម", "tokens": [118, 117, 100, 112, 114, 120, 112], "output": "sramoum"} +{"input": "ចិត្តស៊ូ", "tokens": [102, 107, 104, 119, 119, 100, 118, 100, 114, 120], "output": "chettasaou"} +{"input": "អក្សរខម", "tokens": [100, 110, 118, 100, 100, 117, 114, 110, 107, 112], "output": "aksaarokhm"} +{"input": "ទូទឹកកក", "tokens": [119, 114, 120, 119, 120, 110, 100, 110, 110], "output": "toutukakk"} +{"input": "កាត់ឲ្យខ្លី", "tokens": [110, 100, 119, 35, 100, 114, 124, 35, 110, 107, 111, 104, 108], "output": "kat aoy khlei"} +{"input": "យ៉ាងណាក៏ដោយ", "tokens": [124, 104, 100, 113, 106, 113, 100, 35, 110, 100, 103, 100, 114, 124], "output": "yeangna kadaoy"} +{"input": "បំផុត", "tokens": [101, 100, 112, 115, 107, 114, 119], "output": "bamphot"} +{"input": "រកអ្វីប្រៀបពុំបាន", "tokens": [117, 114, 110, 100, 100, 121, 104, 108, 35, 101, 117, 108, 104, 101, 35, 101, 100, 113], "output": "rokaavei brieb ban"} +{"input": "មកយឺត", "tokens": [112, 114, 110, 124, 120, 119], "output": "mokyut"} +{"input": "អប្សរា", "tokens": [100, 101, 118, 100, 100, 117, 100], "output": "absaara"} +{"input": "បិតុច្ឆា", "tokens": [101, 104, 119, 114, 102, 107, 102, 107, 107, 100], "output": "betochchha"} +{"input": "អស់អញ", "tokens": [100, 118, 100, 100, 113, 107], "output": "asaanh"} +{"input": "ចាង", "tokens": [102, 107, 100, 113, 106], "output": "chang"} +{"input": "មិនអស់ចិត្ត", "tokens": [112, 108, 113, 114, 100, 118, 100, 102, 107, 104, 119, 119], "output": "minoasachett"} diff --git a/models/heb-diac/heb-diac-1.0.README.md b/models/heb-diac/heb-diac-1.0.README.md new file mode 100644 index 0000000..32bd344 --- /dev/null +++ b/models/heb-diac/heb-diac-1.0.README.md @@ -0,0 +1,42 @@ +# heb-diac-1.0 + +Hebrew diacritization (adds nikud). Byte-level seq2seq (ByT5-base): +the tokenizer is raw UTF-8 bytes (pad=0, EOS=1) — no vocab files. +IMF v1 artifact; format spec: interscript/ml-models docs/imf-v1.md. + +- decoder: kv greedy (plain fallback included in the zip) +- metrics: greedy DER 29.0% (the v1 runtime path); beam=4 DER 17.46% + (reference quality — beam search is not in v1 runtimes) — + rababa/docs/RESULTS.md#hebrew-diacritization +- trained from: rababa train_hebrew_seeds.py s43 run-001 + (rababa-checkpoints:/rababa_hebrew_byt5_s43/run-001/best) +- license: BSD-3-Clause + +## Usage + +Ruby (secryst gem, the Ruby binding of interscript-ml): + +```ruby +require "secryst" +translator = Secryst::Translator.new(model: "heb-diac-1.0") +translator.translate("שלום") +``` + +TypeScript (@interscript/ml): + +```ts +import { loadModel } from "@interscript/ml"; +const model = await loadModel("heb-diac-1.0"); +await model.translate("שלום"); +``` + +Python (interscript-ml): + +```python +from interscript_ml import Model +model = Model.load("heb-diac-1.0") +model.translate("שלום") +``` + +All three runtimes verify the sha256 of every ONNX member in this zip +against metadata.yaml before loading. diff --git a/models/heb-diac/heb-diac-1.0.metadata.yaml b/models/heb-diac/heb-diac-1.0.metadata.yaml new file mode 100644 index 0000000..085967c --- /dev/null +++ b/models/heb-diac/heb-diac-1.0.metadata.yaml @@ -0,0 +1,26 @@ +format: imf-v1 +id: heb-diac-1.0 +task: diacritization +source_script: Hebr +target: Hebr +tokenizer: bytes +opset: 14 +decoder: kv +precision: fp32 +license: BSD-3-Clause +trained_from: >- + rababa train_hebrew_seeds.py s43 run-001; checkpoint + rababa-checkpoints:/rababa_hebrew_byt5_s43/run-001/best +metrics: + - name: der_greedy + value: 29.0 + protocol: >- + beam=1 greedy decode (the v1 runtime path); Nakdimon test split, + 5,095 examples; ByT5-base s43 + source: rababa/docs/RESULTS.md#hebrew-diacritization + - name: der_beam4 + value: 17.46 + protocol: >- + beam=4 standard decode (reference quality; beam search is not in + v1 runtimes); Nakdimon test split, 5,095 examples; ByT5-base s43 + source: rababa/docs/RESULTS.md#hebrew-diacritization diff --git a/models/khm-latn/khm-latn-1.0.README.md b/models/khm-latn/khm-latn-1.0.README.md index f48713f..a5b5c28 100644 --- a/models/khm-latn/khm-latn-1.0.README.md +++ b/models/khm-latn/khm-latn-1.0.README.md @@ -5,7 +5,7 @@ the tokenizer is raw UTF-8 bytes (pad=0, EOS=1) — no vocab files. IMF v1 artifact; format spec: interscript/ml-models docs/imf-v1.md. - precision: fp16 (mixed: LayerNorm parameters in fp32) -- decoder: plain greedy (KV-cache variant ships with the WO02 export) +- decoder: kv greedy (plain fallback included in the zip) - metrics: CER 27.42 / EM 59.66 on 895 held-out pairs — secryst/docs/RESULTS.md#khmer-transliteration-2026-08-14 - trained from: secryst train_khmer_byt5.py run-001 diff --git a/models/khm-latn/khm-latn-1.0.metadata.yaml b/models/khm-latn/khm-latn-1.0.metadata.yaml new file mode 100644 index 0000000..20010c2 --- /dev/null +++ b/models/khm-latn/khm-latn-1.0.metadata.yaml @@ -0,0 +1,26 @@ +format: imf-v1 +id: khm-latn-1.0 +task: translit +source_script: Khmr +target: Latn +tokenizer: bytes +opset: 14 +decoder: kv +precision: fp32 +license: BSD-3-Clause +trained_from: >- + secryst train_khmer_byt5.py run-001; checkpoint + secryst-checkpoints:/khmer_byt5/run-001/best +metrics: + - name: cer + value: 27.42 + protocol: >- + greedy decode; 895 held-out pairs; split 16,120/895/895 seed 42; + ByT5-small early stop @ep15 + source: secryst/docs/RESULTS.md#khmer-transliteration-2026-08-14 + - name: em + value: 59.66 + protocol: >- + greedy decode; 895 held-out pairs; split 16,120/895/895 seed 42; + ByT5-small early stop @ep15 + source: secryst/docs/RESULTS.md#khmer-transliteration-2026-08-14 diff --git a/models/urd-diac/urd-diac-1.0.README.md b/models/urd-diac/urd-diac-1.0.README.md new file mode 100644 index 0000000..a698138 --- /dev/null +++ b/models/urd-diac/urd-diac-1.0.README.md @@ -0,0 +1,41 @@ +# urd-diac-1.0 + +Urdu diacritization (adds haraqat). Byte-level seq2seq (ByT5-small): +the tokenizer is raw UTF-8 bytes (pad=0, EOS=1) — no vocab files. +IMF v1 artifact; format spec: interscript/ml-models docs/imf-v1.md. + +- decoder: kv greedy (plain fallback included in the zip) +- metrics: CER 3.74 on 11,940 held-out — + rababa-urdu/docs/RESULTS.md#diacritization-urdu-text--text--haraqat +- trained from: rababa-urdu modal_app_diacrit.py run-001 + (urdu-diacrit-checkpoints:/urdu_diacrit/run-001/best) +- license: BSD-3-Clause + +## Usage + +Ruby (secryst gem, the Ruby binding of interscript-ml): + +```ruby +require "secryst" +translator = Secryst::Translator.new(model: "urd-diac-1.0") +translator.translate("اردو") +``` + +TypeScript (@interscript/ml): + +```ts +import { loadModel } from "@interscript/ml"; +const model = await loadModel("urd-diac-1.0"); +await model.translate("اردو"); +``` + +Python (interscript-ml): + +```python +from interscript_ml import Model +model = Model.load("urd-diac-1.0") +model.translate("اردو") +``` + +All three runtimes verify the sha256 of every ONNX member in this zip +against metadata.yaml before loading. diff --git a/models/urd-diac/urd-diac-1.0.metadata.yaml b/models/urd-diac/urd-diac-1.0.metadata.yaml new file mode 100644 index 0000000..55aae6d --- /dev/null +++ b/models/urd-diac/urd-diac-1.0.metadata.yaml @@ -0,0 +1,20 @@ +format: imf-v1 +id: urd-diac-1.0 +task: diacritization +source_script: Arab +target: Arab +tokenizer: bytes +opset: 14 +decoder: kv +precision: fp32 +license: BSD-3-Clause +trained_from: >- + rababa-urdu modal_app_diacrit.py run-001; checkpoint + urdu-diacrit-checkpoints:/urdu_diacrit/run-001/best +metrics: + - name: cer + value: 3.74 + protocol: >- + greedy decode; 11,940 held-out; labels derived IPA->haraqat + (deterministic conversion, 597K pairs); ByT5-small, 2 epochs + source: rababa-urdu/docs/RESULTS.md#diacritization-urdu-text--text--haraqat diff --git a/models/urd-g2p/urd-g2p-1.0.README.md b/models/urd-g2p/urd-g2p-1.0.README.md new file mode 100644 index 0000000..6a922d1 --- /dev/null +++ b/models/urd-g2p/urd-g2p-1.0.README.md @@ -0,0 +1,41 @@ +# urd-g2p-1.0 + +Urdu → IPA grapheme-to-phoneme conversion. Byte-level seq2seq +(ByT5-small): the tokenizer is raw UTF-8 bytes (pad=0, EOS=1) — no vocab +files. IMF v1 artifact; format spec: interscript/ml-models docs/imf-v1.md. + +- decoder: kv greedy (plain fallback included in the zip) +- metrics: CER 14.77 / EM 33.6 on 12,699 held-out words — + rababa-urdu/docs/RESULTS.md#g2p-urdu-text--ipa +- trained from: rababa-urdu modal_app.py run-001 + (urdu-g2p-checkpoints:/urdu_g2p/run-001/best) +- license: BSD-3-Clause + +## Usage + +Ruby (secryst gem, the Ruby binding of interscript-ml): + +```ruby +require "secryst" +translator = Secryst::Translator.new(model: "urd-g2p-1.0") +translator.translate("اردو") +``` + +TypeScript (@interscript/ml): + +```ts +import { loadModel } from "@interscript/ml"; +const model = await loadModel("urd-g2p-1.0"); +await model.translate("اردو"); +``` + +Python (interscript-ml): + +```python +from interscript_ml import Model +model = Model.load("urd-g2p-1.0") +model.translate("اردو") +``` + +All three runtimes verify the sha256 of every ONNX member in this zip +against metadata.yaml before loading. diff --git a/models/urd-g2p/urd-g2p-1.0.metadata.yaml b/models/urd-g2p/urd-g2p-1.0.metadata.yaml new file mode 100644 index 0000000..d345e0a --- /dev/null +++ b/models/urd-g2p/urd-g2p-1.0.metadata.yaml @@ -0,0 +1,26 @@ +format: imf-v1 +id: urd-g2p-1.0 +task: g2p +source_script: Arab +target: IPA +tokenizer: bytes +opset: 14 +decoder: kv +precision: fp32 +license: BSD-3-Clause +trained_from: >- + rababa-urdu modal_app.py run-001; checkpoint + urdu-g2p-checkpoints:/urdu_g2p/run-001/best +metrics: + - name: cer + value: 14.77 + protocol: >- + greedy decode; 12,699 held-out words from the 635K humair025 + urdu-g2p dictionary; ByT5-small + source: rababa-urdu/docs/RESULTS.md#g2p-urdu-text--ipa + - name: em + value: 33.6 + protocol: >- + greedy decode; 12,699 held-out words from the 635K humair025 + urdu-g2p dictionary; ByT5-small + source: rababa-urdu/docs/RESULTS.md#g2p-urdu-text--ipa diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py new file mode 100644 index 0000000..5daf206 --- /dev/null +++ b/src/gpu/modal_export.py @@ -0,0 +1,229 @@ +"""Modal app: export IMF v1 zips from checkpoints on Modal volumes (WO02), +then gate them with the WO03 parity check — all CPU, never competing +with A100 training. + + modal run --detach src/gpu/modal_export.py --model khm-latn + modal run --detach src/gpu/modal_export.py::parity --model khm-latn + +Watchdog (server evictions happen; both steps are idempotent — retries +are the resume mechanism, and each model's zips are written atomically): + + until modal run --detach src/gpu/modal_export.py --model khm-latn; do sleep 60; done + +Zips land on the secryst-models volume under /imf//; parity is +written into the zip in place (a zip is only shippable strict-validated). +Versions are pinned to the ones the export was verified against locally +(transformers 5.15 breaks T5 tracing with "multiple values for +use_cache"). +""" + +from __future__ import annotations + +from pathlib import Path + +import modal + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent + +IMAGE = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch==2.12.1", + "transformers==5.14.1", + "onnx==1.22.0", + "onnxruntime==1.23.2", + "pyyaml>=6.0", + ) + .add_local_dir(str(REPO_ROOT), "/root/ml-models", copy=True) + .workdir("/root/ml-models") +) + +CHECKPOINT_VOLUMES = { + "/volumes/secryst-checkpoints": modal.Volume.from_name("secryst-checkpoints"), + "/volumes/urdu-g2p-checkpoints": modal.Volume.from_name("urdu-g2p-checkpoints"), + "/volumes/urdu-diacrit-checkpoints": modal.Volume.from_name( + "urdu-diacrit-checkpoints" + ), + "/volumes/rababa-checkpoints": modal.Volume.from_name("rababa-checkpoints"), +} + +DATASET_VOLUMES = { + "/datasets/rababa": modal.Volume.from_name("rababa-datasets"), + "/datasets/secryst": modal.Volume.from_name("secryst-datasets"), + "/datasets/urdu-g2p": modal.Volume.from_name("urdu-g2p-datasets"), + "/datasets/urdu-diacrit": modal.Volume.from_name("urdu-diacrit-datasets"), +} + +MODELS_VOLUME = modal.Volume.from_name("secryst-models") + +MODELS: dict[str, dict[str, str]] = { + "khm-latn": { + "volume": "/volumes/secryst-checkpoints", + "checkpoint": "khmer_byt5/run-001/best", + "metadata": "models/khm-latn/khm-latn-1.0.metadata.yaml", + "readme": "models/khm-latn/khm-latn-1.0.README.md", + "test_volume": "/datasets/secryst", + "test_data": "khmer-translit/test.jsonl", + "probe": "ភាសា", + }, + "urd-g2p": { + "volume": "/volumes/urdu-g2p-checkpoints", + "checkpoint": "urdu_g2p/run-001/best", + "metadata": "models/urd-g2p/urd-g2p-1.0.metadata.yaml", + "readme": "models/urd-g2p/urd-g2p-1.0.README.md", + "test_volume": "/datasets/urdu-g2p", + "test_data": "urdu-g2p/test.jsonl", + "probe": "اردو", + }, + "heb-diac": { + "volume": "/volumes/rababa-checkpoints", + "checkpoint": "rababa_hebrew_byt5_s43/run-001/best", + "metadata": "models/heb-diac/heb-diac-1.0.metadata.yaml", + "readme": "models/heb-diac/heb-diac-1.0.README.md", + "test_volume": "/datasets/rababa", + "test_data": "nakdimon/test-imf.jsonl", + "probe": "שלום", + }, + "urd-diac": { + "volume": "/volumes/urdu-diacrit-checkpoints", + "checkpoint": "urdu_diacrit/run-001/best", + "metadata": "models/urd-diac/urd-diac-1.0.metadata.yaml", + "readme": "models/urd-diac/urd-diac-1.0.README.md", + "test_volume": "/datasets/urdu-diacrit", + "test_data": "urdu-diacrit/test.jsonl", + "probe": "اردو", + }, +} + +app = modal.App("interscript-ml-export", image=IMAGE) + + +def _load_pairs(path: Path) -> list[tuple[str, str]]: + import json + + pairs: list[tuple[str, str]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + row = json.loads(line) + if isinstance(row, dict): + pairs.append( + ( + row.get("input", row.get("src", "")), + row.get("target", row.get("tgt", row.get("gold", ""))), + ) + ) + else: + pairs.append((row[0], row[1] if len(row) > 1 else "")) + return pairs + + +@app.function( + cpu=8, + memory=32 * 1024, + timeout=2 * 3600, + volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, +) +def export_model(model_id: str, precisions: list[str]) -> dict[str, str]: + import sys + + sys.path.insert(0, "/root/ml-models/src") + + spec = MODELS[model_id] + checkpoint = Path(spec["volume"]) / spec["checkpoint"] + metadata_path = Path("/root/ml-models") / spec["metadata"] + readme_path = Path("/root/ml-models") / spec["readme"] + + from imf.export import export_zips, load_byte_seq2seq, onnx_greedy_kv + from imf.validator import validate_zip + + model = load_byte_seq2seq(checkpoint) + out_dir = Path("/outputs/imf") / model_id + zips = export_zips( + model, + metadata_path, + readme_path.read_text(encoding="utf-8"), + out_dir, + precisions=tuple(precisions), + ) + MODELS_VOLUME.commit() + + import zipfile + + import onnxruntime as ort + + report: dict[str, str] = {} + for z in zips: + result = validate_zip(z) + if not result.ok: + raise RuntimeError(f"{z.name} failed validation: {result.errors}") + with zipfile.ZipFile(z) as zf: + zf.extract("encoder.onnx", "/tmp/check") + zf.extract("decoder-kv.onnx", "/tmp/check") + enc = ort.InferenceSession( + "/tmp/check/encoder.onnx", providers=["CPUExecutionProvider"] + ) + kv = ort.InferenceSession( + "/tmp/check/decoder-kv.onnx", providers=["CPUExecutionProvider"] + ) + tokens = onnx_greedy_kv(enc, kv, spec["probe"], max_len=32) + report[z.name] = f"{z.stat().st_size} bytes, probe -> {len(tokens)} tokens" + return report + + +@app.function( + cpu=8, + memory=32 * 1024, + timeout=2 * 3600, + volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, +) +def parity_model(model_id: str, precisions: list[str], limit: int = 0) -> dict[str, str]: + """WO03 gate on Modal: torch reference vs ONNX decode over the test + split; writes the parity block into each zip (strict gate enforced).""" + import sys + + sys.path.insert(0, "/root/ml-models/src") + + spec = MODELS[model_id] + checkpoint = Path(spec["volume"]) / spec["checkpoint"] + test_path = Path(spec["test_volume"]) / spec["test_data"] + + from imf.export import load_byte_seq2seq + from imf.parity import reference_decode, run_parity, write_parity + + model = load_byte_seq2seq(checkpoint) + pairs = _load_pairs(test_path) + if limit: + pairs = pairs[:limit] + + reference = reference_decode(model, [src for src, _ in pairs], max_len=128) + + out_dir = Path("/outputs/imf") / model_id + reports: dict[str, str] = {} + for precision in precisions: + zip_path = out_dir / f"{model_id}-1.0-{precision}.zip" + report = run_parity(model, zip_path, pairs, max_len=128, reference=reference) + reports[precision] = ( + f"samples={report.samples} cer_ref={report.cer_reference}pp " + f"cer_onnx={report.cer_onnx}pp delta={report.cer_delta}pp " + f"mismatches={report.token_mismatches}" + ) + if not report.passed: + raise RuntimeError(f"parity gate FAILED for {zip_path.name}") + write_parity(zip_path, report) + MODELS_VOLUME.commit() + return reports + + +@app.local_entrypoint() +def main(model: str, precisions: str = "fp32,fp16,int8") -> None: + report = export_model.remote(model, precisions.split(",")) + for name, status in report.items(): + print(f"{name}: {status}") + + +@app.local_entrypoint() +def parity(model: str, precisions: str = "fp32,fp16,int8", limit: int = 0) -> None: + reports = parity_model.remote(model, precisions.split(","), limit) + for precision, status in reports.items(): + print(f"{model} [{precision}] {status}") diff --git a/src/imf/cli.py b/src/imf/cli.py index 82e96b6..85ae220 100644 --- a/src/imf/cli.py +++ b/src/imf/cli.py @@ -81,6 +81,69 @@ def _default_readme(metadata: ModelMetadata) -> str: ) +def _load_pairs(path: Path) -> list[tuple[str, str]]: + import json + + pairs: list[tuple[str, str]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + row = json.loads(line) + if isinstance(row, dict): + pairs.append( + ( + row.get("input", row.get("src", "")), + row.get("target", row.get("tgt", row.get("gold", ""))), + ) + ) + else: + pairs.append((row[0], row[1] if len(row) > 1 else "")) + return pairs + + +def _cmd_parity(args: argparse.Namespace) -> int: + from imf.export import load_byte_seq2seq + from imf.parity import run_parity, write_parity + + model = load_byte_seq2seq(args.checkpoint) + pairs = _load_pairs(args.test_data) + if args.limit: + pairs = pairs[: args.limit] + report = run_parity(model, args.zip, pairs, max_len=args.max_len) + print( + f"samples={report.samples} cer_ref={report.cer_reference}pp " + f"cer_onnx={report.cer_onnx}pp delta={report.cer_delta}pp " + f"token_mismatches={report.token_mismatches}" + ) + if not report.passed: + print("error: parity gate FAILED", file=sys.stderr) + return 1 + write_parity(args.zip, report) + print(f"parity written into {args.zip} (strict validation passed)") + return 0 + + +def _cmd_golden(args: argparse.Namespace) -> int: + import json + + from imf.parity import write_golden + + inputs: list[str] = [] + for line in args.inputs.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + row = json.loads(line) + if isinstance(row, dict): + inputs.append(row.get("input", row.get("src", row.get("text", "")))) + elif isinstance(row, str): + inputs.append(row) + else: + inputs.append(row[0]) + out = write_golden(args.zip, inputs, args.out, max_len=args.max_len) + print(f"wrote {len(inputs)} golden cases to {out}") + return 0 + + def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="imf", description=__doc__) sub = parser.add_subparsers(dest="command", required=True) @@ -109,6 +172,32 @@ def build_parser() -> argparse.ArgumentParser: p_pack.add_argument("--out", required=True, type=Path) p_pack.set_defaults(func=_cmd_pack) + p_parity = sub.add_parser( + "parity", help="WO03 gate: ONNX vs torch reference, write parity into the zip" + ) + p_parity.add_argument("zip", type=Path) + p_parity.add_argument( + "--checkpoint", required=True, type=Path, help="HF checkpoint dir (reference)" + ) + p_parity.add_argument( + "--test-data", required=True, type=Path, + help="JSONL: {input, target} pairs (or [src, gold] arrays)", + ) + p_parity.add_argument("--limit", type=int, help="cap sample count") + p_parity.add_argument("--max-len", type=int, default=256) + p_parity.set_defaults(func=_cmd_parity) + + p_golden = sub.add_parser( + "golden", help="emit the cross-runtime golden JSONL from ONNX decode" + ) + p_golden.add_argument("zip", type=Path) + p_golden.add_argument( + "--inputs", required=True, type=Path, help="JSONL of input strings" + ) + p_golden.add_argument("--out", required=True, type=Path) + p_golden.add_argument("--max-len", type=int, default=256) + p_golden.set_defaults(func=_cmd_golden) + return parser diff --git a/src/imf/export.py b/src/imf/export.py new file mode 100644 index 0000000..35ad91d --- /dev/null +++ b/src/imf/export.py @@ -0,0 +1,395 @@ +"""ONNX export of byte-level (ByT5/T5) checkpoints into IMF v1 zips. + +Generalizes secryst PR #44's scripts/export_onnx_byt5.py: + +- encoder.onnx input_ids -> last_hidden_state +- decoder.onnx input_ids, encoder_hidden_states -> logits (fallback) +- decoder-kv.onnx + past_* inputs / present_* outputs (default artifact) +- opset pinned to 14 (the Ruby onnxruntime gem's bundled ORT is old) +- fp16 (keep IO fp32) and int8 (dynamic quantization) variants +- fixture mode: a tiny random T5 for CI, via ``make_fixture_checkpoint`` + +The KV graph caches self-attention only (``past_key_i`` / ``past_value_i`` +inputs, ``present_key_i`` / ``present_value_i`` outputs). Cross-attention +K/V are a deterministic projection of ``encoder_hidden_states``, so they +are recomputed each step instead of being cached — this keeps the +attention-mask length consistent for any past length and spares runtimes +all cross-cache bookkeeping. Step 0 feeds zero-length self pasts. +""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path + +from imf.pack import pack_zip +from imf.schema import ModelMetadata + +OPSET = 14 +GRAPH_NAMES = ("encoder.onnx", "decoder.onnx", "decoder-kv.onnx") + +# Canonical ByT5 byte table (google/byt5): pad=0, eos=1, unk=2, and +# every UTF-8 byte b is token id b + 3, up to 258. The vocab is 384-wide; +# ids > 258 are unused in practice. "tokenizer: bytes" in IMF metadata +# means THIS fixed table — no vocab files, but ids are NOT raw byte values +# (feeding text.bytes directly silently produces garbage). +BYTE_OFFSET = 3 +PAD_ID = 0 +EOS_ID = 1 + + +def encode_bytes(text: str) -> list[int]: + """Canonical byte-level tokenization: byte ids + trailing EOS.""" + return [b + BYTE_OFFSET for b in text.encode("utf-8")] + [EOS_ID] + + +def load_byte_seq2seq(checkpoint_dir: Path | str): + """Load a T5-family checkpoint in eager attention (export-safe).""" + from transformers import AutoModelForSeq2SeqLM + + model = AutoModelForSeq2SeqLM.from_pretrained( + checkpoint_dir, attn_implementation="eager" + ).eval() + return model + + +def make_fixture_checkpoint(out_dir: Path | str, seed: int = 42) -> Path: + """Tiny random T5 with a byte-sized vocab, for CI and --fixture runs.""" + import torch + from transformers import T5Config, T5ForConditionalGeneration + + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + torch.manual_seed(seed) + config = T5Config( + vocab_size=384, + d_model=32, + d_kv=8, + d_ff=64, + num_layers=2, + num_heads=4, + decoder_start_token_id=0, + feed_forward_proj="relu", + relative_attention_num_buckets=8, + relative_attention_max_distance=16, + tie_word_embeddings=False, + ) + config._attn_implementation = "eager" + model = T5ForConditionalGeneration(config).eval() + model.save_pretrained(out_dir) + return out_dir + + +def _decoder_plain(model): + import torch.nn as nn + + class DecoderPlain(nn.Module): + def __init__(self, model): + super().__init__() + self.decoder = model.get_decoder() + self.lm_head = model.lm_head + self.scale = model.config.d_model ** -0.5 + + def forward(self, input_ids, encoder_hidden_states): + hidden = self.decoder( + input_ids=input_ids, encoder_hidden_states=encoder_hidden_states + )[0] + return self.lm_head(hidden * self.scale) + + return DecoderPlain(model) + + +def _decoder_kv(model): + import torch.nn as nn + from transformers.cache_utils import DynamicCache, EncoderDecoderCache + + num_layers = model.config.num_decoder_layers or model.config.num_layers + + class DecoderKV(nn.Module): + def __init__(self, model): + super().__init__() + self.decoder = model.get_decoder() + self.lm_head = model.lm_head + self.scale = model.config.d_model ** -0.5 + + def forward(self, input_ids, encoder_hidden_states, *pasts): + # is_updated must stay unset: EncoderDecoderCache.__init__ bakes + # it as a Python bool from cross-cache lengths, and a True would + # trace the "reuse cross K/V" branch, which mismatches the + # attention mask for any fed cross-past length. Cross K/V are a + # deterministic projection of encoder_hidden_states anyway, so + # the graph recomputes them each step and never caches them. + cache = EncoderDecoderCache(DynamicCache(), DynamicCache()) + for i in range(num_layers): + cache.self_attention_cache.update(pasts[2 * i], pasts[2 * i + 1], i) + hidden = self.decoder( + input_ids=input_ids, + encoder_hidden_states=encoder_hidden_states, + past_key_values=cache, + use_cache=True, + )[0] + logits = self.lm_head(hidden * self.scale) + outputs = [logits] + for i in range(num_layers): + outputs.append(cache.self_attention_cache.layers[i].keys) + outputs.append(cache.self_attention_cache.layers[i].values) + return tuple(outputs) + + return DecoderKV(model) + + +def _kv_io_names(num_layers: int) -> tuple[list[str], list[str]]: + inputs, outputs = [], ["logits"] + for i in range(num_layers): + inputs += [f"past_key_{i}", f"past_value_{i}"] + outputs += [f"present_key_{i}", f"present_value_{i}"] + return inputs, outputs + + +def _sample_pasts(model, encoder_hidden_states): + """Harvest one decoder step's self-attention cache for shape examples.""" + import torch + + out = model.get_decoder()( + input_ids=torch.tensor([[0]]), + encoder_hidden_states=encoder_hidden_states, + use_cache=True, + ) + pasts = [] + for layer in out[1].self_attention_cache.layers: + pasts += [layer.keys.clone(), layer.values.clone()] + return pasts + + +def export_graphs(model, out_dir: Path | str) -> dict[str, Path]: + """Export encoder + plain decoder + KV decoder as fp32 ONNX, opset 14. + + Traces a deepcopy: TorchScript tracing leaves the traced module's + runtime behavior subtly altered (observed empirically), and callers + (e.g. the WO03 parity harness) need the reference model pristine. + """ + import copy + + import torch + + model = copy.deepcopy(model) + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + ids = torch.tensor([[104, 101]]) # "he" + with torch.no_grad(): + hidden = model.get_encoder()(input_ids=ids)[0] + + paths = {} + + torch.onnx.export( + model.get_encoder(), + (ids,), + out_dir / "encoder.onnx", + input_names=["input_ids"], + output_names=["last_hidden_state"], + opset_version=OPSET, + dynamo=False, + dynamic_axes={ + "input_ids": {0: "batch", 1: "seq"}, + "last_hidden_state": {0: "batch", 1: "seq"}, + }, + ) + paths["encoder.onnx"] = out_dir / "encoder.onnx" + + torch.onnx.export( + _decoder_plain(model), + (torch.tensor([[0]]), hidden), + out_dir / "decoder.onnx", + input_names=["input_ids", "encoder_hidden_states"], + output_names=["logits"], + opset_version=OPSET, + dynamo=False, + dynamic_axes={ + "input_ids": {0: "batch", 1: "seq"}, + "encoder_hidden_states": {0: "batch", 1: "seq"}, + "logits": {0: "batch", 1: "seq"}, + }, + ) + paths["decoder.onnx"] = out_dir / "decoder.onnx" + + num_layers = model.config.num_decoder_layers or model.config.num_layers + pasts = _sample_pasts(model, hidden) + kv_inputs = ["input_ids", "encoder_hidden_states"] + _kv_io_names(num_layers)[0] + kv_outputs = _kv_io_names(num_layers)[1] + dynamic = { + "input_ids": {0: "batch", 1: "cur_seq"}, + "encoder_hidden_states": {0: "batch", 1: "enc_seq"}, + "logits": {0: "batch", 1: "cur_seq"}, + } + for name in kv_inputs[2:]: + dynamic[name] = {0: "batch", 2: "past_seq"} + for name in kv_outputs[1:]: + dynamic[name] = {0: "batch", 2: "present_seq"} + + with torch.no_grad(): + torch.onnx.export( + _decoder_kv(model), + (torch.tensor([[0]]), hidden, *pasts), + out_dir / "decoder-kv.onnx", + input_names=kv_inputs, + output_names=kv_outputs, + opset_version=OPSET, + dynamo=False, + dynamic_axes=dynamic, + ) + paths["decoder-kv.onnx"] = out_dir / "decoder-kv.onnx" + return paths + + +def convert_fp16(model): + """A fp16 copy of the model for torch-native half export. + + The onnxruntime float16 CONVERTER is not usable here: on real ByT5 + checkpoints it produces all-zero encoder hiddens (found 2026-08-16, + khm-latn — 1939pp CER); exporting the torch model under .half() is + exact on gold pairs. Graph IO becomes float16 (input_ids stay int64); + runtimes read dtypes from the session, and _zero_pasts follows them. + """ + import copy + + return copy.deepcopy(model).half() + + +def quantize_int8(src: Path | str, dst: Path | str) -> Path: + """fp32 -> dynamically quantized int8 (MatMul weights QInt8). + + MatMul-only: quantizing other ops inserts precision casts that break + ORT's session-time SimplifiedLayerNormFusion, and preprocessing the + graph (quant_pre_process) pins concrete example shapes into + DynamicQuantizeLinear buffers. + """ + from onnxruntime.quantization import QuantType, quantize_dynamic + + quantize_dynamic( + str(src), + str(dst), + weight_type=QuantType.QInt8, + op_types_to_quantize=["MatMul"], + ) + return Path(dst) + + +def onnx_greedy_plain(encoder_sess, decoder_sess, text: str, max_len: int = 256) -> list[int]: + """Greedy decode over ONNX sessions (plain decoder). Self-check helper. + + Returns generated token ids (the vocab is 384-wide; only a trained + byte-level model reliably stays < 256).""" + import numpy as np + + ids = np.array([encode_bytes(text)], dtype=np.int64) + if ids.shape[1] == 1: + return [] + hidden = encoder_sess.run(None, {"input_ids": ids})[0] + dec_ids = np.array([[PAD_ID]], dtype=np.int64) + generated: list[int] = [] + for _ in range(max_len): + logits = decoder_sess.run( + None, {"input_ids": dec_ids, "encoder_hidden_states": hidden} + )[0] + nxt = int(np.argmax(logits[0, -1])) + if nxt == EOS_ID: + break + generated.append(nxt) + dec_ids = np.concatenate([dec_ids, np.array([[nxt]], dtype=np.int64)], axis=1) + return generated + + +def _zero_pasts(kv_sess) -> dict[str, object]: + """Zero-length past inputs for step 0, shapes from session metadata.""" + import numpy as np + + pasts = {} + for meta in kv_sess.get_inputs(): + if not meta.name.startswith("past_"): + continue + shape = meta.shape # [batch, heads, past_seq, d_kv] with str dynamic dims + heads = shape[1] if isinstance(shape[1], int) else 4 + d_kv = shape[3] if isinstance(shape[3], int) else 8 + dtype = np.float16 if meta.type == "tensor(float16)" else np.float32 + pasts[meta.name] = np.zeros((1, heads, 0, d_kv), dtype=dtype) + return pasts + + +def onnx_greedy_kv(encoder_sess, kv_sess, text: str, max_len: int = 256) -> list[int]: + """Greedy decode over ONNX sessions (KV decoder). Self-check helper.""" + import numpy as np + + ids = np.array([encode_bytes(text)], dtype=np.int64) + if ids.shape[1] == 1: + return [] + hidden = encoder_sess.run(None, {"input_ids": ids})[0] + out_names = [o.name for o in kv_sess.get_outputs()] + pasts = _zero_pasts(kv_sess) + cur = np.array([[0]], dtype=np.int64) + generated: list[int] = [] + for _ in range(max_len): + out = kv_sess.run(None, {"input_ids": cur, "encoder_hidden_states": hidden, **pasts}) + results = dict(zip(out_names, out, strict=True)) + nxt = int(np.argmax(results["logits"][0, -1])) + if nxt == EOS_ID: + break + generated.append(nxt) + pasts = { + name: results[name.replace("past_", "present_", 1)] + for name in pasts + } + cur = np.array([[nxt]], dtype=np.int64) + return generated + + +def export_zips( + model, + metadata_path: Path | str, + readme: str, + out_dir: Path | str, + precisions: tuple[str, ...] = ("fp32", "fp16", "int8"), +) -> list[Path]: + """Full pipeline: loaded model -> fp32 graphs -> precision variants -> IMF zips.""" + import tempfile + + metadata = ModelMetadata.from_yaml(Path(metadata_path).read_text(encoding="utf-8")) + if metadata.decoder != "kv": + raise ValueError("WO02 exports declare decoder: kv (plain is the fallback)") + + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + zips = [] + + with tempfile.TemporaryDirectory() as tmp: + tmp = Path(tmp) + graphs = export_graphs(model, tmp / "graphs") + graphs_16 = ( + export_graphs(convert_fp16(model), tmp / "graphs-fp16") + if "fp16" in precisions + else {} + ) + + for precision in precisions: + variant_dir = tmp / precision + variant_dir.mkdir() + sources = graphs if precision != "fp16" else graphs_16 + for name, src in sources.items(): + dst = variant_dir / name + if precision == "fp32" or precision == "fp16": + dst.write_bytes(src.read_bytes()) + elif precision == "int8": + quantize_int8(graphs[name], dst) + else: + raise ValueError(f"unknown precision {precision!r}") + meta = replace(metadata, precision=precision) + zips.append( + pack_zip( + variant_dir, + meta, + readme, + out_dir / f"{metadata.id}-{precision}.zip", + ) + ) + return zips + diff --git a/src/imf/parity.py b/src/imf/parity.py new file mode 100644 index 0000000..b8f2fa2 --- /dev/null +++ b/src/imf/parity.py @@ -0,0 +1,215 @@ +"""WO03 parity gate: ONNX greedy vs the torch reference, CER delta <= 0.2pp. + +The reference is the transformers decoder loop itself (the exact math the +export wraps) rather than ``model.generate`` — generate's behavior is +config-dependent (eos/start-token defaults) and none of it is implemented +by the runtimes. Comparing against the module-level math catches exactly +what export bugs can break. + +``write_parity`` rewrites the parity block inside an existing zip +(metadata.yaml is never sha256-covered, graphs are untouched) and then +requires the zip to pass strict validation — the release gate. +""" + +from __future__ import annotations + +import json +import zipfile +from dataclasses import dataclass +from pathlib import Path + +from framework.evaluator import char_error_rate +from imf.export import BYTE_OFFSET, EOS_ID, PAD_ID, encode_bytes, onnx_greedy_kv +from imf.schema import ModelMetadata, Parity + + +@dataclass(frozen=True) +class ParityReport: + samples: int + cer_reference: float # percentage points + cer_onnx: float + cer_delta: float + token_mismatches: int + + @property + def passed(self) -> bool: + return ( + self.cer_delta <= Parity.MAX_CER_DELTA + and self.samples >= Parity.MIN_SAMPLES + ) + + +def _torch_greedy_tokens(model, text: str, max_len: int) -> list[int]: + import torch + + ids = torch.tensor([encode_bytes(text)], dtype=torch.long) + if ids.shape[1] == 1: + return [] + enc = model.get_encoder()(input_ids=ids)[0] + dec_ids = torch.tensor([[PAD_ID]], dtype=torch.long) + outs: list[int] = [] + for _ in range(max_len): + hidden = model.get_decoder()( + input_ids=dec_ids, encoder_hidden_states=enc + )[0] + logits = model.lm_head(hidden * (model.config.d_model ** -0.5)) + nxt = int(logits[0, -1].argmax()) + if nxt == EOS_ID: + break + outs.append(nxt) + dec_ids = torch.cat([dec_ids, torch.tensor([[nxt]], dtype=torch.long)], 1) + return outs + + +def _decode_tokens(tokens: list[int]) -> str: + return bytes((t - BYTE_OFFSET) % 256 for t in tokens).decode( + "utf-8", errors="replace" + ) + + +def _sessions_from_zip(zip_path: Path): + import tempfile + + import onnxruntime as ort + + with tempfile.TemporaryDirectory() as tmp: + with zipfile.ZipFile(zip_path) as zf: + zf.extract("encoder.onnx", tmp) + decoder = "decoder-kv.onnx" if "decoder-kv.onnx" in zf.namelist() else "decoder.onnx" + zf.extract(decoder, tmp) + enc = ort.InferenceSession( + str(Path(tmp) / "encoder.onnx"), providers=["CPUExecutionProvider"] + ) + dec = ort.InferenceSession( + str(Path(tmp) / decoder), providers=["CPUExecutionProvider"] + ) + return enc, dec + + +def reference_decode(model, sources, max_len: int = 256) -> list[list[int]]: + """Torch-reference greedy decode of many inputs, computed once and + shared across precision variants by run_parity.""" + return [_torch_greedy_tokens(model, source, max_len) for source in sources] + + +def run_parity( + model, zip_path: Path | str, pairs, max_len: int = 256, reference=None +) -> ParityReport: + """pairs: iterable of (source_text, gold_target). Measures both sides + against gold; the gate is the CER distance between the two. + + ``reference`` (from ``reference_decode``) skips the torch side — the + reference is precision-independent, so multi-zip gates decode it once. + """ + zip_path = Path(zip_path) + enc, kv = _sessions_from_zip(zip_path) + + n = 0 + mismatches = 0 + cer_ref_sum = 0.0 + cer_onnx_sum = 0.0 + for i, (source, gold) in enumerate(pairs): + n += 1 + ref = reference[i] if reference is not None else _torch_greedy_tokens( + model, source, max_len + ) + got = onnx_greedy_kv(enc, kv, source, max_len) + if ref != got: + mismatches += 1 + cer_ref_sum += char_error_rate(_decode_tokens(ref), gold) + cer_onnx_sum += char_error_rate(_decode_tokens(got), gold) + + cer_ref = 100.0 * cer_ref_sum / max(n, 1) + cer_onnx = 100.0 * cer_onnx_sum / max(n, 1) + return ParityReport( + samples=n, + cer_reference=round(cer_ref, 4), + cer_onnx=round(cer_onnx, 4), + cer_delta=round(abs(cer_onnx - cer_ref), 4), + token_mismatches=mismatches, + ) + + +def write_parity(zip_path: Path | str, report: ParityReport) -> Path: + """Write the parity block into the zip's metadata and enforce strict + validation. Raises if the gate does not pass.""" + from imf.validator import validate_zip + + zip_path = Path(zip_path) + result = validate_zip(zip_path) + if not result.ok or result.metadata is None: + raise RuntimeError(f"cannot write parity into invalid zip: {result.errors}") + if not report.passed: + raise RuntimeError( + f"parity gate FAILED: cer_delta {report.cer_delta}pp over " + f"{report.samples} samples (limits: <= {Parity.MAX_CER_DELTA}pp, " + f">= {Parity.MIN_SAMPLES} samples)" + ) + + import tempfile + + metadata = result.metadata + updated = ModelMetadata( + format=metadata.format, + id=metadata.id, + task=metadata.task, + source_script=metadata.source_script, + target=metadata.target, + tokenizer=metadata.tokenizer, + opset=metadata.opset, + decoder=metadata.decoder, + precision=metadata.precision, + license=metadata.license, + trained_from=metadata.trained_from, + metrics=metadata.metrics, + parity=Parity(samples=report.samples, cer_delta=report.cer_delta), + sha256=metadata.sha256, + ) + + import yaml + + from imf.pack import _to_dict + + # Same filesystem as the target: os.replace is atomic within one + # filesystem and fails with EXDEV across a volume mount. + with tempfile.TemporaryDirectory(dir=zip_path.parent) as tmp: + rewritten = Path(tmp) / "rewritten.zip" + with zipfile.ZipFile(zip_path) as src, zipfile.ZipFile( + rewritten, "w", zipfile.ZIP_DEFLATED + ) as dst: + for name in src.namelist(): + if name == "metadata.yaml": + dst.writestr( + name, + yaml.safe_dump(_to_dict(updated), sort_keys=False, allow_unicode=True), + ) + else: + dst.writestr(name, src.read(name)) + rewritten.replace(zip_path) + + strict = validate_zip(zip_path, strict=True) + if not strict.ok: + raise RuntimeError(f"strict gate failed after parity write: {strict.errors}") + return zip_path + + +def write_golden(zip_path: Path | str, inputs, out_path: Path | str, max_len: int = 256) -> Path: + """Emit the cross-runtime golden set: fixed inputs + reference outputs + from the ONNX graphs (Python is the reference implementation).""" + + + zip_path = Path(zip_path) + out_path = Path(out_path) + enc, kv = _sessions_from_zip(zip_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w", encoding="utf-8") as fh: + for source in inputs: + tokens = onnx_greedy_kv(enc, kv, source, max_len) + fh.write( + json.dumps( + {"input": source, "tokens": tokens, "output": _decode_tokens(tokens)}, + ensure_ascii=False, + ) + + "\n" + ) + return out_path diff --git a/tests/test_imf_export.py b/tests/test_imf_export.py new file mode 100644 index 0000000..24d5ca8 --- /dev/null +++ b/tests/test_imf_export.py @@ -0,0 +1,136 @@ +"""Tests for ``imf.export`` — the WO02 ONNX export pipeline. + +Real end-to-end over a tiny random T5 (fixture mode): export all three +graphs at opset 14, pack fp32/fp16/int8 zips, and require every +precision to greedy-decode IDENTICALLY to the torch reference model — +plain and KV paths both. No mocks anywhere; skipped when torch / +transformers / onnxruntime are absent. +""" + +from __future__ import annotations + +import zipfile +from pathlib import Path + +import pytest +import yaml + +torch = pytest.importorskip("torch") +transformers = pytest.importorskip("transformers") +ort = pytest.importorskip("onnxruntime") + +from imf.export import ( # noqa: E402 + export_zips, + load_byte_seq2seq, + make_fixture_checkpoint, + onnx_greedy_kv, + onnx_greedy_plain, +) +from imf.parity import _torch_greedy_tokens as torch_greedy # noqa: E402 +from imf.validator import validate_zip # noqa: E402 + +TEXTS = ["he", "hello", "abc"] +MAX_LEN = 12 +PROVIDERS = ["CPUExecutionProvider"] + +METADATA = { + "format": "imf-v1", + "id": "fixture-1.0", + "task": "translit", + "source_script": "Latn", + "target": "Latn", + "tokenizer": "bytes", + "opset": 14, + "decoder": "kv", + "precision": "fp32", + "license": "BSD-3-Clause", + "trained_from": "imf export fixture (seed 42)", + "metrics": [ + { + "name": "cer", + "value": 0.0, + "protocol": "fixture self-check", + "source": "ml-models/tests/test_imf_export.py#fixture", + } + ], +} + + +@pytest.fixture(scope="module") +def reference_model(tmp_path_factory: pytest.TempPathFactory): + ckpt = make_fixture_checkpoint(tmp_path_factory.mktemp("fixture") / "checkpoint") + return load_byte_seq2seq(ckpt) + + +@pytest.fixture(scope="module") +def zips( + tmp_path_factory: pytest.TempPathFactory, reference_model +) -> dict[str, Path]: + root = tmp_path_factory.mktemp("export") + metadata_path = root / "metadata.yaml" + metadata_path.write_text(yaml.safe_dump(METADATA), encoding="utf-8") + paths = export_zips( + reference_model, metadata_path, "# fixture\n", root / "out" + ) + return {p.name: p for p in paths} + + +def _sessions(zip_path: Path): + out = zip_path.parent / zip_path.stem + out.mkdir(exist_ok=True) + with zipfile.ZipFile(zip_path) as zf: + zf.extractall(out) + enc = ort.InferenceSession(str(out / "encoder.onnx"), providers=PROVIDERS) + dec = ort.InferenceSession(str(out / "decoder.onnx"), providers=PROVIDERS) + kv = ort.InferenceSession(str(out / "decoder-kv.onnx"), providers=PROVIDERS) + return enc, dec, kv + + +def test_fixture_exports_match_torch(reference_model) -> None: + """KV and plain ONNX decode must equal the (untraced) torch model.""" + import tempfile + + from imf.export import export_graphs + + with tempfile.TemporaryDirectory() as tmp: + graphs = export_graphs(reference_model, Path(tmp) / "graphs") + enc = ort.InferenceSession(str(graphs["encoder.onnx"]), providers=PROVIDERS) + dec = ort.InferenceSession(str(graphs["decoder.onnx"]), providers=PROVIDERS) + kv = ort.InferenceSession(str(graphs["decoder-kv.onnx"]), providers=PROVIDERS) + for text in TEXTS: + expected = torch_greedy(reference_model, text, MAX_LEN) + assert onnx_greedy_plain(enc, dec, text, MAX_LEN) == expected + assert onnx_greedy_kv(enc, kv, text, MAX_LEN) == expected + + +def test_all_precisions_ship_and_validate(zips: dict[str, Path]) -> None: + assert set(zips) == { + "fixture-1.0-fp32.zip", + "fixture-1.0-fp16.zip", + "fixture-1.0-int8.zip", + } + for name, path in zips.items(): + result = validate_zip(path) + assert result.ok, result.errors + assert result.metadata is not None + assert result.metadata.decoder == "kv" + assert result.metadata.opset == 14 + assert result.metadata.precision in name + + +def test_precisions_are_token_exact(zips: dict[str, Path]) -> None: + """fp16/int8 must not change a single generated token vs fp32.""" + enc32, _, kv32 = _sessions(zips["fixture-1.0-fp32.zip"]) + reference = { + text: onnx_greedy_kv(enc32, kv32, text, MAX_LEN) for text in TEXTS + } + for precision in ("fp16", "int8"): + enc, _, kv = _sessions(zips[f"fixture-1.0-{precision}.zip"]) + for text in TEXTS: + assert onnx_greedy_kv(enc, kv, text, MAX_LEN) == reference[text] + + +def test_fp16_smaller_than_fp32(zips: dict[str, Path]) -> None: + assert zips["fixture-1.0-fp16.zip"].stat().st_size < zips[ + "fixture-1.0-fp32.zip" + ].stat().st_size diff --git a/tests/test_imf_parity.py b/tests/test_imf_parity.py new file mode 100644 index 0000000..c02292c --- /dev/null +++ b/tests/test_imf_parity.py @@ -0,0 +1,105 @@ +"""Tests for ``imf.parity`` — the WO03 gate. + +Runs the real gate end-to-end on the fixture model (torch reference vs +ONNX KV decode over 600 pairs), then exercises the gate's failure modes +with real ParityReport values. No mocks. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml + +torch = pytest.importorskip("torch") +transformers = pytest.importorskip("transformers") +ort = pytest.importorskip("onnxruntime") + +from imf.export import ( # noqa: E402 + export_zips, + load_byte_seq2seq, + make_fixture_checkpoint, +) +from imf.parity import ParityReport, run_parity, write_golden, write_parity # noqa: E402 +from imf.validator import validate_zip # noqa: E402 + +METADATA = { + "format": "imf-v1", + "id": "fixture-1.0", + "task": "translit", + "source_script": "Latn", + "target": "Latn", + "tokenizer": "bytes", + "opset": 14, + "decoder": "kv", + "precision": "fp32", + "license": "BSD-3-Clause", + "trained_from": "imf export fixture (seed 42)", + "metrics": [ + { + "name": "cer", + "value": 0.0, + "protocol": "fixture self-check", + "source": "ml-models/tests/test_imf_parity.py#fixture", + } + ], +} + +PAIRS = [(text, "xxxxx") for text in ("he", "hello", "abc", "world")] + + +@pytest.fixture(scope="module") +def gated_zip(tmp_path_factory: pytest.TempPathFactory) -> Path: + ckpt = make_fixture_checkpoint(tmp_path_factory.mktemp("ckpt") / "fixture") + model = load_byte_seq2seq(ckpt) + root = tmp_path_factory.mktemp("parity") + metadata_path = root / "metadata.yaml" + metadata_path.write_text(yaml.safe_dump(METADATA), encoding="utf-8") + zips = export_zips(model, metadata_path, "# fixture\n", root / "out") + report = run_parity(model, zips[0], PAIRS * 150, max_len=12) + assert report.passed + write_parity(zips[0], report) + return zips[0] + + +def test_parity_report_is_exact_for_fp32(gated_zip: Path) -> None: + result = validate_zip(gated_zip, strict=True) + assert result.ok, result.errors + assert result.metadata is not None + assert result.metadata.parity is not None + assert result.metadata.parity.samples == 600 + assert result.metadata.parity.cer_delta == 0.0 + + +def test_gate_rejects_high_cer_delta(gated_zip: Path) -> None: + bad = ParityReport( + samples=600, cer_reference=10.0, cer_onnx=10.5, + cer_delta=0.5, token_mismatches=3, + ) + with pytest.raises(RuntimeError, match="parity gate FAILED"): + write_parity(gated_zip, bad) + + +def test_gate_rejects_small_sample(gated_zip: Path) -> None: + small = ParityReport( + samples=100, cer_reference=10.0, cer_onnx=10.0, + cer_delta=0.0, token_mismatches=0, + ) + with pytest.raises(RuntimeError, match="parity gate FAILED"): + write_parity(gated_zip, small) + + +def test_golden_jsonl_roundtrip(gated_zip: Path, tmp_path: Path) -> None: + inputs = tmp_path / "inputs.jsonl" + inputs.write_text( + "\n".join(json.dumps(t) for t in ("he", "hello", "abc")) + "\n", + encoding="utf-8", + ) + out = write_golden(gated_zip, ["he", "hello", "abc"], tmp_path / "golden.jsonl", max_len=12) + rows = [json.loads(line) for line in out.read_text(encoding="utf-8").splitlines()] + assert [r["input"] for r in rows] == ["he", "hello", "abc"] + for row in rows: + assert set(row) == {"input", "tokens", "output"} + assert all(isinstance(t, int) for t in row["tokens"])