# lora_basic.py
"""LoRA 基础使用示例"""

from vllm import LLM, SamplingParams

# 启用 LoRA 支持
llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",
    enable_lora=True,           # 启用 LoRA
    max_lora_rank=64,           # 最大 LoRA rank
    max_loras=4,                # 最大同时加载的 LoRA 数量
    max_model_len=4096,
    gpu_memory_utilization=0.9,
)

# 加载并使用 LoRA 适配器
from vllm.lora.request import LoRARequest

lora_adapter = LoRARequest(
    lora_name="medical-lora",          # 适配器名称
    lora_int_id=1,                      # 整数 ID
    lora_path="/data/lora/medical-lora", # 适配器路径
)

sampling_params = SamplingParams(temperature=0.7, max_tokens=256)

# 使用 LoRA 生成
outputs = llm.generate(
    ["患者出现持续性头痛，可能是什么原因？"],
    sampling_params,
    lora_request=lora_adapter,
)

print(outputs[0].outputs[0].text)

# 启动带 LoRA 支持的 API 服务
vllm serve Qwen/Qwen2.5-7B-Instruct \
    --enable-lora \
    --max-lora-rank 64 \
    --max-loras 4 \
    --lora-modules \
        medical-lora=/data/lora/medical-lora \
        legal-lora=/data/lora/legal-lora \
        finance-lora=/data/lora/finance-lora

# 使用指定 LoRA 发送请求
curl http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "qwen-7b",
        "messages": [{"role": "user", "content": "诊断：慢性胃炎的治疗方案"}],
        "max_tokens": 300,
        "temperature": 0.7,
        "model_extra": {
            "lora_name": "medical-lora"
        }
    }'

# Python 客户端使用 LoRA
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="none")

# 指定 LoRA 适配器
response = client.chat.completions.create(
    model="qwen-7b",
    messages=[{"role": "user", "content": "分析这个合同条款的风险。"}],
    max_tokens=300,
    extra_body={"lora_name": "legal-lora"},
)
print(response.choices[0].message.content)

# multi_lora.py
"""多 LoRA 同时加载"""

from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest

llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",
    enable_lora=True,
    max_lora_rank=64,
    max_loras=8,            # 最多同时加载 8 个 LoRA
    max_model_len=4096,
)

# 定义多个 LoRA 适配器
loras = {
    "medical": LoRARequest("medical", 1, "/data/lora/medical"),
    "legal": LoRARequest("legal", 2, "/data/lora/legal"),
    "finance": LoRARequest("finance", 3, "/data/lora/finance"),
    "code": LoRARequest("code", 4, "/data/lora/code"),
}

# 不同请求使用不同的 LoRA
prompts_and_loras = [
    ("头痛的常见原因有哪些？", loras["medical"]),
    ("分析合同中的违约条款。", loras["legal"]),
    ("评估这个投资方案的风险。", loras["finance"]),
    ("写一个快速排序算法。", loras["code"]),
    ("你好，今天天气怎么样？", None),  # 不使用 LoRA
]

sampling_params = SamplingParams(temperature=0.7, max_tokens=256)

for prompt, lora in prompts_and_loras:
    output = llm.generate([prompt], sampling_params, lora_request=lora)
    lora_name = lora.lora_name if lora else "base"
    print(f"[{lora_name}] {output[0].outputs[0].text[:100]}...")

# hot_swap.py
"""运行时动态加载 LoRA"""

import time
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
from vllm.lora.request import LoRARequest

async def demo_hot_swap():
    # 创建引擎
    engine_args = AsyncEngineArgs(
        model="Qwen/Qwen2.5-7B-Instruct",
        enable_lora=True,
        max_lora_rank=64,
        max_loras=4,
    )
    engine = AsyncLLMEngine.from_engine_args(engine_args)
    
    sampling_params = SamplingParams(temperature=0.7, max_tokens=100)
    
    # 初始 LoRA
    lora_v1 = LoRARequest("model-v1", 1, "/data/lora/medical-v1")
    
    result = await engine.generate(
        "头痛的诊断流程是什么？",
        sampling_params,
        request_id="req-1",
        lora_request=lora_v1,
    )
    print(f"V1: {result.outputs[0].text[:100]}")
    
    # 热切换到新版本的 LoRA
    lora_v2 = LoRARequest("model-v2", 2, "/data/lora/medical-v2")
    
    result = await engine.generate(
        "头痛的诊断流程是什么？",
        sampling_params,
        request_id="req-2",
        lora_request=lora_v2,
    )
    print(f"V2: {result.outputs[0].text[:100]}")

