# SPDX-License-Identifier: Apache-2.0
# Independent streaming GGUF converter, prepared 2026-10-03.
# llama.cpp/GGUF: Copyright (c) 2023-2026 The ggml authors (MIT).
# Architecture reference: Copyright 2026 Aleph Alpha GmbH (Apache-2.0).
# See THIRD_PARTY_NOTICES.txt for attribution and separate license scopes.
"""Streaming Kolibri BF16 -> GGUF using the official gguf writer.
Sources: Aleph-Alpha/aleph-alpha-inference @049a6a7; ggml-org/llama.cpp @edd6e2b.
"""
import argparse, json, math, re, shutil, struct, sys, time
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).parent / 'llama.cpp' / 'gguf-py'))
import gguf

def metadata(w, cfg, model):
    w.add_name('Kolibri-1')
    w.add_block_count(cfg['num_hidden_layers'])
    w.add_context_length(cfg['max_position_embeddings'])
    w.add_embedding_length(cfg['hidden_size'])
    w.add_feed_forward_length(cfg['moe_intermediate_size'])
    w.add_head_count(cfg['num_attention_heads'])
    w.add_head_count_kv(cfg['num_key_value_heads'])
    w.add_key_length(cfg['head_dim']); w.add_value_length(cfg['head_dim'])
    w.add_rope_dimension_count(cfg['head_dim'])
    w.add_rope_freq_base(cfg['rope_theta'])
    w.add_rope_scaling_type(gguf.RopeScalingType.NONE)
    w.add_layer_norm_rms_eps(cfg['rms_norm_eps'])
    w.add_expert_count(cfg['num_experts'])
    w.add_expert_used_count(cfg['num_experts_per_tok'])
    w.add_expert_shared_count(1)
    w.add_expert_feed_forward_length(cfg['moe_intermediate_size'])
    w.add_expert_shared_feed_forward_length(cfg['shared_expert_intermediate_size'])
    w.add_expert_weights_norm(cfg['norm_topk_prob'])
    w.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
    w.add_sliding_window(cfg['sliding_window'])
    w.add_sliding_window_pattern([x == 'sliding_attention' for x in cfg['layer_types']])
    w.add_vocab_size(cfg['vocab_size'])
    w.add_file_type(gguf.LlamaFileType.MOSTLY_BF16)
    t = json.loads((model/'tokenizer.json').read_text(encoding='utf-8'))
    tc = json.loads((model/'tokenizer_config.json').read_text(encoding='utf-8'))
    if t['normalizer'] is not None or t['model']['type'] != 'BPE' or t['model'].get('ignore_merges', False):
        raise ValueError('Unsupported tokenizer settings')
    expected = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
    if t['pre_tokenizer']['pretokenizers'][0]['pattern']['Regex'] != expected:
        raise ValueError('Unexpected tokenizer regex')
    vocab = t['model']['vocab']
    tokens = [f'[PAD{i}]' for i in range(cfg['vocab_size'])]
    types = [gguf.TokenType.UNUSED] * len(tokens)
    for token, idx in vocab.items():
        tokens[idx] = token; types[idx] = gguf.TokenType.NORMAL
    for token in t['added_tokens']:
        if token['lstrip'] or token['rstrip'] or token['single_word']:
            raise ValueError('Unsupported added token behavior')
        tokens[token['id']] = token['content']
        types[token['id']] = gguf.TokenType.CONTROL if token['special'] else gguf.TokenType.USER_DEFINED
    w.add_tokenizer_model('gpt2'); w.add_tokenizer_pre('qwen2')
    w.add_token_list(tokens); w.add_token_types(types)
    w.add_token_merges([' '.join(x) if isinstance(x, list) else x for x in t['model']['merges']])
    w.add_add_bos_token(False); w.add_add_eos_token(False); w.add_add_space_prefix(False)
    w.add_eos_token_id(cfg['eos_token_id']); w.add_pad_token_id(cfg['pad_token_id'])
    w.add_eot_token_id(cfg['eos_token_id'])
    w.add_chat_template(tc.get('chat_template'))

