"""Generic Hugging Face text model -> RKLLM w8a8 for RK3588 (run on the 3090 WSL box).

  python convert_rkllm.py MODEL_DIR OUT.rkllm [--max-context 4096] [--device cuda] [--fla]

--fla: import flash-linear-attention first so `rwkv7` (fla-hub checkpoints) is registered with transformers. The fla
pinned by requirements_rwkv7.txt also registers names (bitnet, ...) that transformers 5.8 now ships, which raises
"already used"; those are skipped.
"""
import argparse
import os

from rkllm.api import RKLLM

ap = argparse.ArgumentParser()
ap.add_argument("model_dir")
ap.add_argument("out")
ap.add_argument("--max-context", type=int, default=4096)
ap.add_argument("--device", default="cuda")
ap.add_argument("--dtype", default="w8a8")
ap.add_argument("--fla", action="store_true")
args = ap.parse_args()
if args.fla:
    from transformers.models.auto import auto_factory, configuration_auto

    def skip_taken(orig):
        def register(self, key, value, exist_ok=False):
            try:
                orig(self, key, value, exist_ok=exist_ok)
            except ValueError as e:
                if "already" not in str(e):
                    raise
        return register

    configuration_auto._LazyConfigMapping.register = skip_taken(configuration_auto._LazyConfigMapping.register)
    auto_factory._LazyAutoMapping.register = skip_taken(auto_factory._LazyAutoMapping.register)
    import fla.models
    import transformers
    # rkllm looks the class up by its architectures[] name on the transformers module
    transformers.RWKV7ForCausalLM = fla.models.RWKV7ForCausalLM
llm = RKLLM()
assert llm.load_huggingface(model=args.model_dir, device=args.device) == 0, "load failed"
assert llm.build(do_quantization=True, optimization_level=1, quantized_dtype=args.dtype, quantized_algorithm="normal",
                 target_platform="rk3588", num_npu_core=3, max_context=args.max_context) == 0, "build failed"
assert llm.export_rkllm(args.out) == 0, "export failed"
print("EXPORTED", args.out, os.path.getsize(args.out))