# -*- coding: utf-8 -*-
"""Qwen3.5-4B -> .rkllm (RK3588), Parameter identisch zum laufenden Qwen3-4B.

Damit der A/B-Test das MODELL misst und nicht die Konvertierung, sind alle
Stellschrauben so gesetzt wie beim produktiven qwen3-4b-w8a8g128-16k.rkllm:
w8a8 mit Gruppengroesse 128, 3 NPU-Kerne, max_context 16384 (das Maximum des
Toolkits; die Runtime auf dem Board nutzt davon 12288, siehe rkllm_server.cpp:52).
"""
import argparse, os, sys, time

from rkllm.api import RKLLM

ap = argparse.ArgumentParser()
ap.add_argument('--model', default=os.environ.get('MODEL_ID', 'Qwen/Qwen3.5-4B'),
                help='HF-Repo oder lokaler Pfad')
ap.add_argument('--out', default='/data/out/qwen35-4b-w8a8g128-16k.rkllm')
ap.add_argument('--dtype', default='w8a8_g128',
                help='w8a8, w8a8_g128, w8a8_g256, w8a8_g512, w4a16, w4a16_g32/64/128')
ap.add_argument('--max-context', type=int, default=16384, help='max. 16384, Vielfaches von 32')
ap.add_argument('--npu-cores', type=int, default=3)
ap.add_argument('--platform', default='rk3588')
ap.add_argument('--device', default='cuda', help='cuda oder cpu')
ap.add_argument('--dataset', default=None,
                help='Kalibrier-JSON [{"input":"","target":""}]. Ohne Angabe nimmt das '
                     'Toolkit seinen eingebauten Datensatz.')
ap.add_argument('--opt-level', type=int, default=1)
ap.add_argument('--hybrid-rate', type=float, default=0)
args = ap.parse_args()

os.makedirs(os.path.dirname(args.out), exist_ok=True)
t0 = time.time()

print('== Lade %s (device=%s)' % (args.model, args.device), flush=True)
llm = RKLLM()
ret = llm.load_huggingface(model=args.model, device=args.device)
if ret != 0:
    sys.exit('Laden fehlgeschlagen (ret=%d)' % ret)
print('   geladen nach %.0f s' % (time.time() - t0), flush=True)

print('== Quantisiere: %s, max_context=%d, %d NPU-Kerne, Ziel %s'
      % (args.dtype, args.max_context, args.npu_cores, args.platform), flush=True)
ret = llm.build(do_quantization=True,
                optimization_level=args.opt_level,
                quantized_dtype=args.dtype,
                quantized_algorithm='normal',
                target_platform=args.platform,
                num_npu_core=args.npu_cores,
                extra_qparams=None,
                dataset=args.dataset,
                hybrid_rate=args.hybrid_rate,
                max_context=args.max_context)
if ret != 0:
    sys.exit('Build fehlgeschlagen (ret=%d)' % ret)
print('   quantisiert nach %.0f s' % (time.time() - t0), flush=True)

print('== Exportiere nach %s' % args.out, flush=True)
ret = llm.export_rkllm(args.out)
if ret != 0:
    sys.exit('Export fehlgeschlagen (ret=%d)' % ret)

size = os.path.getsize(args.out) / (1024 ** 3)
print('FERTIG: %s (%.2f GB) in %.0f s' % (args.out, size, time.time() - t0))