hf/export_hf.py

"""Export TinyLiquid to a Hugging Face repo directory.

Produces:
  hf_repo/config.json               TinyLiquidConfig + HF fields
  hf_repo/model.safetensors         fp32 weights
  hf_repo/modeling_tinyliquid.py    self-contained trust_remote_code model
  hf_repo/tokenizer.json            (copy of our HF-format tokenizer)
  hf_repo/tokenizer_config.json     special tokens + chat template
  hf_repo/special_tokens_map.json
  hf_repo/generation_config.json
  hf_repo/quantized/q8.safetensors  our Q8 int8-storage weights (near-lossless)

Usage:
  .venv/bin/python hf/export_hf.py --ckpt ckpt/dpo --out hf_repo
"""

import argparse
import ast
import json
import shutil
from pathlib import Path

import torch
from safetensors.torch import save_file

from model.config import TinyLiquidConfig
from model.tiny_liquid import TinyLiquid
from model.utils import latest_ckpt
from model.quant import quantize_q8
from data.tokenizer import load_tokenizer, PERSONA_TOKENS


def _dataclass_fields(cfg_path: Path):
    tree = ast.parse(cfg_path.read_text(encoding="utf-8"))
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef) and node.name == "TinyLiquidConfig":
            fields = []
            for stmt in node.body:
                if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
                    default = None
                    if stmt.value is not None:
                        try:
                            default = ast.literal_eval(stmt.value)
                        except ValueError:
                            default = None
                    fields.append((stmt.target.id, default))
            return fields
    raise SystemExit("TinyLiquidConfig not found in config.py")


def _emit_config_class(cfg_path: Path) -> str:
    fields = _dataclass_fields(cfg_path)
    params = ", ".join(f"{n}={v!r}" if v is not None else f"{n}=None"
                       for n, v in fields)
    assigns = "\n".join(f"        self.{n} = {n}" for n, _ in fields)
    return f'''class TinyLiquidConfig(PretrainedConfig):
    """Architecture config for TinyLiquid (HF-compatible)."""

    model_type = "tiny_liquid"

    def __init__(
        self,
        {params},
        **kwargs,
    ):
        super().__init__(**kwargs)
{assigns}

    # --- standard aliases used by transformers internals ---
    @property
    def num_hidden_layers(self):
        return self.n_blocks

    @property
    def hidden_size(self):
        return self.d_model

    @property
    def num_attention_heads(self):
        return 1

    @property
    def max_position_embeddings(self):
        return self.max_seq_len
'''


def gen_modeling_file(dst: Path, tiny_arch: Path, cfg_arch: Path):
    """Emit a self-contained modeling_tinyliquid.py from our arch source."""
    src = tiny_arch.read_text(encoding="utf-8")

    header = '''"""TinyLiquid for Hugging Face (trust_remote_code).

Self-contained copy of the TinyLiquid non-transformer architecture
(basis-expansion liquid blocks with causal recurrence + gated MLP), wrapped
for transformers-compatible loading.

Load with:
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained("your-org/tiny-liquid-analyst")
    model = AutoModelForCausalLM.from_pretrained(
        "your-org/tiny-liquid-analyst", trust_remote_code=True)
    model.persona_id = 1  # 0 none, 1 analyst, 2 skeptic
"""
import json
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import CausalLMOutputWithPast
'''
    lines = [l for l in src.splitlines() if not l.startswith("from .config")]
    cut = next(i for i, l in enumerate(lines) if l.startswith("class RMSNorm"))
    imports_part = "\n".join(lines[:cut])
    body = "\n".join(lines[cut:])

    wrapper = '''

class TinyLiquidForCausalLM(PreTrainedModel):
    """transformers-compatible wrapper around TinyLiquid."""

    config_class = TinyLiquidConfig
    _tied_weights_keys = []
    all_tied_weights_keys = {}

    def __init__(self, config: TinyLiquidConfig):
        super().__init__(config)
        self.model = TinyLiquid(config)
        self.persona_id = 1  # default analyst; 0 none, 2 skeptic

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        persona_ids: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        if persona_ids is None:
            persona_ids = torch.tensor([self.persona_id], device=input_ids.device)
        logits = self.model(input_ids, persona_ids=persona_ids)
        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1), ignore_index=-100)
        return CausalLMOutputWithPast(
            loss=loss, logits=logits, past_key_values=None, hidden_states=None)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids, "persona_ids": kwargs.get("persona_ids")}
'''
    cfg_class = _emit_config_class(cfg_arch)
    dst.write_text(header + "\n\n" + cfg_class + "\n\n" + imports_part + "\n" + body + wrapper,
                   encoding="utf-8")
    print(f"wrote {dst}")


