#!/usr/bin/env perl
#
# ascii_condense_gguf.pl
#
# Reproduces the "ASCII-Condensed" vocabulary-pruning trick used to build
# bsaleh03/Qwen3.8-27B-ASCII-Condensed from unsloth/Qwen3.8-27B-GGUF.
#
# What it does to a llama.cpp GGUF file, IN QUANTIZED SPACE (no dequantize/
# requantize step, ever):
#
#   1. Reads tokenizer.ggml.{tokens,token_type,scores,merges} from the GGUF
#      metadata.
#   2. Decodes every token's GPT2 byte-level-BPE surface form back to the
#      raw bytes it represents, and decides whether to KEEP it:
#        - any CONTROL-type token (special/added tokens)         -> keep
#        - any token that decodes to exactly one raw byte         -> keep
#          (these are the atomic "byte-fallback" tokens that guarantee
#          every possible byte, even non-ASCII ones, stays representable)
#        - any other token whose full decoded byte sequence is
#          all-ASCII (< 0x80)                                     -> keep
#        - everything else (multi-byte tokens that decode to any
#          non-ASCII byte)                                        -> drop
#   3. Every token referenced by a *_token_id metadata field (bos/eos/eot/
#      eom/unk/sep/pad/mask/fim_*) is force-kept, then all of those ids are
#      remapped to their new positions.
#   4. tokenizer.ggml.merges rules are kept only if BOTH parent tokens
#      survive step 2.
#   5. token_embd.weight and output.weight are row-gathered: for each
#      surviving vocabulary row, the *raw quantized bytes* of that row are
#      copied byte-for-byte (using the ggml block_size/type_size of
#      whatever quant type the tensor already is) into the new tensor.
#      Rows are streamed in original order, so this never needs to hold
#      more than one row in memory.
#   6. Every other tensor is copied byte-for-byte, unchanged, straight
#      from the input file to the output file (streamed in chunks), so
#      the result is bit-identical for all non-vocabulary tensors.
#   7. Tensor offsets/padding are recomputed to satisfy GGUF alignment,
#      since the two shrunk tensors change where later tensors start.
#
# Usage:
#   perl ascii_condense_gguf.pl input.gguf [output.gguf]
#
# If output.gguf is omitted, it defaults to
#   <input-without-.gguf>-ASCII-Condensed.gguf
#
use strict;
use warnings;
use v5.10;
use Encode qw(decode FB_QUIET);
use Fcntl qw(SEEK_SET);

$| = 1;

# ---------------------------------------------------------------------
# ggml type table: id => [ name, block_size (elements/block), type_size (bytes/block) ]
# This is the standard, stable ggml.c type_traits table (llama.cpp master).
# ---------------------------------------------------------------------
my %GGML_TYPE = (
    0  => ['F32',      1,   4],
    1  => ['F16',      1,   2],
    2  => ['Q4_0',     32,  18],
    3  => ['Q4_1',     32,  20],
    6  => ['Q5_0',     32,  22],
    7  => ['Q5_1',     32,  24],
    8  => ['Q8_0',     32,  34],
    9  => ['Q8_1',     32,  40],
    10 => ['Q2_K',     256, 84],
    11 => ['Q3_K',     256, 110],
    12 => ['Q4_K',     256, 144],
    13 => ['Q5_K',     256, 176],
    14 => ['Q6_K',     256, 210],
    15 => ['Q8_K',     256, 292],
    16 => ['IQ2_XXS',  256, 66],
    17 => ['IQ2_XS',   256, 74],
    18 => ['IQ3_XXS',  256, 98],
    19 => ['IQ1_S',    256, 50],
    20 => ['IQ4_NL',   32,  18],
    21 => ['IQ3_S',    256, 110],
    22 => ['IQ2_S',    256, 82],
    23 => ['IQ4_XS',   256, 136],
    24 => ['I8',       1,   1],
    25 => ['I16',      1,   2],
    26 => ['I32',      1,   4],
    27 => ['I64',      1,   8],
    28 => ['F64',      1,   8],
    29 => ['IQ1_M',    256, 56],
    30 => ['BF16',     1,   2],
    34 => ['TQ1_0',    256, 54],
    35 => ['TQ2_0',    256, 66],
);

