Han2Han pre-trained checkpoint (uniform average of the last five checkpoints, steps 6347 to 6725)

README.md

@@ -0,0 +1,128 @@
+---
+license: apache-2.0
+language:
+- ko
+library_name: transformers
+pipeline_tag: feature-extraction
+tags:
+- han2han
+- hanja
+- hangul
+- historical-korean
+- encoder-decoder
+- custom_code
+---
+
+# Han2Han PT
+
+The pre-trained (PT) checkpoint of [Han2Han](https://github.com/cadazar/han2han), a
+169M-parameter encoder-decoder model that learns script-invariant
+representations of Korean text: a document written in Hanja and its Hangul
+transcription land at the same point in embedding space. The recipe (jamo and
+character-level embedding fusion, morpheme-aware denoising, bidirectional
+Hanja-Hangul transcription) is described in the paper, accepted to Findings of
+EMNLP 2026.
+
+This is the starting point for fine-tuning, and the first of three checkpoints:
+
+| Repo | Stage |
+| --- | --- |
+| `cadazar/han2han-pt` (this one) | pre-training, 35B tokens |
+| [`cadazar/han2han-it`](https://huggingface.co/cadazar/han2han-it) | instruction tuning from these weights |
+| [`cadazar/han2han-rl`](https://huggingface.co/cadazar/han2han-rl) | reinforcement learning on Hanja-Hangul transcription from `han2han-it` |
+
+## These weights are an average
+
+The weights are the uniform average of the last five checkpoints of the
+pre-training run (steps 6347, 6441, 6536, 6630 and 6725, saved every 0.5B
+tokens from 33B to 35B), not the final checkpoint alone. The average is what
+the instruction tuning started from, chosen after comparing starting points,
+so it is the pre-trained model that `han2han-it` and `han2han-rl` descend
+from. The learning rate had nearly finished its cooldown over those steps: no
+weight differs from the final checkpoint by more than 0.0015.
+
+## Intended use
+
+Fine-tuning on downstream tasks, and sentence embeddings that treat Hanja and
+Hangul spellings of the same text alike.
+
+It is not a chat model. The tokenizer and the chat template are the same as in
+the other two repos, but pre-training never used the chat tokens
+(`<|system|>`, `<|user|>`, `<|assistant|>`, `<|think|>`, `<|end_of_turn|>`):
+their embedding rows are still nearly one shared vector here, as for every
+other token pre-training never saw. For generation from a prompt use
+`han2han-it` or `han2han-rl`.
+
+## Usage
+
+Runtime requirements: `torch` and `transformers` (tested with torch 2.14.1 and
+transformers 5.18.0 on CPU). The model class ships in this repo, so loading needs
+`trust_remote_code=True`; the tokenizer itself runs no custom code, but without
+the flag `AutoTokenizer` stops to ask.
+
+`output_sentence_embeddings=True` returns the mean-pooled encoder states, one
+640-dimensional vector per input.
+
+```python
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+repo = "cadazar/han2han-pt"
+tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+model = AutoModelForSeq2SeqLM.from_pretrained(repo, trust_remote_code=True).eval()
+
+texts = [
+    "會場을 一巡하고 돌아올 때까지도 日本人側 畵家까지도 一人도 發見할 수가 없었다.",
+    "회장을 일순하고 돌아올 때까지도 일본인측 화가까지도 일인도 발견할 수가 없었다.",
+    "南畵나 四君子에서는 破墨의 妙法으로 白雪을 象徵할 수 있다.",
+]
+inputs = tokenizer(texts, return_tensors="pt", padding=True)
+with torch.no_grad():
+    embeddings = model(**inputs, output_sentence_embeddings=True)[0]
+embeddings = F.normalize(embeddings, dim=-1)
+print(embeddings @ embeddings.T)
+# tensor([[1.0000, 0.9315, 0.8253],
+#         [0.9315, 1.0000, 0.7910],
+#         [0.8253, 0.7910, 1.0000]])
+```
+
+The first two rows are the same 1920s newspaper sentence in mixed script and in
+Hangul; the third is a different sentence.
+
+### Tokenizer
+
+`tokenizer.json` is a `tokenizers`-library build of the SentencePiece model
+(`spiece.model`, kept here as the source), made by
+`scripts/build_hf_tokenizer.py` in the GitHub repo. Calling the tokenizer does
+not add BOS or EOS tokens, and special tokens written into the text are mapped
+to their ids. It encodes like the SentencePiece wrapper the training code uses,
+with one known difference: where two segmentations of a span have exactly the
+same score (runs of digits, mostly), the same pieces can come out in a different
+order.
+
+## Files
+
+| File | Contents |
+| --- | --- |
+| `model.safetensors` | fp32 weights, 169.2M parameters, plus the `jbu` / `cbu` subword bucket tables |
+| `config.json`, `generation_config.json` | model and generation config, with `auto_map` entries for the Auto classes |
+| `tokenizer.json`, `tokenizer_config.json`, `chat_template.jinja` | fast tokenizer (38400 pieces) and chat template, the same as in `han2han-it` |
+| `spiece.model` | the SentencePiece model `tokenizer.json` was built from |
+| `modeling_han2han.py`, `han2han_config.py` | modeling code from the GitHub repo at commit [`af1330e`](https://github.com/cadazar/han2han/commit/af1330e49f1edf46e26b63937724ddbaeffc3bee), the same files as in `han2han-it` |
+
+## Citation
+
+```bibtex
+@inproceedings{han2han2026,
+  title     = {Han2Han: Efficient Language-Specific Character Representation
+               through Script-Aware Pre-Training for Historical Text Analysis},
+  author    = {Adams, Cellik and Jo, EunKyoung and Kim, Juae},
+  booktitle = {Findings of the Association for Computational Linguistics: EMNLP 2026},
+  year      = {2026}
+}
+```
+
+## License
+
+Apache License 2.0, the same as the [GitHub repo](https://github.com/cadazar/han2han).

chat_template.jinja

@@ -0,0 +1,24 @@
+{%- if messages and messages[0]['role'] == 'system' -%}
+    {%- if messages[0]['content'] is string -%}
+        {{- '<|system|>' + messages[0]['content'] -}}
+    {%- else -%}
+        {{- '<|system|>' + messages[0]['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join('') -}}
+    {%- endif -%}
+    {%- set loop_messages = messages[1:] -%}
+{%- else -%}
+    {%- set loop_messages = messages -%}
+{%- endif -%}
+{%- for message in loop_messages -%}
+    {%- if message['role'] not in ['user', 'assistant'] -%}
+        {{- raise_exception('Han2Han chat supports a leading system message, then user and assistant turns; got role ' + message['role']) -}}
+    {%- endif -%}
+    {%- if message['content'] is string -%}
+        {%- set content = message['content'] -%}
+    {%- else -%}
+        {%- set content = message['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join('') -%}
+    {%- endif -%}
+    {{- '<|' + message['role'] + '|>' + content + '<|end_of_turn|>' -}}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{- '<|think|>' if enable_thinking is defined and enable_thinking else '<|assistant|>' -}}
+{%- endif -%}

config.json

@@ -0,0 +1,111 @@
+{
+  "apply_legacy_rope_quirk": false,
+  "architectures": [
+    "AutoModelForCausalLM"
+  ],
+  "attention_mechanism": null,
+  "attn_pdrop": 0.0,
+  "auto_map": {
+    "AutoConfig": "han2han_config.Han2HanConfig",
+    "AutoModel": "modeling_han2han.Han2Han",
+    "AutoModelForMultipleChoice": "modeling_han2han.Han2HanForMultipleChoice",
+    "AutoModelForQuestionAnswering": "modeling_han2han.Han2HanForQuestionAnswering",
+    "AutoModelForSeq2SeqLM": "modeling_han2han.Han2Han",
+    "AutoModelForSequenceClassification": "modeling_han2han.Han2HanForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling_han2han.Han2HanForTokenClassification",
+    "AutoModelForCausalLM": "modeling_han2han.Han2HanForCausalLM"
+  },
+  "bos_token_id": 2,
+  "char_is_unified_cjk": false,
+  "char_subwords": true,
+  "char_vocab_size": 5376,
+  "classf_pdrop": 0.1,
+  "classifier_head_type": "linear",
+  "cross_attn_num_heads": 4,
+  "cross_attn_num_kv_heads": 2,
+  "cross_attn_pdrop": 0.0,
+  "d_ff": 2048,
+  "d_model": 640,
+  "d_prime": 1024,
+  "decoder_attention_types": [
+    "mha-sliding",
+    "mha-sliding",
+    "mha-sliding",
+    "mha-sliding",
+    "mha-sliding",
+    "mha"
+  ],
+  "decoder_cross_attention_types": [
+    "mha"
+  ],
+  "decoder_nlayer": 18,
+  "decoder_norm_type": "rmsnorm",
+  "decoder_start_token_id": 3,
+  "dense_ffn_activation": "geglu",
+  "dtype": "float32",
+  "embd_pdrop": 0.0,
+  "embedding_dropout_rate": 0.0,
+  "encoder_attention_types": [
+    "mha-sliding",
+    "mha-sliding",
+    "mha-sliding",
+    "mha-sliding",
+    "mha-sliding",
+    "mha"
+  ],
+  "encoder_nlayer": 18,
+  "encoder_norm_type": "rmsnorm",
+  "eos_token_id": 3,
+  "ffn_activation": "geglu",
+  "head_dim": 256,
+  "init_biases_normal": true,
+  "initializer_range": 0.02,
+  "is_encoder_decoder": true,
+  "jamo_subwords": true,
+  "jamo_vocab_size": 4992,
+  "kernel_init_scale": 0.1,
+  "kernel_init_type": "variance_scaling",
+  "label_smoothing": 0.0,
+  "layer_norm_epsilon": 1e-05,
+  "layer_pdrop": 0.0,
+  "model_type": "han2han",
+  "n_positions": 2048,
+  "num_heads": 4,
+  "num_kv_heads": 1,
+  "pad_token_id": 0,
+  "qk_norm_post_rope": true,
+  "query_pre_attn_scalar": 256,
+  "remat_policy": "full",
+  "resid_pdrop": 0.0,
+  "return_dict": true,
+  "rope_theta": 500000,
+  "rope_theta_sliding": 10000,
+  "seed": 42,
+  "sft_decoder_start_token_id_default": null,
+  "sft_decoder_start_token_id_thinking": null,
+  "sft_eos_token_id": null,
+  "sliding_window_size": 128,
+  "subword_embed_dim": 384,
+  "subword_entry_fusion": true,
+  "entry_pool_mode": "legacy",
+  "entry_pad_bias_slots": 0,
+  "entry_slot_width": 32,
+  "entry_pool_unroll": true,
+  "tie_encoder_decoder": true,
+  "tie_input_output_embeddings": true,
+  "tie_subtoken_embeddings": false,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.18.0",
+  "use_bart_collator": true,
+  "use_bart_training": true,
+  "use_bias": true,
+  "use_fla_fused_mlp": false,
+  "use_fla_fused_norm": false,
+  "use_fla_fused_rotary": false,
+  "use_han2han_transcription": "false",
+  "use_learned_bidirectional": true,
+  "use_qk_norm": true,
+  "use_scan_layers": true,
+  "use_sub_ln": true,
+  "vocab_size": 38400
+}

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "bos_token_id": 2,
+  "decoder_start_token_id": 9,
+  "eos_token_id": 10,
+  "pad_token_id": 0,
+  "use_cache": true,
+  "transformers_version": "5.18.0"
+}

han2han_config.py

@@ -0,0 +1,565 @@
+#!/usr/bin/env python3
+# coding: utf-8
+
+import logging
+
+from transformers.configuration_utils import PretrainedConfig
+from typing import Optional, List
+
+logger = logging.getLogger(__name__)
+
+class Han2HanConfig(PretrainedConfig):
+    model_type = "han2han"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    has_no_defaults_at_init = True
+    attribute_map = {
+        "hidden_size": "d_model",
+        "num_attention_heads": "num_heads",
+        "num_hidden_layers": "decoder_nlayer",  # HF expects this for decoder
+        "num_layers": "decoder_nlayer",
+        "encoder_layers": "encoder_nlayer",
+        "decoder_layers": "decoder_nlayer",
+        "d_kv": "d_prime",
+        "intermediate_size": "d_ff",
+        "hidden_dropout_prob": "resid_pdrop",
+        "attention_probs_dropout_prob": "attn_pdrop",
+    }
+
+    # Fields that should always be saved, even if they match PretrainedConfig defaults
+    _always_save = ["tie_word_embeddings", "tie_encoder_decoder"]
+
+    # Explicitly list attributes that should always be saved
+    def to_diff_dict(self):
+        """
+        Override to ensure certain fields are always saved, even if they match
+        PretrainedConfig defaults.
+        """
+        # Get the diff dict from parent
+        output = super().to_diff_dict()
+        
+        # Force inclusion of fields we always want to save
+        for field in self._always_save:
+            if hasattr(self, field):
+                output[field] = getattr(self, field)
+
+        return output
+
+    def get_text_config(self, decoder=None, encoder=None):
+        """Return this config for either side of the model.
+
+        The encoder and decoder share one flat config. The base implementation
+        treats flat encoder-decoder configs as legacy and strips the `decoder_`
+        prefix from every key (`decoder_nlayer` -> `nlayer`), which leaves
+        `num_hidden_layers` unresolvable when `generate()` sizes its KV cache.
+        """
+        return self
+
+    def __init__(
+        self,
+
+        jamo_subwords: bool = False,
+        char_subwords: bool = False,
+        use_han2han_transcription: str = 'false',  # 'false', 'true', 'hangul_only', 'reverse'
+
+        vocab_size: int = 38000,
+        jamo_vocab_size: int = 8000,
+        char_vocab_size: int = 8000,
+        subword_embed_dim: Optional[int] = None,  # hidden dim of wje/wce subtoken embeddings; None defaults to d_model // 2
+        char_is_unified_cjk: bool = False,
+        decoder_nlayer: int = 4,
+        encoder_nlayer: int = 4,
+        n_positions: int = 1024,
+        d_model: int = 768,
+        d_prime: Optional[int] = None,
+        rope_theta: float = 10000.0,
+        rope_theta_sliding: Optional[float] = None,  # per-layer theta override for 'sliding'/'local' attention layers (None = inherit rope_theta)
+        apply_legacy_rope_quirk: Optional[bool] = None,  # None = auto-detect legacy scan_rope_theta quirk; False = opt out (use rope_theta/rope_theta_sliding as specified, for V2+); True = force-apply
+        d_ff: int = 3072,
+        attention_mechanism: str = 'mha',          # global default attention type for all layers
+        encoder_attention_types: Optional[List[str]] = None,  # per-layer encoder self-attention types
+        decoder_attention_types: Optional[List[str]] = None,  # per-layer decoder self-attention types
+        decoder_cross_attention_types: Optional[List[str]] = None,  # per-layer decoder cross-attention types
+        sliding_window_size: int = 256,             # window size for 'mha-sliding' layers (0 = full attention)
+        ffn_activation: str = "swiglu",     # 'swiglu', 'geglu', 'reglu2', 'gelu', 'gelu_new', 'relu2'
+        dense_ffn_activation: Optional[str] = None,  # activation override; None = follow ffn_activation
+        use_fla_fused_mlp: bool = False,
+        use_fla_fused_norm: bool = False,
+        use_fla_fused_rotary: bool = False,
+        num_heads: Optional[int] = None,
+        head_dim: Optional[int] = None,                       # MHA only: per-head dim. If both set with d_prime, must be consistent. Required for GQA.
+        num_kv_heads: Optional[int] = None,                   # MHA only: KV heads for self-attn (None = num_heads = full MHA). 1 = MQA. Must divide num_heads.
+        cross_attn_num_heads: Optional[int] = None,           # MHA only: Q heads for cross-attn (None = num_heads). Lets cross-attn use full d_model fidelity (cross_attn_num_heads * head_dim) while self-attn stays compressed.
+        cross_attn_num_kv_heads: Optional[int] = None,        # MHA only: KV heads for cross-attn (None = num_kv_heads). Must divide cross_attn_num_heads.
+        use_qk_norm: bool = False,                            # MHA only: per-head RMSNorm on Q and K (Gemma 3 / T5Gemma 2 style)
+        query_pre_attn_scalar: Optional[float] = None,        # MHA only: HF-Gemma 3 semantics. Q multiplier = scalar ** -0.5. None = head_dim ** -0.5.
+        num_labels: int = 3,
+        use_learned_bidirectional: float = True,
+        remat_policy: str = "full",
+        layer_pdrop: float = 0.1,
+        resid_pdrop: float = 0.1,
+        embd_pdrop: float = 0.1,
+        attn_pdrop: float = 0.1,
+        cross_attn_pdrop: float = 0.15,
+        classf_pdrop: float = 0.1,
+        classifier_head_type: str = 'linear',  # 'linear' (T5Gemma 2) or 'mlp' (RoBERTa-style tanh)
+        embedding_dropout_rate: float = 0.0,  # Probability of dropping each type of embedding module (wte/wje/wce)
+        layer_norm_epsilon: float = 1e-5,
+        decoder_norm_type: str = 'rmsnorm',  # 'rmsnorm', 'rmsnorm_bias', 'layernorm'
+        encoder_norm_type: str = 'rmsnorm',  # 'rmsnorm', 'rmsnorm_bias', 'layernorm'
+        initializer_range: float = 0.02,
+        kernel_init_type: str = 'normal',      # 'normal' or 'variance_scaling'
+        kernel_init_scale: float = 0.1,        # scale for variance_scaling (0.1=Switch-style, 1.0=lecun)
+        init_biases_normal: bool = False,      # if True, init biases as normal(stddev=initializer_range) (V1 behavior); else zeros
+        init_cache: bool = False,
+        pad_token_id: int = 1,
+        decoder_start_token_id: int = 0,
+        tie_word_embeddings: bool = True,
+        tie_encoder_decoder: bool = False,
+        tie_input_output_embeddings: bool = False,
+        tie_subtoken_embeddings: bool = True,
+        return_dict: bool = True,
+        seed: int = 0,
+        eos_token_id: int = 2,
+        bos_token_id: int = 0,
+        use_bart_training: bool = True,
+        use_bart_collator: bool = True,
+
+        # SFT-only token ids. Read by ChatSFTCollator; not used by pretraining.
+        # Decoder is primed with <|think|> when reasoning is enabled, otherwise
+        # with <|assistant|>; turn closes with <|end_of_turn|>. None means look
+        # up from tokenizer at collator init time (recommended).
+        sft_decoder_start_token_id_thinking: Optional[int] = None,
+        sft_decoder_start_token_id_default: Optional[int] = None,
+        sft_eos_token_id: Optional[int] = None,
+
+        use_sub_ln: bool = False,                   # SubLN: RMSNorm before output projections in attn and FFN
+
+        # Bias configuration
+        use_bias: bool = True,                      # global toggle for biases in all linear layers
+
+        label_smoothing: float = 0.1,               # label smoothing alpha for cross-entropy loss
+
+        # scan layers (compile one layer body, repeat N times via jax.lax.scan)
+        use_scan_layers: bool = False,
+
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.jamo_vocab_size = jamo_vocab_size
+        self.char_vocab_size = char_vocab_size
+        # subtoken embedding hidden dim: defaults to d_model // 2 when unspecified
+        self.subword_embed_dim = subword_embed_dim if subword_embed_dim is not None else d_model // 2
+        self.char_is_unified_cjk = char_is_unified_cjk
+        self.jamo_subwords = jamo_subwords
+        self.char_subwords = char_subwords
+        self.use_han2han_transcription = use_han2han_transcription
+        self.decoder_nlayer = decoder_nlayer
+        self.encoder_nlayer = encoder_nlayer
+        self.n_positions = n_positions
+        self.d_model = d_model
+        self.rope_theta = rope_theta
+        self.rope_theta_sliding = rope_theta_sliding
+        self.apply_legacy_rope_quirk = apply_legacy_rope_quirk
+        self.d_prime = d_prime
+        self.d_ff = d_ff
+        self.ffn_activation = ffn_activation
+        self.dense_ffn_activation = dense_ffn_activation if dense_ffn_activation is not None else ffn_activation
+        # Only set num_heads for MHA attention
+        self.num_heads = num_heads if attention_mechanism == 'mha' or any(
+            'mha' in attn_type for attn_type in (encoder_attention_types or [])) else None
+        self.head_dim = head_dim
+        self.num_kv_heads = num_kv_heads
+        self.cross_attn_num_heads = cross_attn_num_heads
+        self.cross_attn_num_kv_heads = cross_attn_num_kv_heads
+        self.use_qk_norm = use_qk_norm
+        self.query_pre_attn_scalar = query_pre_attn_scalar
+        self.use_fla_fused_mlp = use_fla_fused_mlp
+        self.use_fla_fused_norm = use_fla_fused_norm
+        self.use_fla_fused_rotary = use_fla_fused_rotary
+        self.remat_policy = remat_policy
+        self.layer_pdrop = layer_pdrop
+        self.resid_pdrop = resid_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.attn_pdrop = attn_pdrop
+        self.cross_attn_pdrop = cross_attn_pdrop
+        self.classf_pdrop = classf_pdrop
+        self.classifier_head_type = classifier_head_type
+        self.embedding_dropout_rate = embedding_dropout_rate
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.decoder_norm_type = decoder_norm_type
+        self.encoder_norm_type = encoder_norm_type
+        self.initializer_range = initializer_range
+        self.kernel_init_type = kernel_init_type
+        self.kernel_init_scale = kernel_init_scale
+        self.init_biases_normal = init_biases_normal
+        self.use_cache = init_cache
+        self.tie_input_output_embeddings = tie_input_output_embeddings
+        self.use_sub_ln = use_sub_ln
+
+        # bias configuration
+        self.use_bias = use_bias
+
+        self.label_smoothing = label_smoothing
+        self.use_scan_layers = use_scan_layers
+
+        super().__init__(
+            pad_token_id = pad_token_id,
+            decoder_start_token_id = decoder_start_token_id,
+            tie_encoder_decoder = tie_encoder_decoder,
+            tie_word_embeddings = tie_word_embeddings,
+            bos_token_id = bos_token_id,
+            eos_token_id = eos_token_id,
+            **kwargs
+        )
+
+        self.num_labels = num_labels
+        self.is_encoder_decoder = True
+        self.tie_word_embeddings = tie_word_embeddings
+        self.tie_encoder_decoder = tie_encoder_decoder
+        self.tie_subtoken_embeddings = tie_subtoken_embeddings
+        self.seed = seed
+        self.return_dict = return_dict
+        self.eos_token_id = eos_token_id
+        self.sft_decoder_start_token_id_thinking = sft_decoder_start_token_id_thinking
+        self.sft_decoder_start_token_id_default = sft_decoder_start_token_id_default
+        self.sft_eos_token_id = sft_eos_token_id
+        self.attention_mechanism = attention_mechanism
+        self.encoder_attention_types = encoder_attention_types
+        self.decoder_attention_types = decoder_attention_types
+        self.decoder_cross_attention_types = decoder_cross_attention_types
+        self.sliding_window_size = sliding_window_size
+        self.use_learned_bidirectional = use_learned_bidirectional
+        self.use_bart_training = use_bart_training
+        self.use_bart_collator = use_bart_collator
+
+        self._sanitize_d_prime(encoder_attention_types, decoder_attention_types, decoder_cross_attention_types)
+        self._resolve_mha_attention_fields(encoder_attention_types, decoder_attention_types, decoder_cross_attention_types)
+
+        if self.d_prime > 256 and self.use_fla_fused_rotary:
+            logger.warning(
+                "WARNING: `d_prime` is greater than 256 and `use_fla_fused_rotary` is set to `True`. "
+                "This is not supported. Falling back to non-fused rotary embeddings."
+            )
+            self.use_fla_fused_rotary = False
+
+        self._validate_attention_config()
+
+        if self.use_scan_layers:
+            self._validate_scan_config()
+
+        self._apply_legacy_rope_quirk()
+
+    def _apply_legacy_rope_quirk(self) -> None:
+        """Rewrite rope_theta when the legacy Flax `_scan_key` quirk applied.
+
+        Pre-fix Flax `_identify_scan_groups` collapsed all MHA variants
+        ('mha', 'mha-sliding', 'mha-local') under a single scan key. So when
+        `use_scan_layers=True` and the attention pattern mixes 'mha' with
+        'mha-sliding' or 'mha-local', all layers in the resulting scan stack
+        share one `FlaxHan2HanBlock` instance whose `effective_rope_theta` is
+        baked from `position_specs[0]`. In practice the configured
+        `rope_theta` field was overwritten by `rope_theta_sliding` whenever
+        the first attention layer was sliding.
+
+        Checkpoints trained under that quirk have a saved `rope_theta` that
+        doesn't reflect what the weights actually saw. This sanitizer detects
+        the pattern and rewrites the in-memory config so the runtime loads
+        the model with values matching its training behavior. Flax post-fix
+        resolves `effective_rope_theta` per call from `override_window_size`,
+        so the sanitizer's rewrite is what makes legacy checkpoints behave
+        identically under the corrected dispatch path.
+
+        Honors `self.apply_legacy_rope_quirk`:
+          - None (default): auto-apply when the heuristic triggers (existing
+            checkpoints, no opt-out specified -- safe default).
+          - False: explicit opt-out (V2+ runs that want hybrid rope_theta).
+            Sanitizer becomes a no-op even when the trigger pattern matches.
+          - True: force-apply (mostly for testing).
+        """
+        opt = getattr(self, 'apply_legacy_rope_quirk', None)
+        if opt is False:
+            return
+
+        if not getattr(self, 'use_scan_layers', False):
+            return
+        if getattr(self, 'rope_theta_sliding', None) is None:
+            return
+        if self.rope_theta == self.rope_theta_sliding:
+            return
+
+        def _has_mixed_mha(types):
+            if not types:
+                return False
+            has_sliding = any(
+                ('sliding' in t) or ('local' in t)
+                for t in types if t
+            )
+            has_full_mha = any(t == 'mha' for t in types)
+            return has_sliding and has_full_mha
+
+        triggered = opt is True or (
+            _has_mixed_mha(self.encoder_attention_types)
+            or _has_mixed_mha(self.decoder_attention_types)
+        )
+        if not triggered:
+            return
+
+        logger.warning(
+            "[Han2HanConfig] Detected legacy scan_rope_theta quirk "
+            "(use_scan_layers=True + mixed mha/mha-sliding pattern + "
+            "rope_theta != rope_theta_sliding). Pre-fix Flax "
+            "_identify_scan_groups collapsed MHA variants into one scan "
+            "stack whose rope_theta was baked from position_specs[0], so "
+            "the saved rope_theta=%s was overwritten by rope_theta_sliding=%s "
+            "during training. Rewriting in-memory rope_theta=%s to match "
+            "actual training behavior. To opt out for new training, set "
+            "apply_legacy_rope_quirk=False (CLI: --no_apply_legacy_rope_quirk).",
+            self.rope_theta, self.rope_theta_sliding, self.rope_theta_sliding,
+        )
+        self.rope_theta = self.rope_theta_sliding
+
+    def _sanitize_d_prime(
+        self,
+        encoder_attention_types: Optional[List[str]],
+        decoder_attention_types: Optional[List[str]],
+        decoder_cross_attention_types: Optional[List[str]],
+    ) -> None:
+        """Sanitize d_prime based on attention configuration.
+
+        - Layerwise config: d_prime defaults to d_model if not specified
+        - Global MHA: d_prime defaults to d_model if not specified
+        """
+        has_layerwise = (
+            encoder_attention_types is not None or
+            decoder_attention_types is not None or
+            decoder_cross_attention_types is not None
+        )
+
+        if has_layerwise:
+            if self.d_prime is None and self.head_dim is None:
+                self.d_prime = self.d_model
+                logger.info(f"d_prime not specified for layerwise config, defaulting to d_model={self.d_model}")
+        elif self.attention_mechanism == 'mha':
+            if self.d_prime is None and self.head_dim is None:
+                self.d_prime = self.d_model
+                logger.info(f"d_prime not specified for MHA, defaulting to d_model={self.d_model}")
+        else:
+            if self.d_prime is None:
+                raise ValueError(
+                    f"d_prime must be explicitly set for attention_mechanism='{self.attention_mechanism}'."
+                )
+
+    def _resolve_mha_attention_fields(
+        self,
+        encoder_attention_types: Optional[List[str]],
+        decoder_attention_types: Optional[List[str]],
+        decoder_cross_attention_types: Optional[List[str]],
+    ) -> None:
+        """Resolve head_dim/d_prime consistency and num_kv_heads defaults for MHA."""
+        all_types = []
+        if encoder_attention_types is not None:
+            all_types.extend(encoder_attention_types)
+        if decoder_attention_types is not None:
+            all_types.extend(decoder_attention_types)
+        if decoder_cross_attention_types is not None:
+            all_types.extend(t for t in decoder_cross_attention_types if t is not None)
+        has_layerwise = bool(all_types)
+        has_mha = (
+            any(t == 'mha' or t.startswith('mha-') for t in all_types)
+            if has_layerwise
+            else self.attention_mechanism == 'mha'
+        )
+
+        if not has_mha:
+            return
+
+        if self.num_heads is None:
+            raise ValueError(
+                "num_heads must be specified for MHA attention."
+            )
+
+        if self.head_dim is not None and self.d_prime is not None:
+            if self.head_dim * self.num_heads != self.d_prime:
+                raise ValueError(
+                    f"head_dim ({self.head_dim}) * num_heads ({self.num_heads}) = "
+                    f"{self.head_dim * self.num_heads} does not match d_prime ({self.d_prime})."
+                )
+        elif self.head_dim is not None:
+            self.d_prime = self.head_dim * self.num_heads
+        elif self.d_prime is not None:
+            if self.d_prime % self.num_heads != 0:
+                raise ValueError(
+                    f"d_prime ({self.d_prime}) must be divisible by num_heads ({self.num_heads})."
+                )
+            self.head_dim = self.d_prime // self.num_heads
+        else:
+            raise ValueError(
+                "MHA requires at least one of head_dim or d_prime to be set."
+            )
+
+        if self.num_kv_heads is None:
+            self.num_kv_heads = self.num_heads
+        if self.num_kv_heads < 1 or self.num_heads % self.num_kv_heads != 0:
+            raise ValueError(
+                f"num_kv_heads ({self.num_kv_heads}) must be >= 1 and divide "
+                f"num_heads ({self.num_heads}) evenly."
+            )
+
+        if self.cross_attn_num_heads is None:
+            self.cross_attn_num_heads = self.num_heads
+        if self.cross_attn_num_heads < 1:
+            raise ValueError(
+                f"cross_attn_num_heads ({self.cross_attn_num_heads}) must be >= 1."
+            )
+
+        if self.cross_attn_num_kv_heads is None:
+            self.cross_attn_num_kv_heads = self.num_kv_heads
+        if self.cross_attn_num_kv_heads < 1 or self.cross_attn_num_heads % self.cross_attn_num_kv_heads != 0:
+            raise ValueError(
+                f"cross_attn_num_kv_heads ({self.cross_attn_num_kv_heads}) must be >= 1 and divide "
+                f"cross_attn_num_heads ({self.cross_attn_num_heads}) evenly."
+            )
+
+        if self.query_pre_attn_scalar is not None and self.query_pre_attn_scalar <= 0:
+            raise ValueError(
+                f"query_pre_attn_scalar must be positive, got {self.query_pre_attn_scalar}."
+            )
+
+    def _validate_attention_config(self) -> None:
+        """Validate mutual exclusivity of global vs layerwise attention config."""
+        has_layerwise = (
+            self.encoder_attention_types is not None or
+            self.decoder_attention_types is not None or
+            self.decoder_cross_attention_types is not None
+        )
+
+        if has_layerwise:
+            missing = []
+            if self.encoder_attention_types is None:
+                missing.append('encoder_attention_types')
+            if self.decoder_attention_types is None:
+                missing.append('decoder_attention_types')
+            if self.decoder_cross_attention_types is None:
+                missing.append('decoder_cross_attention_types')
+
+            if missing:
+                raise ValueError(
+                    f"Per-layer attention mode requires all three lists to be specified. "
+                    f"Missing: {missing}. Either specify all layerwise types, or remove "
+                    f"all and use 'attention_mechanism' for a global default."
+                )
+
+            if self.attention_mechanism is not None:
+                logger.warning(
+                    f"Both attention_mechanism ('{self.attention_mechanism}') and layerwise "
+                    f"attention types are set. Layerwise types take precedence; "
+                    f"attention_mechanism will be ignored."
+                )
+        else:
+            if self.attention_mechanism is None:
+                raise ValueError(
+                    "No attention configuration specified. Either set 'attention_mechanism' "
+                    "for a global default, or specify all three layerwise types: "
+                    "encoder_attention_types, decoder_attention_types, decoder_cross_attention_types."
+                )
+
+    def _validate_scan_config(self) -> None:
+        """Validate config constraints for scanned layer execution."""
+        # layerdrop breaks scan (stochastic layer skip)
+        if self.layer_pdrop > 0:
+            raise ValueError(
+                f"use_scan_layers is incompatible with layerdrop > 0. "
+                f"Got layer_pdrop={self.layer_pdrop}. Set layerdrop: 0.0."
+            )
+
+    def make_kernel_init(self, dtype=None):
+        """Create kernel initializer based on config."""
+        # try block keeps transformers' remote-code import check from requiring flax
+        try:
+            from flax import nnx
+        except ImportError:
+            raise
+        kwargs = {'dtype': dtype} if dtype is not None else {}
+        if self.kernel_init_type == 'variance_scaling':
+            return nnx.initializers.variance_scaling(
+                self.kernel_init_scale, 'fan_in', 'truncated_normal', **kwargs
+            )
+        return nnx.initializers.normal(stddev=self.initializer_range, **kwargs)
+
+    def make_bias_init(self):
+        """Create bias initializer based on config.
+
+        V1 behavior is normal(stddev=initializer_range); current default is zeros.
+        """
+        # try block keeps transformers' remote-code import check from requiring flax
+        try:
+            from flax import nnx
+        except ImportError:
+            raise
+        if self.init_biases_normal:
+            return nnx.initializers.normal(stddev=self.initializer_range)
+        return nnx.initializers.zeros_init()
+
+    def get_encoder_attention_types(self) -> List[str]:
+        """Get expanded per-layer encoder attention types.
+
+        If encoder_attention_types is None, returns [attention_mechanism] * encoder_nlayer.
+        If specified, expands short patterns by repetition (pattern length must divide layer count).
+        """
+        return self._expand_attention_types(
+            self.encoder_attention_types,
+            self.encoder_nlayer,
+            "encoder_attention_types"
+        )
+
+    def get_decoder_attention_types(self) -> List[str]:
+        """Get expanded per-layer decoder self-attention types."""
+        return self._expand_attention_types(
+            self.decoder_attention_types,
+            self.decoder_nlayer,
+            "decoder_attention_types"
+        )
+
+    def get_decoder_cross_attention_types(self) -> List[str]:
+        """Get expanded per-layer decoder cross-attention types."""
+        return self._expand_attention_types(
+            self.decoder_cross_attention_types,
+            self.decoder_nlayer,
+            "decoder_cross_attention_types"
+        )
+
+    def _expand_attention_types(
+        self,
+        attention_types: Optional[List[str]],
+        num_layers: int,
+        field_name: str
+    ) -> List[str]:
+        """Expand short attention type patterns to full layer count.
+
+        Args:
+            attention_types: List of attention types or None
+            num_layers: Number of layers to expand to
+            field_name: Name of field for error messages
+
+        Returns:
+            List of attention types with length == num_layers
+        """
+        if attention_types is None:
+            if self.attention_mechanism is None:
+                raise ValueError(
+                    f"{field_name} is None and no global attention_mechanism fallback. "
+                    f"This should have been caught by _validate_attention_config()."
+                )
+            return [self.attention_mechanism] * num_layers
+
+        if len(attention_types) == num_layers:
+            return attention_types
+
+        if num_layers % len(attention_types) != 0:
+            raise ValueError(
+                f"{field_name} length ({len(attention_types)}) must divide "
+                f"layer count ({num_layers}) evenly for pattern repetition"
+            )
+
+        repeats = num_layers // len(attention_types)
+        return attention_types * repeats

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bbdf4f887074486d60cce4296cc42e4c00bbf3f3873ac562bf7c87e88269f9c9
+size 755547592

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca2b6bd503cd712bb538a5619a688dfee434b6a32a165635de4940e0a9fe22d4
+size 912333

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "mask_token": "<mask>",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 2048,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}