update dockerfile: enttypoint outside docker

2026-02-04 19:48:59 +07:00
parent c7887b70b3
commit 9f1327cf1c
3 changed files with 101 additions and 40 deletions
--- a/entrypoint.py
+++ b/entrypoint.py
@@ -1,37 +1,91 @@
-from rkllm.api import RKLLM
-from tqdm import tqdm
-import torch
-from torch import nn
+#!/usr/bin/env python3
+import argparse
 import os
+from rkllm.api import RKLLM

-modelpath = '/models/Qwen3-4B-Instruct-2507-F16/Qwen3-4B-Instruct-2507-F16.gguf'
-output = '/output/Qwen3-4B-Instruct-2507-F16-Q.rkllm'
-os.makedirs("/output", exist_ok=True)
+def parse_args():
+    parser = argparse.ArgumentParser(description='Компиляция LLM для Rockchip NPU')
+    
+    # Обязательные параметры
+    parser.add_argument('--model', type=str, required=True,
+                        help='Путь к модели (Hugging Face или GGUF)')
+    parser.add_argument('--output', type=str, required=True,
+                        help='Путь для сохранения .rkllm файла')
+    
+    # Опциональные параметры с умными значениями по умолчанию
+    parser.add_argument('--format', type=str, choices=['huggingface', 'gguf'], default='huggingface',
+                        help='Формат входной модели (по умолчанию: huggingface)')
+    parser.add_argument('--enable-q', type=str, choices=['True', 'False'], default='True',
+                        help='Включить квантование (по умолчаниюЖ  True)')                    
+    parser.add_argument('--quant', type=str, default='w8a8_g256',
+                        choices=['w4a16', 'w4a16_g32', 'w4a16_g64', 'w4a16_g128', 
+                                 'w8a8', 'w8a8_g128', 'w8a8_g256', 'w8a8_g512'],
+                        help='Тип квантования (по умолчанию: w8a8_g256)')
+    parser.add_argument('--platform', type=str, default='rk3588',
+                        choices=['rk3588', 'rk3576', 'rk3562', 'rv1126b'],
+                        help='Целевая платформа (по умолчанию: rk3588)')
+    parser.add_argument('--npu-cores', type=int, default=3, choices=[1, 2, 3],
+                        help='Количество ядер NPU (по умолчанию: 3)')
+    parser.add_argument('--max-context', type=int, default=8192,
+                        help='Максимальная длина контекста (по умолчанию: 8192)')
+    
+    return parser.parse_args()

-print(f"Загрузка модели из: {modelpath}")
+def main():
+    args = parse_args()
+    os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True)
+    
+    print(f"🔧 Компиляция модели:")
+    print(f"   Модель: {args.model}")
+    print(f"   Формат: {args.format}")
+    print(f"   Включено квантование: {args.enable_q}")
+    print(f"   Квантование: {args.quant}")
+    print(f"   Платформа: {args.platform}")
+    print(f"   Ядра NPU: {args.npu_cores}")
+    print(f"   Длина контекста: {args.max_context}")
+    
+    llm = RKLLM()
+    
+    # Загрузка в зависимости от формата
+    if args.format == 'huggingface':
+        ret = llm.load_huggingface(
+            model=args.model,
+            device='cpu',
+            trust_remote_code=True,
+            quantized_model=False
+        )
+    else:  # gguf
+        ret = llm.load_gguf(model=args.model)
+    
+    if ret != 0:
+        print(f"❌ Ошибка загрузки модели (код {ret})")
+        exit(ret)
+    
+    # Определение алгоритма квантования
+    algo = 'grq' if 'w4a16' in args.quant else 'normal'
+    
+    ret = llm.build(
+        do_quantization=args.enable_q,
+        optimization_level=1,
+        quantized_dtype=args.quant,
+        quantized_algorithm=algo,
+        target_platform=args.platform,
+        num_npu_core=args.npu_cores,
+        max_context=args.max_context
+    )
+    
+    if ret != 0:
+        print(f"❌ Ошибка компиляции (код {ret})")
+        exit(ret)
+    
+    ret = llm.export_rkllm(args.output)
+    if ret != 0:
+        print(f"❌ Ошибка экспорта (код {ret})")
+        exit(ret)
+    
+    size_gb = os.path.getsize(args.output) / (1024**3)
+    print(f"\n✅ Успешно! Сохранено: {args.output}")
+    print(f"📦 Размер: {size_gb:.2f} GB")

-llm = RKLLM()
-ret = llm.load_gguf(model=modelpath)
-# ret = llm.load_huggingface(model=modelpath, model_lora = None, device='cpu')
-
-if ret != 0:
-    print(f"❌ Ошибка загрузки модели: {ret}")
-    exit(ret)
-
-print("Компиляция для RK3588 (NPU)...")
-ret = llm.build(do_quantization=True, optimization_level=1, quantized_dtype='w8a8',
-                quantized_algorithm='normal', target_platform='rk3588', num_npu_core=3, extra_qparams=None) 
-
-
-if ret != 0:
-    print(f"❌ Ошибка компиляции: {ret}")
-    exit(ret)
-
-# Export rkllm model
-ret = llm.export_rkllm(output)
-if ret != 0:
-    print(f"❌ Ошибка экспорта: {ret}")
-    exit(ret)
-
-print(f"✅ Модель скомпилирована: {output}")
-print(f"Размер: {os.path.getsize(output) / 1024**3:.2f} GB")
+if __name__ == '__main__':
+    main()