# GGUF metadata value type ids
use constant {
    T_UINT8 => 0, T_INT8 => 1, T_UINT16 => 2, T_INT16 => 3,
    T_UINT32 => 4, T_INT32 => 5, T_FLOAT32 => 6, T_BOOL => 7,
    T_STRING => 8, T_ARRAY => 9, T_UINT64 => 10, T_INT64 => 11,
    T_FLOAT64 => 12,
};

# llama_token_type enum (tokenizer.ggml.token_type values)
use constant { TT_NORMAL => 1, TT_UNKNOWN => 2, TT_CONTROL => 3,
               TT_USER_DEFINED => 4, TT_UNUSED => 5, TT_BYTE => 6 };

# ---------------------------------------------------------------------
# Low level binary readers (sequential filehandle, little-endian)
# ---------------------------------------------------------------------
sub r_bytes {
    my ($fh, $n) = @_;
    return '' if $n == 0;
    my $buf;
    my $got = read($fh, $buf, $n);
    die "unexpected EOF (wanted $n bytes)\n" unless defined($got) && $got == $n;
    return $buf;
}
sub r_u8   { unpack('C',   r_bytes($_[0], 1)) }
sub r_i8   { unpack('c',   r_bytes($_[0], 1)) }
sub r_u16  { unpack('v',   r_bytes($_[0], 2)) }
sub r_i16  { unpack('s<',  r_bytes($_[0], 2)) }
sub r_u32  { unpack('V',   r_bytes($_[0], 4)) }
sub r_i32  { unpack('l<',  r_bytes($_[0], 4)) }
sub r_u64  { unpack('Q<',  r_bytes($_[0], 8)) }
sub r_i64  { unpack('q<',  r_bytes($_[0], 8)) }
sub r_f32  { unpack('f<',  r_bytes($_[0], 4)) }
sub r_f64  { unpack('d<',  r_bytes($_[0], 8)) }
sub r_bool { r_u8($_[0]) ? 1 : 0 }
sub r_str  { my $fh = shift; my $len = r_u64($fh); return $len ? r_bytes($fh, $len) : ''; }

sub read_value {
    my ($fh, $type) = @_;
    return r_u8($fh)   if $type == T_UINT8;
    return r_i8($fh)   if $type == T_INT8;
    return r_u16($fh)  if $type == T_UINT16;
    return r_i16($fh)  if $type == T_INT16;
    return r_u32($fh)  if $type == T_UINT32;
    return r_i32($fh)  if $type == T_INT32;
    return r_f32($fh)  if $type == T_FLOAT32;
    return r_bool($fh) if $type == T_BOOL;
    return r_str($fh)  if $type == T_STRING;
    return r_u64($fh)  if $type == T_UINT64;
    return r_i64($fh)  if $type == T_INT64;
    return r_f64($fh)  if $type == T_FLOAT64;
    if ($type == T_ARRAY) {
        my $atype = r_u32($fh);
        my $alen  = r_u64($fh);
        my @data;
        for (1 .. $alen) { push @data, read_value($fh, $atype); }
        return { _array => 1, type => $atype, data => \@data };
    }
    die "unknown GGUF value type $type\n";
}

# ---------------------------------------------------------------------
# Low level binary writers (mirror of the above)
# ---------------------------------------------------------------------
sub w_bytes { my ($fh, $data) = @_; print {$fh} $data or die "write failed: $!\n"; }
sub w_u8   { w_bytes($_[0], pack('C',  $_[1])) }
sub w_i8   { w_bytes($_[0], pack('c',  $_[1])) }
sub w_u16  { w_bytes($_[0], pack('v',  $_[1])) }
sub w_i16  { w_bytes($_[0], pack('s<', $_[1])) }
sub w_u32  { w_bytes($_[0], pack('V',  $_[1])) }
sub w_i32  { w_bytes($_[0], pack('l<', $_[1])) }
sub w_u64  { w_bytes($_[0], pack('Q<', $_[1])) }
sub w_i64  { w_bytes($_[0], pack('q<', $_[1])) }
sub w_f32  { w_bytes($_[0], pack('f<', $_[1])) }
sub w_f64  { w_bytes($_[0], pack('d<', $_[1])) }
sub w_bool { w_u8($_[0], $_[1] ? 1 : 0) }
sub w_str  { my ($fh, $s) = @_; w_u64($fh, length($s)); w_bytes($fh, $s) if length($s); }