def build_config(sd_cfg: dict, vocab_size: int) -> TinyLiquidConfig:
    cfg = TinyLiquidConfig(vocab_size=vocab_size,
                           **{k: v for k, v in sd_cfg.items() if k != "vocab_size"})
    return cfg


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="ckpt/dpo")
    ap.add_argument("--out", default="hf_repo")
    ap.add_argument("--tok", default="data/tokenizer.json")
    args = ap.parse_args()

    out = Path(args.out)
    (out / "quantized").mkdir(parents=True, exist_ok=True)

    ckpt = latest_ckpt(args.ckpt)
    assert ckpt, f"no checkpoints in {args.ckpt}"
    sd = torch.load(ckpt, map_location="cpu")
    tok = load_tokenizer(args.tok)
    cfg = build_config(sd["config"], tok.get_vocab_size())
    model = TinyLiquid(cfg)
    model.load_state_dict(sd["model"])
    model.eval()

    tensors = {"model." + k: v.detach().contiguous() for k, v in model.state_dict().items()}
    save_file(tensors, out / "model.safetensors")
    print(f"wrote {out / 'model.safetensors'} ({sum(v.numel() for v in tensors.values())} params)")

    qs = quantize_q8(model)
    flat = {}
    for name, st in qs.items():
        flat[name + ".q"] = st["q"].contiguous()
        flat[name + ".scale"] = st["scale"].contiguous()
    save_file(flat, out / "quantized" / "q8.safetensors")
    print(f"wrote {out / 'quantized' / 'q8.safetensors'} ({len(qs)} linear layers)")

    import dataclasses
    hf_cfg = dataclasses.asdict(cfg)
    hf_cfg.update({
        "architectures": ["TinyLiquidForCausalLM"],
        "model_type": "tiny_liquid",
        "auto_map": {"AutoConfig": "modeling_tinyliquid.TinyLiquidConfig",
                "AutoModelForCausalLM": "modeling_tinyliquid.TinyLiquidForCausalLM"},
        "torch_dtype": "float32",
        "transformers_version": "4.x",
        "persona_tokens": PERSONA_TOKENS,
    })
    (out / "config.json").write_text(json.dumps(hf_cfg, indent=2), encoding="utf-8")

    shutil.copy(args.tok, out / "tokenizer.json")
    special = {}
    for name in ["<|endoftext|>", "<|user|>", "<|assistant|>", "<|scratchpad|>",
                 "<|final|>", "<|analyst|>", "<|skeptic|>"]:
        special[name] = tok.token_to_id(name)
    tok_cfg = {
        "tokenizer_class": "PreTrainedTokenizerFast",
        "model_max_length": cfg.max_seq_len,
        "bos_token": None,
        "eos_token": "<|endoftext|>",
        "unk_token": None,
        "pad_token": "<|endoftext|>",
        "added_tokens_decoder": {str(i): {"content": n, "special": True} for n, i in special.items()},
        "chat_template": (
            "{% for m in messages %}"
            "{% if m['role'] == 'system' %}<|analyst|>{% endif %}"
            "{% if m['role'] == 'user' %}<|user|>{{ m['content'] }}<|assistant|>{% endif %}"
            "{% if m['role'] == 'assistant' %}{{ m['content'] }}<|endoftext|>{% endif %}"
            "{% endfor %}"
        ),
    }
    (out / "tokenizer_config.json").write_text(json.dumps(tok_cfg, indent=2), encoding="utf-8")
    smap = {k: {"content": v, "lstrip": False, "rstrip": False, "single_word": False}
            for k, v in special.items()}
    (out / "special_tokens_map.json").write_text(json.dumps(smap, indent=2), encoding="utf-8")

    gen = {"max_new_tokens": 220, "temperature": 0.6, "top_k": 40,
           "repetition_penalty": 1.4, "do_sample": True}
    (out / "generation_config.json").write_text(json.dumps(gen, indent=2), encoding="utf-8")

    gen_modeling_file(out / "modeling_tinyliquid.py",
                      Path("model/tiny_liquid.py"), Path("model/config.py"))
    print(f"export complete -> {out}")


if __name__ == "__main__":
    main()