kwja-mlx

kwja-mlx

This repository contains MLX-converted weights for upstream KWJA checkpoints.

Source

The files here are converted from the official KWJA checkpoint releases under kwja==2.5.1 / model version v2.4. The original checkpoint files are downloaded from:

• https://lotus.kuee.kyoto-u.ac.jp/kwja/v2.4/

Converted modules are provided for three model sizes:

• tiny
• base
• large

Each size includes three modules:

• char
• seq2seq
• word

Upstream module mapping

Conversion method

The conversion is a format conversion from PyTorch Lightning checkpoints to MLX runtime bundles.

High-level steps:

1. Load the original .ckpt file on CPU.
2. Read model metadata from the checkpoint hyper_parameters.
3. Extract the relevant module weights from the checkpoint state_dict.
4. Remap or sanitize tensor names where needed so they match the MLX module layout.
5. Save weights as .mlx.safetensors.
6. Save loader metadata as a sidecar .mlx.json file.

Module-specific notes:

• char: exports the DeBERTa encoder and tagging heads; the convolution weight is transposed to match the MLX layout.
• seq2seq: exports the encoder-decoder weights after stripping the encoder_decoder. prefix and sanitizing them for the MLX T5 implementation.
• word: exports the DeBERTa encoder, prediction heads, and CRF transition parameters used at inference time.

The conversion command used in the companion tooling is:

python -m hira kwja-export-mlx --model-size base --module all

The same conversion flow is used for tiny and large.

File layout

Each converted module is stored as a pair of files:

• <size>/<module>.mlx.safetensors: MLX weights
• <size>/<module>.mlx.json: metadata needed to reconstruct the tokenizer/config/runtime bundle

Examples:

• base/char_deberta-v2-base-wwm.mlx.safetensors
• base/char_deberta-v2-base-wwm.mlx.json
• base/seq2seq_t5-base.mlx.safetensors
• base/seq2seq_t5-base.mlx.json
• base/word_deberta-v2-base.mlx.safetensors
• base/word_deberta-v2-base.mlx.json

Scope

This repository is intended to host converted MLX artifacts. It does not add new training, fine-tuning, or evaluation results beyond the upstream KWJA checkpoints.