sub write_value {
    my ($fh, $type, $val) = @_;
    return w_u8($fh, $val)   if $type == T_UINT8;
    return w_i8($fh, $val)   if $type == T_INT8;
    return w_u16($fh, $val)  if $type == T_UINT16;
    return w_i16($fh, $val)  if $type == T_INT16;
    return w_u32($fh, $val)  if $type == T_UINT32;
    return w_i32($fh, $val)  if $type == T_INT32;
    return w_f32($fh, $val)  if $type == T_FLOAT32;
    return w_bool($fh, $val) if $type == T_BOOL;
    return w_str($fh, $val)  if $type == T_STRING;
    return w_u64($fh, $val)  if $type == T_UINT64;
    return w_i64($fh, $val)  if $type == T_INT64;
    return w_f64($fh, $val)  if $type == T_FLOAT64;
    if ($type == T_ARRAY) {
        my $atype = $val->{type};
        my $data  = $val->{data};
        w_u32($fh, $atype);
        w_u64($fh, scalar(@$data));
        write_value($fh, $atype, $_) for @$data;
        return;
    }
    die "unknown GGUF value type $type on write\n";
}

# ---------------------------------------------------------------------
# GPT2 byte-level BPE alphabet (bytes_to_unicode) and its inverse.
# Qwen (like GPT2/GPT-NeoX/most BPE tokenizers) stores vocabulary strings
# in this "surface" form, e.g. a literal space (0x20) is stored as the
# character U+0120 ('Ġ'). To know what raw bytes a token really represents
# we have to invert this mapping.
# ---------------------------------------------------------------------
sub build_unicode_to_byte {
    my @bs;
    push @bs, $_ for (ord('!') .. ord('~'));   # 33..126
    push @bs, $_ for (0xA1 .. 0xAC);           # ¡..¬
    push @bs, $_ for (0xAE .. 0xFF);           # ®..ÿ
    my %is_base = map { $_ => 1 } @bs;
    my @cs = @bs;
    my $n = 0;
    for my $b (0 .. 255) {
        next if $is_base{$b};
        push @bs, $b;
        push @cs, 256 + $n;
        $n++;
    }
    my %u2b;
    @u2b{@cs} = @bs;
    return \%u2b;
}

# Decode a token's raw on-disk (UTF-8) surface bytes into the list of
# original bytes it represents. Returns undef for any codepoint that
# cannot be mapped back (treated conservatively as "not ASCII" by the
# caller, since we can't prove it's safe).
sub decode_token_bytes {
    my ($surface, $u2b) = @_;
    my $str = decode('UTF-8', $surface, FB_QUIET);
    my @out;
    for my $ch (split //, $str) {
        my $cp = ord($ch);
        if (exists $u2b->{$cp}) {
            push @out, $u2b->{$cp};
        } elsif ($cp < 0x80) {
            # Literal ASCII character not part of the byte-alphabet
            # (this is how special/control token text like
            # "<|im_start|>" is stored - verbatim, not byte-mapped).
            push @out, $cp;
        } else {
            push @out, undef;   # unmappable -> caller treats as non-ASCII
        }
    }
    return \@out;
}

sub token_survives {
    my ($surface, $ttype, $u2b) = @_;
    return 1 if defined($ttype) && $ttype == TT_CONTROL;
    return 1 if defined($ttype) && $ttype == TT_BYTE;
    my $bytes = decode_token_bytes($surface, $u2b);
    return 1 if scalar(@$bytes) == 1;               # atomic byte token
    for my $b (@$bytes) {
        return 0 if !defined($b) || $b >= 0x80;      # any non-ASCII byte -> drop
    }
    return 1;
}