def tensor_plan(model, cfg):
    idx = json.loads((model/'model.safetensors.index.json').read_text())['weight_map']
    headers = {}; records = {}
    for filename in sorted(set(idx.values())):
        path = model / filename
        with path.open('rb') as f:
            n = struct.unpack('<Q',f.read(8))[0]
            if n > 16*1024*1024: raise ValueError('Invalid header length')
            header = json.loads(f.read(n))
        for name, entry in header.items():
            if name == '__metadata__': continue
            if name in records or idx.get(name) != filename: raise ValueError('Index mismatch')
            if entry['dtype'] != 'BF16': raise ValueError('Expected BF16')
            start,end = entry['data_offsets']; shape=entry['shape']
            if end-start != math.prod(shape)*2 or 8+n+end > path.stat().st_size:
                raise ValueError('Tensor size mismatch')
            records[name]=(path,8+n+start,end-start,shape)
    if records.keys() != idx.keys(): raise ValueError('Missing tensors')
    d=cfg['hidden_size']; q=cfg['num_attention_heads']*cfg['head_dim']
    kv=cfg['num_key_value_heads']*cfg['head_dim']; e=cfg['num_experts']
    ff=cfg['moe_intermediate_size']; fs=cfg['shared_expert_intermediate_size']; v=cfg['vocab_size']
    plan=[]; used=set()
    def add(target, sources, shape, f32=False):
        chunks=[]
        for source in sources:
            if source in used: raise ValueError('Duplicate source')
            used.add(source); chunks.append(records[source])
        expected_shape=shape[1:] if len(sources)>1 else shape
        if any(c[3] != list(expected_shape) for c in chunks):
            raise ValueError(f'Unexpected shape: {target}')
        plan.append((target,tuple(shape),chunks,f32))
    add('token_embd.weight',['model.embed_tokens.weight'],[v,d])
    add('output_norm.weight',['model.norm.weight'],[d],True)
    add('output.weight',['lm_head.weight'],[v,d])
    specs={
        'self_attn.q_proj.weight':('attn_q.weight',[q,d]),
        'self_attn.k_proj.weight':('attn_k.weight',[kv,d]),
        'self_attn.v_proj.weight':('attn_v.weight',[kv,d]),
        'self_attn.o_proj.weight':('attn_output.weight',[d,q]),
        'self_attn.q_norm.weight':('attn_q_norm.weight',[cfg['head_dim']]),
        'self_attn.k_norm.weight':('attn_k_norm.weight',[cfg['head_dim']]),
        'input_layernorm.weight':('attn_norm.weight',[d]),
        'post_attn_norm.weight':('post_attention_norm.weight',[d]),
        'post_attention_layernorm.weight':('ffn_norm.weight',[d]),
        'post_ffn_norm.weight':('post_ffw_norm.weight',[d]),
        'mlp.gate.weight':('ffn_gate_inp.weight',[e,d]),
        'moe.router.expert_bias':('exp_probs_b.bias',[e]),
        'mlp.shared_experts.gate_proj.weight':('ffn_gate_shexp.weight',[fs,d]),
        'mlp.shared_experts.up_proj.weight':('ffn_up_shexp.weight',[fs,d]),
        'mlp.shared_experts.down_proj.weight':('ffn_down_shexp.weight',[d,fs]),
    }
    # Resolve standardized tensor names from the official GGUF definitions.
    specs['post_attn_norm.weight']= (gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.ATTN_POST_NORM]+'.weight',[d])
    specs['post_ffn_norm.weight']= (gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.FFN_POST_NORM]+'.weight',[d])
    specs['moe.router.expert_bias']= (gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.FFN_EXP_PROBS_B]+'.bias',[e])
    for i in range(cfg['num_hidden_layers']):
        prefix=f'model.layers.{i}.'
        for source,(target,shape) in specs.items():
            dest=target.format(bid=i) if '{bid}' in target else f'blk.{i}.{target}'
            add(dest,[prefix+source],shape,len(shape)==1 or source=='mlp.gate.weight')
        for proj,ty,shape in [('gate_proj',gguf.MODEL_TENSOR.FFN_GATE_EXP,[e,ff,d]),('up_proj',gguf.MODEL_TENSOR.FFN_UP_EXP,[e,ff,d]),('down_proj',gguf.MODEL_TENSOR.FFN_DOWN_EXP,[e,d,ff])]:
            dest=gguf.TENSOR_NAMES[ty].format(bid=i)+'.weight'
            add(dest,[prefix+f'mlp.experts.{j}.{proj}.weight' for j in range(e)],shape)
    if used != records.keys(): raise ValueError(f'Unmapped tensors: {records.keys()-used}')
    return plan

