"""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()