# ---------------------------------------------------------------------
# Stream-copy N bytes from one filehandle to another, in chunks.
# ---------------------------------------------------------------------
sub copy_bytes {
    my ($in, $out, $n, $chunk) = @_;
    $chunk //= 16 * 1024 * 1024;
    while ($n > 0) {
        my $want = $n < $chunk ? $n : $chunk;
        my $buf;
        my $got = read($in, $buf, $want);
        die "unexpected EOF while copying tensor data\n" unless defined($got) && $got == $want;
        w_bytes($out, $buf);
        $n -= $want;
    }
}

# ===========================================================================
# MAIN
# ===========================================================================

my ($in_path, $out_path) = @ARGV;
if (!defined $in_path) {
    die "Usage: $0 input.gguf [output.gguf]\n";
}
if (!defined $out_path) {
    $out_path = $in_path;
    if ($out_path =~ s/\.gguf$//i) {
        $out_path .= '-ASCII-Condensed.gguf';
    } else {
        $out_path .= '.ASCII-Condensed.gguf';
    }
}
die "Refusing to overwrite input file\n" if $out_path eq $in_path;

open(my $in, '<:raw', $in_path) or die "can't open $in_path: $!\n";

# --- header ---
my $magic = r_bytes($in, 4);
die "$in_path is not a GGUF file (bad magic)\n" unless $magic eq 'GGUF';
my $version = r_u32($in);
die "unsupported GGUF version $version (expected 2 or 3)\n" unless $version == 2 || $version == 3;
my $tensor_count = r_u64($in);
my $kv_count     = r_u64($in);

say STDERR "GGUF v$version, $tensor_count tensors, $kv_count metadata entries";

# --- metadata ---
my @kvs;       # [ {key=>, type=>, val=>}, ... ]  in original order
my %kv_index;  # key -> index into @kvs
for my $i (0 .. $kv_count - 1) {
    my $key   = r_str($in);
    my $vtype = r_u32($in);
    my $val   = read_value($in, $vtype);
    push @kvs, { key => $key, type => $vtype, val => $val };
    $kv_index{$key} = $#kvs;
}

my $alignment = 32;
if (exists $kv_index{'general.alignment'}) {
    $alignment = $kvs[$kv_index{'general.alignment'}]{val};
}

# --- tensor infos ---
my @tensors;   # [ {name=>, ndim=>, ne=>[...], type=>, offset=>}, ... ] original order
for my $i (0 .. $tensor_count - 1) {
    my $name = r_str($in);
    my $ndim = r_u32($in);
    my @ne;
    push @ne, r_u64($in) for (1 .. $ndim);
    my $ttype  = r_u32($in);
    my $offset = r_u64($in);
    push @tensors, { name => $name, ndim => $ndim, ne => \@ne, type => $ttype, offset => $offset };
}

# absolute file position where tensor data begins (after alignment padding)
my $header_end = tell($in);
my $old_data_start = $header_end + ((-$header_end) % $alignment + $alignment) % $alignment;
# simpler equivalent:
$old_data_start = $header_end + ((($alignment - ($header_end % $alignment)) % $alignment));

# ---------------------------------------------------------------------
# Locate tokenizer fields
# ---------------------------------------------------------------------
sub need_kv { my $k = shift; die "missing required metadata key '$k'\n" unless exists $kv_index{$k}; return $kvs[$kv_index{$k}]; }

my $tok_model = exists $kv_index{'tokenizer.ggml.model'} ? $kvs[$kv_index{'tokenizer.ggml.model'}]{val} : '(unknown)';
say STDERR "tokenizer.ggml.model = $tok_model";
if ($tok_model ne 'gpt2') {
    warn "WARNING: tokenizer.ggml.model is '$tok_model', not 'gpt2'. This script assumes\n" .
         "         a GPT2-style byte-level BPE vocabulary (which Qwen uses). Proceeding,\n" .
         "         but the ASCII/byte-token detection may be wrong for other tokenizer types.\n";
}