def convert(model, out, dry=False, all_f32=False):
    cfg=json.loads((model/'config.json').read_text())
    if cfg['architectures'] != ['Kolibri1ForCausalLM']: raise ValueError('Wrong architecture')
    plan=tensor_plan(model,cfg)
    if all_f32: plan=[(name,shape,chunks,True) for name,shape,chunks,_ in plan]
    total=sum(math.prod(shape)*(4 if f32 else 2) for _,shape,_,f32 in plan)
    print(f'Validated {sum(len(x[2]) for x in plan)} source tensors -> {len(plan)} GGUF tensors; {total/1e9:.3f} GB',flush=True)
    if dry: return
    if out.exists() or out.with_suffix(out.suffix+'.partial').exists(): raise FileExistsError(out)
    out.parent.mkdir(parents=True,exist_ok=True)
    if shutil.disk_usage(out.parent).free < total + 2*1024**3: raise RuntimeError('Insufficient disk space')
    partial=out.with_suffix(out.suffix+'.partial')
    w=gguf.GGUFWriter(partial,'kolibri1'); metadata(w,cfg,model)
    for name,shape,chunks,f32 in plan:
        w.add_tensor_info(name,shape,np.dtype('float32' if f32 else 'uint16'),math.prod(shape)*(4 if f32 else 2),raw_dtype=gguf.GGMLQuantizationType.F32 if f32 else gguf.GGMLQuantizationType.BF16)
    w.write_header_to_file(); w.write_kv_data_to_file(); w.write_ti_data_to_file()
    fp=w.fout[0]; handles={}; start=time.monotonic(); done=0; last=0
    try:
        for name,shape,chunks,f32 in plan:
            w.write_padding(fp,fp.tell())
            size=0
            for path,offset,nbytes,_ in chunks:
                if path not in handles: handles[path]=path.open('rb')
                src=handles[path]; src.seek(offset)
                if f32:
                    raw=src.read(nbytes)
                    if len(raw)!=nbytes: raise EOFError(path)
                    a=(np.frombuffer(raw,dtype='<u2').astype(np.uint32)<<16).view(np.float32)
                    fp.write(a.tobytes()); size+=a.nbytes
                else:
                    left=nbytes
                    while left:
                        raw=src.read(min(left,8*1024*1024))
                        if not raw: raise EOFError(path)
                        fp.write(raw); left-=len(raw); size+=len(raw)
            w.write_padding(fp,size); done+=size
            now=time.monotonic()
            if now-last>10 or done==total:
                print(f'{done/1e9:.2f}/{total/1e9:.2f} GB, {done/(now-start)/1e6:.0f} MB/s, tensor {name}',flush=True); last=now
    finally:
        for f in handles.values(): f.close()
        w.close()
    partial.rename(out)
    print(f'COMPLETE: {out}',flush=True)

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('model',type=Path);p.add_argument('output',type=Path);p.add_argument('--dry-run',action='store_true');a=p.parse_args()
    convert(a.model,a.output,a.dry_run)