# lora_manager.py
"""LoRA 适配器管理"""

import os
from pathlib import Path

class LoRAManager:
    """管理 LoRA 适配器的加载和切换"""
    
    def __init__(self, lora_base_dir: str):
        self.base_dir = Path(lora_base_dir)
        self.loaded_adapters: dict[str, LoRARequest] = {}
        self._next_id = 1
    
    def discover_adapters(self) -> list[str]:
        """发现可用的 LoRA 适配器"""
        adapters = []
        for d in self.base_dir.iterdir():
            if d.is_dir() and (d / "adapter_config.json").exists():
                adapters.append(d.name)
        return adapters
    
    def load_adapter(self, name: str) -> LoRARequest:
        """加载指定的适配器"""
        if name in self.loaded_adapters:
            return self.loaded_adapters[name]
        
        adapter_path = self.base_dir / name
        if not adapter_path.exists():
            raise FileNotFoundError(f"适配器 {name} 不存在")
        
        lora_request = LoRARequest(
            lora_name=name,
            lora_int_id=self._next_id,
            lora_path=str(adapter_path),
        )
        self._next_id += 1
        self.loaded_adapters[name] = lora_request
        
        return lora_request
    
    def get_adapter(self, name: str) -> LoRARequest:
        """获取已加载的适配器"""
        if name not in self.loaded_adapters:
            return self.load_adapter(name)
        return self.loaded_adapters[name]

# 使用示例
manager = LoRAManager("/data/lora-adapters")
print("可用适配器:", manager.discover_adapters())

medical_lora = manager.load_adapter("medical-v3")

# train_lora.py
"""训练 LoRA 适配器示例"""

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from peft import LoraConfig, get_peft_model, TaskType
from datasets import load_dataset
from trl import SFTTrainer

# 1. 加载基础模型
model_name = "Qwen/Qwen2.5-7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 2. 配置 LoRA
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,                    # LoRA rank
    lora_alpha=32,           # 缩放因子
    lora_dropout=0.05,       # Dropout
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    bias="none",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# 输出: trainable params: 13,107,200 || all params: 7,614,939,136 || trainable%: 0.172

# 3. 准备数据
dataset = load_dataset("json", data_files="train_data.jsonl")

# 4. 训练
training_args = TrainingArguments(
    output_dir="./lora-output",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    warmup_ratio=0.1,
    logging_steps=10,
    save_strategy="epoch",
    fp16=True,
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    tokenizer=tokenizer,
    dataset_text_field="text",
    max_seq_length=2048,
)

trainer.train()

# 5. 保存 LoRA 适配器
model.save_pretrained("./medical-lora-adapter")
tokenizer.save_pretrained("./medical-lora-adapter")

# awq_lora.py
"""AWQ 量化 + LoRA 组合使用"""

from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest

# 使用 AWQ 量化模型作为基础模型
llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct-AWQ",
    quantization="awq",
    enable_lora=True,
    max_lora_rank=64,
    max_loras=4,
    max_model_len=4096,
)

lora = LoRARequest("medical", 1, "/data/lora/medical")
sampling_params = SamplingParams(temperature=0.7, max_tokens=256)

# 注意：LoRA 适配器需要在原始 FP16 模型上训练
# vLLM 会自动处理精度转换
outputs = llm.generate(
    ["患者症状：头痛、发热、乏力。请给出诊断建议。"],
    sampling_params,
    lora_request=lora,
)
print(outputs[0].outputs[0].text)

# 性能优化配置
llm = LLM(
    model="model",
    enable_lora=True,
    max_lora_rank=32,       # 不要设太大，16-32 通常足够
    max_loras=4,            # 根据实际需求设置
    max_model_len=4096,     # 减小可增加并发
    gpu_memory_utilization=0.9,
    # LoRA 相关的额外显存需要预留
)

# 同时运行两个版本的 LoRA，对比效果
lora_v1 = LoRARequest("model-v1", 1, "/lora/sentiment-v1")
lora_v2 = LoRARequest("model-v2", 2, "/lora/sentiment-v2")

# 50/50 分流
import random
for prompt in test_prompts:
    lora = lora_v1 if random.random() < 0.5 else lora_v2
    result = llm.generate([prompt], params, lora_request=lora)