my $tokens_kv = need_kv('tokenizer.ggml.tokens');
my @tokens    = @{ $tokens_kv->{val}{data} };
my $n_vocab_old = scalar @tokens;

my @ttypes;
if (exists $kv_index{'tokenizer.ggml.token_type'}) {
    @ttypes = @{ $kvs[$kv_index{'tokenizer.ggml.token_type'}]{val}{data} };
}

my $merges_idx = $kv_index{'tokenizer.ggml.merges'};
my @merges = $merges_idx ? @{ $kvs[$merges_idx]{val}{data} } : ();

say STDERR "vocab size (old) = $n_vocab_old, merge rules (old) = " . scalar(@merges);

# ---------------------------------------------------------------------
# Decide which vocab rows survive
# ---------------------------------------------------------------------
my $u2b = build_unicode_to_byte();
my @survive = (0) x $n_vocab_old;
for my $i (0 .. $n_vocab_old - 1) {
    $survive[$i] = token_survives($tokens[$i], $ttypes[$i], $u2b) ? 1 : 0;
}

# Force-keep every token referenced by a *_token_id metadata field, so we
# never orphan bos/eos/pad/etc.
my @special_id_keys = qw(
    tokenizer.ggml.bos_token_id
    tokenizer.ggml.eos_token_id
    tokenizer.ggml.eot_token_id
    tokenizer.ggml.eom_token_id
    tokenizer.ggml.unknown_token_id
    tokenizer.ggml.seperator_token_id
    tokenizer.ggml.padding_token_id
    tokenizer.ggml.cls_token_id
    tokenizer.ggml.mask_token_id
    tokenizer.ggml.fim_pre_token_id
    tokenizer.ggml.fim_suf_token_id
    tokenizer.ggml.fim_mid_token_id
    tokenizer.ggml.fim_pad_token_id
    tokenizer.ggml.fim_rep_token_id
    tokenizer.ggml.fim_sep_token_id
    tokenizer.ggml.prefix_token_id
    tokenizer.ggml.suffix_token_id
    tokenizer.ggml.middle_token_id
);
for my $k (@special_id_keys) {
    next unless exists $kv_index{$k};
    my $id = $kvs[$kv_index{$k}]{val};
    next if !defined($id) || $id < 0 || $id >= $n_vocab_old;
    $survive[$id] = 1;
}
if (exists $kv_index{'tokenizer.ggml.suppress_tokens'}) {
    for my $id (@{ $kvs[$kv_index{'tokenizer.ggml.suppress_tokens'}]{val}{data} }) {
        $survive[$id] = 1 if $id >= 0 && $id < $n_vocab_old;
    }
}

# old id -> new id map (-1 if dropped)
my @old2new = (-1) x $n_vocab_old;
my $n_vocab_new = 0;
for my $i (0 .. $n_vocab_old - 1) {
    if ($survive[$i]) { $old2new[$i] = $n_vocab_new++; }
}
say STDERR "vocab size (new) = $n_vocab_new (dropped " . ($n_vocab_old - $n_vocab_new) . ")";

# ---------------------------------------------------------------------
# Build filtered tokenizer arrays
# ---------------------------------------------------------------------
sub filter_parallel_array {
    my ($arrval, $survive_ref) = @_;
    return undef unless $arrval;
    my @data = @{ $arrval->{data} };
    my @out;
    for my $i (0 .. $#data) {
        push @out, $data[$i] if $survive_ref->[$i];
    }
    return { _array => 1, type => $arrval->{type}, data => \@out };
}

my $new_tokens_val = filter_parallel_array($tokens_kv->{val}, \@survive);
my $new_ttype_val;
if (exists $kv_index{'tokenizer.ggml.token_type'}) {
    $new_ttype_val = filter_parallel_array($kvs[$kv_index{'tokenizer.ggml.token_type'}]{val}, \@survive);
}
my $new_scores_val;
if (exists $kv_index{'tokenizer.ggml.scores'}) {
    $new_scores_val = filter_parallel_array($kvs[$kv_index{'tokenizer.ggml.scores'}]{val}, \@survive);
}

