#!/usr/bin/env python3
"""本地 GPU 推理服务（ONNX Runtime + CUDA）
RTX 3050 4GB 最优方案：Qwen2.5-0.5B ONNX + CUDAExecutionProvider
用法:
  python3 gpu_local_infer.py "你的问题" [max_tokens]
  python3 gpu_local_infer.py --bench   # 性能基准
"""
import os
import sys
import time
import subprocess
import numpy as np

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['LD_LIBRARY_PATH'] = '/home/muc/.local/lib:/usr/local/cuda/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')

import onnxruntime as ort
from transformers import AutoTokenizer

MODEL_DIR = os.path.expanduser("~/models/onnx-community-Qwen2.5-0.5B-Instruct")
ONNX_FILE = os.path.join(MODEL_DIR, "onnx", "model.onnx")

NUM_LAYERS = 24
NUM_KV_HEADS = 2
HEAD_DIM = 64

class LocalGPUInfer:
    def __init__(self, verbose=True):
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
        self.session = ort.InferenceSession(
            ONNX_FILE,
            providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
        )
        self.providers = self.session.get_providers()
        if verbose:
            print(f"✅ ONNX Runtime: {ort.__version__}")
            print(f"✅ Provider: {self.providers}")
            if 'CUDAExecutionProvider' in self.providers:
                print("🎉 CUDA 加速已启用！")

    def generate(self, prompt, max_new_tokens=100, temperature=0.7):
        inputs = self.tokenizer(prompt, return_tensors="np")
        input_ids = inputs['input_ids'].astype(np.int64)
        seq_len = input_ids.shape[1]
        batch = 1

        # 初始 past_key_values 为空
        past_key_values = {}
        for i in range(NUM_LAYERS):
            for kv in ['key', 'value']:
                past_key_values[f'past_key_values.{i}.{kv}'] = np.zeros(
                    (batch, NUM_KV_HEADS, 0, HEAD_DIM), dtype=np.float32
                )

        generated = []
        current_input_ids = input_ids
        past_len = 0

        for step in range(max_new_tokens):
            seq = current_input_ids.shape[1]
            ort_inputs = {
                'input_ids': current_input_ids,
                'attention_mask': np.ones((batch, past_len + seq), dtype=np.int64),
                'position_ids': np.arange(past_len, past_len + seq).reshape(1, seq).astype(np.int64),
            }
            for name, val in past_key_values.items():
                ort_inputs[name] = val

            outputs = self.session.run(None, ort_inputs)
            logits = outputs[0]

            if temperature <= 0:
                next_token = np.argmax(logits[:, -1, :], axis=-1)
            else:
                # 温度采样
                probs = logits[:, -1, :] / temperature
                probs = np.exp(probs - probs.max())
                probs = probs / probs.sum()
                next_token = np.array([np.random.choice(len(probs[0]), p=probs[0])], dtype=np.int64)

            next_token = next_token.astype(np.int64)
            generated.append(next_token[0])

            for i in range(NUM_LAYERS):
                past_key_values[f'past_key_values.{i}.key'] = outputs[1 + i*2]
                past_key_values[f'past_key_values.{i}.value'] = outputs[2 + i*2]

            current_input_ids = next_token.reshape(1, 1)
            past_len += seq

            if next_token[0] == self.tokenizer.eos_token_id:
                break

        return self.tokenizer.decode(generated, skip_special_tokens=True)

    def benchmark(self):
        """性能基准"""
        prompts = [
            "用一句话解释什么是过拟合",
            "你好",
            "中国的首都是哪里？",
        ]
        print("\n=== 性能基准 ===")
        for prompt in prompts:
            # 预热
            self.generate(prompt, max_new_tokens=10)
            # 计时
            start = time.time()
            result = self.generate(prompt, max_new_tokens=50)
            elapsed = time.time() - start
            print(f"\n输入: {prompt[:20]}...")
            print(f"输出: {result[:60]}...")
            print(f"耗时: {elapsed:.3f}s")

        print("\n=== GPU 状态 ===")
        r = subprocess.run(
            ["nvidia-smi", "--query-gpu=memory.used,memory.total,utilization.gpu", "--format=csv,noheader"],
            capture_output=True, text=True
        )
        print(f"  {r.stdout.strip()}")

def main():
    if len(sys.argv) < 2:
        print(__doc__)
        return

    if sys.argv[1] == "--bench":
        infer = LocalGPUInfer()
        infer.benchmark()
        return

    prompt = sys.argv[1]
    max_tokens = int(sys.argv[2]) if len(sys.argv) > 2 else 100

    infer = LocalGPUInfer()
    start = time.time()
    result = infer.generate(prompt, max_new_tokens=max_tokens)
    elapsed = time.time() - start
    print(f"\n输入: {prompt}")
    print(f"输出: {result}")
    print(f"\n耗时: {elapsed:.3f}s")

if __name__ == "__main__":
    main()