#!/usr/bin/env python3 """Fix GGUF newline crash by adding byte-fallback tokens to the tokenizer. Problem: The SentencePiece Unigram tokenizer was trained without byte_fallback=True, so characters like \n have no token representation. llama.cpp crashes when it encounters these characters because there's no byte-fallback. Fix: 1. Add 256 byte-fallback tokens (<0x00> .. <0xFF>) to tokenizer.json 2. Resize model embeddings from 64000 -> 64256 3. Update config.json vocab_size 4. Copy tokenizer.model for proper GGUF conversion Usage: python scripts/fix_tokenizer_byte_fallback.py \ --input outputs/hf_checkpoint-best \ --output outputs/hf_checkpoint-best-fixed \ --sp_model tokenizer/korean_sp/tokenizer.model """ import argparse import json import shutil from pathlib import Path import torch from safetensors.torch import load_file, save_file BYTE_FALLBACK_COUNT = 256 BYTE_TOKEN_TEMPLATE = "<0x{:02X}>" def fix_tokenizer_json(input_path: Path, output_path: Path): """Add byte_fallback=True and 256 byte tokens to tokenizer.json.""" with open(input_path) as f: tok = json.load(f) model = tok["model"] vocab = model["vocab"] # list of [piece, score] original_size = len(vocab) # Enable byte_fallback model["byte_fallback"] = True # Add 256 byte tokens with very low score (they're fallback only) for i in range(BYTE_FALLBACK_COUNT): byte_token = BYTE_TOKEN_TEMPLATE.format(i) vocab.append([byte_token, 0.0]) new_size = len(vocab) print(f" Vocab: {original_size} -> {new_size} (+{BYTE_FALLBACK_COUNT} byte tokens)") print(f" byte_fallback: False -> True") # Also add byte tokens to added_tokens list added = tok.get("added_tokens", []) for i in range(BYTE_FALLBACK_COUNT): byte_token = BYTE_TOKEN_TEMPLATE.format(i) added.append({ "id": original_size + i, "content": byte_token, "single_word": False, "lstrip": False, "rstrip": False, "normalized": False, "special": True, }) tok["added_tokens"] = added with open(output_path, "w") as f: json.dump(tok, f, ensure_ascii=False, indent=2) return original_size, new_size def fix_config_json(input_path: Path, output_path: Path, new_vocab_size: int): """Update vocab_size in config.json.""" with open(input_path) as f: config = json.load(f) old_size = config["vocab_size"] config["vocab_size"] = new_vocab_size print(f" config.json vocab_size: {old_size} -> {new_vocab_size}") with open(output_path, "w") as f: json.dump(config, f, indent=2) def resize_embeddings(input_path: Path, output_path: Path, old_vocab: int, new_vocab: int, tie_embeddings: bool): """Resize embedding and lm_head weights to accommodate new tokens.""" print(f" Loading model weights from {input_path} ...") state_dict = load_file(str(input_path)) embed_key = "model.embed_tokens.weight" lm_head_key = "lm_head.weight" if embed_key not in state_dict: raise KeyError(f"{embed_key} not found in state_dict. Keys: {list(state_dict.keys())[:10]}") embed = state_dict[embed_key] print(f" embed_tokens shape: {embed.shape}") hidden_size = embed.shape[1] extra = new_vocab - old_vocab # Initialize new embeddings as mean of existing (better than random for byte tokens) mean_embed = embed.mean(dim=0, keepdim=True) # Add small noise to avoid identical embeddings noise = torch.randn(extra, hidden_size, dtype=embed.dtype) * 0.01 new_rows = mean_embed.expand(extra, -1) + noise new_embed = torch.cat([embed, new_rows], dim=0) state_dict[embed_key] = new_embed print(f" embed_tokens resized: {embed.shape} -> {new_embed.shape}") if tie_embeddings: # When tie_word_embeddings=True, lm_head shares embed_tokens # Remove lm_head if present (it will be tied automatically) if lm_head_key in state_dict: del state_dict[lm_head_key] print(f" lm_head removed (tie_word_embeddings=True)") else: if lm_head_key in state_dict: lm_head = state_dict[lm_head_key] mean_lm = lm_head.mean(dim=0, keepdim=True) noise_lm = torch.randn(extra, hidden_size, dtype=lm_head.dtype) * 0.01 new_lm = torch.cat([lm_head, mean_lm.expand(extra, -1) + noise_lm], dim=0) state_dict[lm_head_key] = new_lm print(f" lm_head resized: {lm_head.shape} -> {new_lm.shape}") print(f" Saving to {output_path} ...") save_file(state_dict, str(output_path)) def main(): parser = argparse.ArgumentParser(description="Fix tokenizer byte-fallback for GGUF") parser.add_argument("--input", type=Path, required=True, help="Input HF checkpoint dir") parser.add_argument("--output", type=Path, required=True, help="Output fixed HF checkpoint dir") parser.add_argument("--sp_model", type=Path, default=None, help="SentencePiece .model file to copy (for GGUF conversion)") args = parser.parse_args() input_dir = args.input output_dir = args.output if not input_dir.exists(): print(f"ERROR: Input directory not found: {input_dir}") return 1 output_dir.mkdir(parents=True, exist_ok=True) # Load config to check tie_word_embeddings with open(input_dir / "config.json") as f: config = json.load(f) old_vocab = config["vocab_size"] new_vocab = old_vocab + BYTE_FALLBACK_COUNT tie_embeddings = config.get("tie_word_embeddings", False) print(f"=== Byte-Fallback Fix ===") print(f"Input: {input_dir}") print(f"Output: {output_dir}") print(f"Old vocab: {old_vocab}, New vocab: {new_vocab}") print(f"tie_word_embeddings: {tie_embeddings}") print() # 1. Fix tokenizer.json print("[1/4] Fixing tokenizer.json ...") fix_tokenizer_json( input_dir / "tokenizer.json", output_dir / "tokenizer.json", ) # 2. Fix config.json print("[2/4] Fixing config.json ...") fix_config_json( input_dir / "config.json", output_dir / "config.json", new_vocab, ) # 3. Resize model weights print("[3/4] Resizing embeddings ...") resize_embeddings( input_dir / "model.safetensors", output_dir / "model.safetensors", old_vocab, new_vocab, tie_embeddings, ) # 4. Copy other files print("[4/4] Copying remaining files ...") for fname in ["tokenizer_config.json", "generation_config.json"]: src = input_dir / fname if src.exists(): shutil.copy2(src, output_dir / fname) print(f" Copied {fname}") # Copy SentencePiece model if provided (needed for GGUF conversion) if args.sp_model and args.sp_model.exists(): shutil.copy2(args.sp_model, output_dir / "tokenizer.model") print(f" Copied tokenizer.model from {args.sp_model}") elif (input_dir / "tokenizer.model").exists(): shutil.copy2(input_dir / "tokenizer.model", output_dir / "tokenizer.model") print(f" Copied tokenizer.model from input dir") # Update tokenizer_config.json to add added_tokens_decoder for byte tokens tc_path = output_dir / "tokenizer_config.json" if tc_path.exists(): with open(tc_path) as f: tc = json.load(f) added_tokens_decoder = tc.get("added_tokens_decoder", {}) for i in range(BYTE_FALLBACK_COUNT): token_id = old_vocab + i byte_token = BYTE_TOKEN_TEMPLATE.format(i) added_tokens_decoder[str(token_id)] = { "content": byte_token, "lstrip": False, "normalized": False, "rstrip": False, "single_word": False, "special": True, } tc["added_tokens_decoder"] = added_tokens_decoder with open(tc_path, "w") as f: json.dump(tc, f, indent=2) print(f" Updated tokenizer_config.json with {BYTE_FALLBACK_COUNT} byte tokens") print() print(f"=== Done! Fixed checkpoint at: {output_dir} ===") print(f"Next: python outputs/llama.cpp/convert_hf_to_gguf.py {output_dir} --outfile outputs/gguf/frankenstallm-3b-f16.gguf --outtype f16") return 0 if __name__ == "__main__": raise SystemExit(main())