# Merge filtering: keep a rule only if both parent tokens survive.
my %surface2id;
for my $i (0 .. $n_vocab_old - 1) {
    $surface2id{$tokens[$i]} = $i unless exists $surface2id{$tokens[$i]};
}
my @new_merges;
for my $m (@merges) {
    my ($a, $b) = split(/ /, $m, 2);
    next unless defined($a) && defined($b);
    next unless exists $surface2id{$a} && exists $surface2id{$b};
    next unless $survive[$surface2id{$a}] && $survive[$surface2id{$b}];
    push @new_merges, $m;
}
say STDERR "merge rules (new) = " . scalar(@new_merges) . " (dropped " . (scalar(@merges) - scalar(@new_merges)) . ")";

# ---------------------------------------------------------------------
# Apply all metadata edits into @kvs (in place, preserving key order)
# ---------------------------------------------------------------------
$kvs[$kv_index{'tokenizer.ggml.tokens'}]{val} = $new_tokens_val;
if ($new_ttype_val) { $kvs[$kv_index{'tokenizer.ggml.token_type'}]{val} = $new_ttype_val; }
if ($new_scores_val) { $kvs[$kv_index{'tokenizer.ggml.scores'}]{val} = $new_scores_val; }
if (defined $merges_idx) {
    $kvs[$merges_idx]{val} = { _array => 1, type => T_STRING, data => \@new_merges };
}
for my $k (@special_id_keys) {
    next unless exists $kv_index{$k};
    my $old_id = $kvs[$kv_index{$k}]{val};
    next if !defined($old_id) || $old_id < 0 || $old_id >= $n_vocab_old;
    $kvs[$kv_index{$k}]{val} = $old2new[$old_id];
}
if (exists $kv_index{'tokenizer.ggml.suppress_tokens'}) {
    my $arr = $kvs[$kv_index{'tokenizer.ggml.suppress_tokens'}]{val};
    $arr->{data} = [ map { $old2new[$_] } grep { $_ >= 0 && $_ < $n_vocab_old && $survive[$_] } @{ $arr->{data} } ];
}
# Any "{arch}.vocab_size" scalar key
for my $kv (@kvs) {
    if ($kv->{key} =~ /\.vocab_size$/) {
        say STDERR "updating '$kv->{key}': $kv->{val} -> $n_vocab_new";
        $kv->{val} = $n_vocab_new;
    }
}

# ---------------------------------------------------------------------
# Figure out which tensors need row-gathering
# ---------------------------------------------------------------------
my %VOCAB_TENSOR_NAMES = map { $_ => 1 } ('token_embd.weight', 'output.weight');
my %tensor_plan;  # name -> { row_bytes, new_ne1, old_ne1 }

for my $t (@tensors) {
    next unless $VOCAB_TENSOR_NAMES{$t->{name}};
    my ($tname, $blk, $tsize) = @{ $GGML_TYPE{$t->{type}} // die "unknown ggml type id $t->{type} for tensor $t->{name}\n" };
    my $row_elems = $t->{ne}[0];
    die "$t->{name}: row length $row_elems not a multiple of $tname block size $blk\n"
        if $row_elems % $blk != 0;
    my $row_bytes = ($row_elems / $blk) * $tsize;
    my $old_rows  = $t->{ne}[1];
    if ($old_rows != $n_vocab_old) {
        warn "WARNING: $t->{name} has $old_rows rows but vocab has $n_vocab_old tokens - skipping row-gather for this tensor.\n";
        next;
    }
    $tensor_plan{$t->{name}} = {
        type_name => $tname, row_bytes => $row_bytes, old_rows => $old_rows, new_rows => $n_vocab_new,
    };
    say STDERR "will row-gather $t->{name}: type=$tname row_bytes=$row_bytes rows $old_rows -> $n_vocab_new";
}
for my $n (keys %VOCAB_TENSOR_NAMES) {
    say STDERR "note: tensor '$n' not found in file (tied or absent) - nothing to do for it."
        unless grep { $_->{name} eq $n } @tensors;
}

# ---------------------------------------------------------------------
# Compute new tensor sizes / offsets (in original tensor order)
# ---------------------------------------------------------------------
sub tensor_nelements { my $t = shift; my $n = 1; $n *= $_ for @{ $t->{ne} }; return $n; }
sub tensor_bytesize {
    my $t = shift;
    my ($tname, $blk, $tsize) = @{ $GGML_TYPE{$t->{type}} // die "unknown ggml type id $t->{type}\n" };
    my $n = tensor_nelements($t);
    die "$t->{name}: total element count not divisible by block size\n" if $n % $blk != 0;
    return ($n / $blk) * $tsize;
}

my $running = 0;
for my $t (@tensors) {
    my $plan = $tensor_plan{$t->{name}};
    my $old_offset = $t->{offset};
    my $old_ne1;
    if ($plan) {
        $old_ne1 = $t->{ne}[1];
        $t->{new_ne} = [ $t->{ne}[0], $plan->{new_rows}, (@{ $t->{ne} } > 2 ? @{ $t->{ne} }[2 .. $#{ $t->{ne} }] : ()) ];
        $t->{new_size} = $plan->{row_bytes} * $plan->{new_rows};
    } else {
        $t->{new_ne}   = $t->{ne};
        $t->{new_size} = tensor_bytesize($t);
    }
    $t->{old_abs_offset} = $old_data_start + $old_offset;
    $t->{old_size}       = $plan ? $plan->{row_bytes} * $plan->{old_rows} : $t->{new_size};

    $t->{new_offset} = $running;
    $running += $t->{new_size};
    my $pad = ($alignment - ($running % $alignment)) % $alignment;
    $t->{pad_after} = $pad;
    $running += $pad;
}
my $new_data_total = $running;
say STDERR "new tensor-data section: $new_data_total bytes";

# ===========================================================================
# WRITE OUTPUT
# ===========================================================================
open(my $out, '>:raw', $out_path) or die "can't create $out_path: $!\n";

w_bytes($out, 'GGUF');
w_u32($out, $version);
w_u64($out, $tensor_count);         # tensor count is unchanged - same tensors, just resized
w_u64($out, scalar @kvs);           # kv count unchanged - same keys, some values edited
for my $kv (@kvs) {
    w_str($out, $kv->{key});
    w_u32($out, $kv->{type});
    write_value($out, $kv->{type}, $kv->{val});
}
for my $t (@tensors) {
    w_str($out, $t->{name});
    w_u32($out, scalar @{ $t->{new_ne} });
    w_u64($out, $_) for @{ $t->{new_ne} };
    w_u32($out, $t->{type});
    w_u64($out, $t->{new_offset});
}

my $new_header_end = tell($out);
my $new_pad = ($alignment - ($new_header_end % $alignment)) % $alignment;
w_bytes($out, "\x00" x $new_pad) if $new_pad;
my $new_data_start = tell($out);
die "internal error: new_data_start misaligned\n" if $new_data_start % $alignment != 0;

# --- write tensor data, tensor by tensor, in original order ---
for my $t (@tensors) {
    my $plan = $tensor_plan{$t->{name}};
    if ($plan) {
        say STDERR "writing $t->{name} (row-gathered) ...";
        seek($in, $t->{old_abs_offset}, SEEK_SET) or die "seek failed: $!\n";
        my $row_bytes = $plan->{row_bytes};
        for my $i (0 .. $plan->{old_rows} - 1) {
            my $row = r_bytes($in, $row_bytes);
            w_bytes($out, $row) if $survive[$i];
        }
    } else {
        seek($in, $t->{old_abs_offset}, SEEK_SET) or die "seek failed: $!\n";
        copy_bytes($in, $out, $t->{old_size});
    }
    w_bytes($out, "\x00" x $t->{pad_after}) if $t->{pad_after};
}

close($out) or die "error closing $out_path: $!\n";
close($in);

say STDERR "done -> $out_path";
say STDERR "  vocab: $n_vocab_old -> $n_vocab_new rows";
say STDERR "  merges: " . scalar(@merges) . " -> " . scalar(@new_merges);
say STDERR "  tensors row-gathered: " . join(', ', sort keys %tensor_plan);