101_参数高效微调_QLoRA技术深度解析与实践

# NF4量化的核心代码逻辑
def quantize_to_nf4(weights):
    # 计算权重的均值和标准差
    mean = weights.mean()
    std = weights.std()
    
    # 标准化权重到标准正态分布
    normalized_weights = (weights - mean) / std
    
    # 将标准化权重映射到NF4编码空间
    # NF4编码点设计为最佳匹配标准正态分布
    scale = calculate_optimal_scale(normalized_weights)
    zero_point = calculate_optimal_zero_point(normalized_weights)
    
    # 执行量化
    quantized = ((normalized_weights - zero_point) / scale).round().clamp(-8, 7).astype(np.int8)
    
    return quantized, scale, zero_point, mean, std

# 低秩矩阵初始化代码
def initialize_lora_adapters(model, r=16, lora_alpha=32, target_modules=None):
    if target_modules is None:
        target_modules = ["q_proj", "v_proj"]  # 通常只对注意力模块应用
        
    for name, module in model.named_modules():
        if any(mm in name for mm in target_modules) and hasattr(module, "weight"):
            # 获取原始权重形状
            in_features, out_features = module.weight.shape
            
            # 创建低秩适应矩阵
            A = torch.nn.Parameter(torch.zeros(in_features, r))
            B = torch.nn.Parameter(torch.zeros(r, out_features))
            
            # Kaiming初始化A矩阵
            torch.nn.init.kaiming_uniform_(A, a=math.sqrt(5))
            
            # 保存适配器参数
            module.lora_A = A
            module.lora_B = B
            module.lora_alpha = lora_alpha
            module.lora_r = r
            
            # 冻结原始权重
            module.weight.requires_grad = False
            
    return model

# QLoRA前向传播代码
def lora_forward(module, input):
    # 原始权重的前向计算
    output = F.linear(input, module.weight, module.bias)
    
    # 低秩适配器的贡献
    lora_output = F.linear(input, module.lora_B @ module.lora_A, None) * (module.lora_alpha / module.lora_r)
    
    return output + lora_output

# 创建并激活虚拟环境
conda create -n qlora-env python=3.10
conda activate qlora-env

# 安装PyTorch（CUDA版本）
pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 --index-url https://download.pytorch.org/whl/cu121

# 安装PEFT库（包含QLoRA实现）
pip install peft==0.10.0

# 安装Transformers
git clone https://github.com/huggingface/transformers.git
cd transformers
pip install -e .
cd ..

# 安装其他依赖
pip install bitsandbytes==0.43.1 datasets==2.18.0 accelerate==0.30.1 scipy==1.13.0

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import prepare_model_for_kbit_training

# 配置4位量化
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# 加载预训练模型和分词器
model_name = "meta-llama/Llama-2-7b-hf"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 为k-bit训练准备模型
model = prepare_model_for_kbit_training(model)

from peft import LoraConfig, get_peft_model

# 配置LoRA参数
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# 应用LoRA适配器
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 打印可训练参数数量

from datasets import load_dataset

# 加载数据集
dataset = load_dataset("timdettmers/openassistant-guanaco")

# 数据预处理函数
def preprocess_function(examples):
    # 格式化输入输出对
    instructions = examples["instruction"]
    inputs = examples["input"]
    outputs = examples["output"]
    
    texts = []
    for i in range(len(instructions)):
        text = f"### 指令:\n{instructions[i]}"
        if inputs[i]:
            text += f"\n### 输入:\n{inputs[i]}"
        text += f"\n### 输出:\n{outputs[i]}"
        texts.append(text)
    
    # 编码文本
    return tokenizer(texts, padding="max_length", truncation=True, max_length=1024)

# 处理数据集
tokenized_dataset = dataset.map(preprocess_function, batched=True)

import transformers

trainer = transformers.Trainer(
    model=model,
    train_dataset=tokenized_dataset["train"],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        warmup_steps=100,
        max_steps=1000,
        learning_rate=2e-4,
        fp16=True,
        logging_steps=10,
        output_dir="./qlora-output",
        save_strategy="steps",
        save_steps=100,
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

model.config.use_cache = False  # 禁用缓存以进行训练
trainer.train()

# 针对不同模型的目标模块配置
target_modules_config = {
    "llama": ["q_proj", "v_proj"],
    "mistral": ["q_proj", "v_proj", "gate_proj"],
    "falcon": ["query_key_value", "dense"],
    "bloom": ["query_key_value"],
    "gpt2": ["c_attn"],
    "gptj": ["q_proj", "v_proj"],
    "gpt_neox": ["query_key_value"],
    "opt": ["q_proj", "v_proj"]
}

# 选择适合特定模型的配置
model_family = "llama"  # 可替换为其他模型家族
lora_config = LoraConfig(
    r=32,
    lora_alpha=64,
    target_modules=target_modules_config[model_family],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# 多模态模型QLoRA配置
multimodal_lora_config = LoraConfig(
    r=32,
    lora_alpha=64,
    target_modules=[
        "text_model.encoder.layer.23.self_attn.q_proj",
        "text_model.encoder.layer.23.self_attn.v_proj",
        "vision_model.encoder.layer.23.mlp.fc1",
        "vision_model.encoder.layer.23.mlp.fc2",
        "multi_modal_projector"
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="MULTIMODAL_CAUSAL_LM"
)

# 梯度检查点优化
model.gradient_checkpointing_enable()
model.config.use_cache = False  # 训练时必须禁用缓存

# 配置分布式训练
trainer = transformers.Trainer(
    model=model,
    train_dataset=tokenized_dataset["train"],
    args=transformers.TrainingArguments(
        # 基本参数...
        gradient_checkpointing=True,
        gradient_accumulation_steps=8,
        ddp_find_unused_parameters=False,
        optim="paged_adamw_8bit",  # 使用8位优化器进一步节省内存
    ),
    # 其他配置...
)

# 启用混合精度训练
training_args = transformers.TrainingArguments(
    # 其他参数...
    fp16=True,  # 或bf16=True
    tf32=True,  # 对支持TF32的GPU启用
)

# 动态批量大小计算
def find_optimal_batch_size(model, max_memory=40):  # max_memory单位为GB
    # 初始批量大小
    batch_size = 1
    memory_used = 0
    
    while True:
        try:
            # 创建虚拟输入
            input_ids = torch.randint(0, 32000, (batch_size, 1024)).cuda()
            
            # 前向传播
            outputs = model(input_ids, labels=input_ids)
            
            # 反向传播
            outputs.loss.backward()
            
            # 清理
            torch.cuda.empty_cache()
            
            # 增加批量大小
            batch_size *= 2
        except RuntimeError as e:
            if "out of memory" in str(e):
                # 内存不足，回退到前一个可用大小
                batch_size = batch_size // 2
                return max(1, batch_size)
            else:
                raise e

# 更精细的梯度检查点控制
def set_gradient_checkpointing(model, checkpoint_ratio=0.5):
    # 计算需要应用检查点的层数
    total_layers = len(model.model.layers)
    checkpoint_layers = int(total_layers * checkpoint_ratio)
    
    # 对指定层应用梯度检查点
    for i, layer in enumerate(model.model.layers):
        # 通常只在中间层应用梯度检查点
        if i % (total_layers // checkpoint_layers) == 0:
            layer.gradient_checkpointing_enable()

# 启用TF32和BF16混合精度
torch.backends.cuda.matmul.allow_tf32 = True  # 对支持TF32的GPU启用

# 配置学习率调度器
training_args = transformers.TrainingArguments(
    # 其他参数...
    learning_rate=2e-4,
    warmup_steps=100,  # 预热步数
    lr_scheduler_type="cosine",  # 使用余弦衰减
    weight_decay=0.01,
)

# 配置梯度裁剪
training_args = transformers.TrainingArguments(
    # 其他参数...
    gradient_clipping=1.0,
)

# 使用AdamW优化器的变体
training_args = transformers.TrainingArguments(
    # 其他参数...
    optim="adamw_torch",  # 或者"paged_adamw_32bit"/"paged_adamw_8bit"
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
)

# 数据过滤函数
def filter_high_quality_data(examples, min_length=100, max_length=2000):
    # 过滤过短或过长的样本
    filtered = []
    for text in examples["text"]:
        if min_length <= len(text) <= max_length:
            filtered.append(True)
        else:
            filtered.append(False)
    return filtered

# 应用过滤
high_quality_dataset = dataset.filter(filter_high_quality_data)

# 简单的数据增强函数
def augment_data(examples):
    augmented = []
    for text in examples["text"]:
        # 原始文本
        augmented.append(text)
        
        # 同义词替换增强（示例）
        if len(text) > 100:
            words = text.split()
            # 替换少量词语为同义词（实际实现需要同义词词典）
            augmented.append(" ".join(words))
    return {"text": augmented}

# 应用数据增强
augmented_dataset = dataset.map(augment_data, batched=True)

# 分层采样以保持类别平衡
from collections import Counter

def stratified_sampling(dataset, target_column, sample_ratio=0.1):
    # 计算每个类别的样本数
    class_counts = Counter(dataset[target_column])
    
    # 为每个类别确定采样数量
    sample_indices = []
    for label, count in class_counts.items():
        # 对每个类别采样固定比例
        label_indices = [i for i, x in enumerate(dataset[target_column]) if x == label]
        sample_size = max(1, int(len(label_indices) * sample_ratio))
        sample_indices.extend(random.sample(label_indices, sample_size))
    
    return dataset.select(sample_indices)

# 保存LoRA适配器权重
model.save_pretrained("qlora-adapter")

from transformers import AutoModelForCausalLM
from peft import PeftModel

# 加载基础模型
base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    return_dict=True,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 加载并合并LoRA适配器
merged_model = PeftModel.from_pretrained(
    base_model,
    "qlora-adapter",
    device_map="auto"
)
merged_model = merged_model.merge_and_unload()

# 保存合并后的模型
merged_model.save_pretrained("merged-model")
tokenizer.save_pretrained("merged-model")

# 使用ONNX进行推理优化
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 加载合并后的模型
model = AutoModelForCausalLM.from_pretrained(
    "merged-model",
    load_in_8bit=True,  # 加载为8位量化模型进行推理
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("merged-model")

# 启用Flash Attention加速
model = model.to_bettertransformer()

# 启用KV缓存
model.config.use_cache = True

# 批量推理示例
def batch_inference(model, tokenizer, prompts, max_length=100, batch_size=4):
    results = []
    
    # 按批次处理提示
    for i in range(0, len(prompts), batch_size):
        batch = prompts[i:i+batch_size]
        
        # 编码批次
        inputs = tokenizer(
            batch,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512
        ).to(model.device)
        
        # 生成文本
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_length=max_length,
                temperature=0.7,
                do_sample=True,
                num_return_sequences=1
            )
        
        # 解码结果
        for output in outputs:
            results.append(tokenizer.decode(output, skip_special_tokens=True))
    
    return results

# 服务器端部署优化
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载优化后的模型
model = AutoModelForCausalLM.from_pretrained(
    "merged-model",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained("merged-model")

# 预热模型（减少首次推理延迟）
def warmup_model(model, tokenizer, iterations=5):
    warmup_prompt = "Hello, how are you?"
    for _ in range(iterations):
        inputs = tokenizer(warmup_prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            model.generate(**inputs, max_length=50)

# 启用自动混合精度
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    warmup_model(model, tokenizer)

# 完整的聊天机器人微调示例
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset
import transformers
import os

# 设置随机种子以确保可重复性
torch.manual_seed(42)

# 1. 加载和量化模型
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token  # 设置填充token

# 2. 准备模型进行k位训练
model = prepare_model_for_kbit_training(model)

# 3. 配置LoRA参数
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 打印可训练参数数量

# 4. 加载和处理数据集
dataset = load_dataset("timdettmers/openassistant-guanaco")

# 数据预处理函数
def preprocess_function(examples):
    # 格式化聊天数据
    conversations = []
    for instruction, input_text, output in zip(examples["instruction"], examples["input"], examples["output"]):
        # 构建对话历史
        conversation = f"<s>[INST] {instruction}\n"
        if input_text.strip():
            conversation += f"{input_text}\n"
        conversation += f"[/INST] {output}</s>"
        conversations.append(conversation)
    
    # 编码文本
    return tokenizer(conversations, padding="max_length", truncation=True, max_length=512)

# 处理数据集
tokenized_dataset = dataset.map(preprocess_function, batched=True)

# 5. 配置训练参数
training_args = transformers.TrainingArguments(
    output_dir="./qlora-llama2-chatbot",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    logging_steps=10,
    max_steps=1000,
    save_strategy="steps",
    save_steps=100,
    warmup_steps=100,
    bf16=True,
    gradient_checkpointing=True,
    optim="paged_adamw_8bit",
    report_to="tensorboard"
)

# 6. 创建Trainer实例
trainer = transformers.Trainer(
    model=model,
    train_dataset=tokenized_dataset["train"],
    args=training_args,
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
)

# 7. 开始训练
model.config.use_cache = False  # 训练时禁用缓存
trainer.train()

# 8. 保存模型
model.save_pretrained("./qlora-llama2-chatbot-adapter")

# 9. 合并模型进行推理
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map="auto"
)

merged_model = PeftModel.from_pretrained(base_model, "./qlora-llama2-chatbot-adapter")
merged_model = merged_model.merge_and_unload()
merged_model.save_pretrained("./merged-llama2-chatbot")
tokenizer.save_pretrained("./merged-llama2-chatbot")

# 10. 测试聊天机器人
def chat_with_model(prompt, model=merged_model, tokenizer=tokenizer, max_length=200):
    # 格式化输入
    formatted_prompt = f"<s>[INST] {prompt} [/INST]"
    
    # 编码并生成回复
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            temperature=0.7,
            do_sample=True
        )
    
    # 解码并返回结果
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # 提取模型回复部分
    response = response.split("[/INST]")[-1].strip()
    return response

# 测试几个问题
test_prompts = [
    "请解释什么是QLoRA技术？",
    "如何使用QLoRA微调一个大型语言模型？",
    "QLoRA相比LoRA有哪些优势？"
]

for prompt in test_prompts:
    response = chat_with_model(prompt)
    print(f"用户: {prompt}")
    print(f"模型: {response}\n")

# 医学领域模型微调示例
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset
import transformers

# 1. 加载和量化模型
model_name = "mistralai/Mistral-7B-v0.1"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained(model_name)

# 2. 准备模型和配置LoRA
model = prepare_model_for_kbit_training(model)
lora_config = LoraConfig(
    r=32,  # 医学领域需要更大的秩以捕捉专业知识
    lora_alpha=64,
    target_modules=["q_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.1,  # 医学应用需要更多正则化
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)

# 3. 加载医学数据集
# 这里假设已经有一个医学问答数据集
# 实际应用中，您需要使用真实的医学数据集
def load_medical_dataset():
    # 示例：创建一个简单的医学数据集
    medical_data = [
        {
            "question": "什么是高血压的主要风险因素？",
            "answer": "高血压的主要风险因素包括：年龄增长、家族史、高钠饮食、低钾摄入、肥胖、缺乏体力活动、吸烟、过量饮酒、长期压力以及某些慢性疾病如糖尿病和肾脏疾病。"
        },
        {
            "question": "心肌梗塞的典型症状有哪些？",
            "answer": "心肌梗塞的典型症状包括：胸部中央持续疼痛或不适，可能放射到手臂、颈部、下巴或背部；呼吸困难；出汗；恶心或呕吐；头晕或昏厥；焦虑感。值得注意的是，女性可能表现出不典型症状。"
        },
        # 更多医学问答数据...
    ]
    
    # 将数据转换为数据集格式
    from datasets import Dataset
    return Dataset.from_list(medical_data)

medical_dataset = load_medical_dataset()

# 4. 数据预处理
def preprocess_medical_data(examples):
    # 格式化医学问答数据
    texts = []
    for question, answer in zip(examples["question"], examples["answer"]):
        text = f"### 医学问题：\n{question}\n\n### 专业回答：\n{answer}"
        texts.append(text)
    
    # 编码文本
    return tokenizer(texts, padding="max_length", truncation=True, max_length=1024)

tokenized_medical_dataset = medical_dataset.map(preprocess_medical_data, batched=True)

# 5. 配置训练参数
training_args = transformers.TrainingArguments(
    output_dir="./qlora-medical-model",
    per_device_train_batch_size=2,  # 医学数据更复杂，使用更小的批量
    gradient_accumulation_steps=8,
    learning_rate=1e-4,  # 医学领域使用更小的学习率以稳定训练
    logging_steps=5,
    max_steps=500,
    save_strategy="steps",
    save_steps=50,
    warmup_steps=50,
    bf16=True,
    gradient_checkpointing=True,
    optim="paged_adamw_8bit"
)

# 6. 训练模型
trainer = transformers.Trainer(
    model=model,
    train_dataset=tokenized_medical_dataset,
    args=training_args,
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
)

model.config.use_cache = False
trainer.train()

# 7. 保存和使用模型
model.save_pretrained("./qlora-medical-adapter")

# 测试医学问答
def medical_qa(question):
    inputs = tokenizer(
        f"### 医学问题：\n{question}\n\n### 专业回答：\n",
        return_tensors="pt"
    ).to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=512,
            temperature=0.3,  # 医学应用使用较低温度以获得确定性答案
            do_sample=True
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("### 专业回答：\n")[-1]

# 测试医学问题
test_medical_questions = [
    "糖尿病患者的饮食注意事项有哪些？",
    "如何识别和应对药物过敏反应？",
    "心脏病患者的康复运动指南是什么？"
]

for question in test_medical_questions:
    answer = medical_qa(question)
    print(f"问题: {question}")
    print(f"专业回答: {answer}\n")

# 尝试更小的批量大小
training_args = transformers.TrainingArguments(
    per_device_train_batch_size=1,  # 尝试最小批量
    gradient_accumulation_steps=16,  # 增加累积步数以保持有效批量
    # 其他参数...
)

# 对所有层应用梯度检查点
model.gradient_checkpointing_enable()

training_args = transformers.TrainingArguments(
    optim="paged_adamw_8bit",  # 使用分页8位优化器
    # 其他参数...
)

# 减小序列长度
def preprocess_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)  # 减小max_length

# 尝试更小的学习率
training_args = transformers.TrainingArguments(
    learning_rate=1e-4,  # 降低学习率
    warmup_steps=200,  # 增加预热步数
    # 其他参数...
)

# 增加dropout和权重衰减
lora_config = LoraConfig(
    lora_dropout=0.1,  # 增加dropout
    # 其他参数...
)

training_args = transformers.TrainingArguments(
    weight_decay=0.01,  # 增加权重衰减
    # 其他参数...
)

training_args = transformers.TrainingArguments(
    gradient_clipping=1.0,  # 启用梯度裁剪
    # 其他参数...
)

# 使用更大的秩值
lora_config = LoraConfig(
    r=64,  # 增加秩值
    lora_alpha=128,  # 相应增加alpha
    # 其他参数...
)

# 对更多模块应用LoRA
lora_config = LoraConfig(
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    # 其他参数...
)

training_args = transformers.TrainingArguments(
    max_steps=2000,  # 增加训练步数
    # 其他参数...
)

# 推荐的版本组合
peft==0.10.0
transformers==4.40.0
bitsandbytes==0.43.1

# 针对特定模型家族的配置
if model_family == "mistral":
    lora_config = LoraConfig(
        target_modules=["q_proj", "v_proj", "gate_proj"],
        # 其他参数...
    )

# 手动指定模块名映射
if hasattr(model, "model") and hasattr(model.model, "layers"):
    # 处理嵌套模型结构
    for layer in model.model.layers:
        # 检查并适配不同的模块命名
        for name, module in layer.named_modules():
            if "attention" in name.lower():
                # 自定义处理...

# QLoRA与RAG结合示例
from langchain.chains import RetrievalQA
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

# 1. 加载经过QLoRA微调的模型
base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", torch_dtype=torch.float16)
adapter_model = PeftModel.from_pretrained(base_model, "qlora-finetuned-adapter")
merged_model = adapter_model.merge_and_unload()
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 2. 创建检索组件
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = FAISS.load_local("faiss_index", embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# 3. 创建自定义LLM包装器
from langchain.llms import HuggingFacePipeline
from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model=merged_model,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    device_map="auto",
    max_new_tokens=512,
    temperature=0.3
)

llm = HuggingFacePipeline(pipeline=pipe)

# 4. 创建RAG链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    return_source_documents=True
)

# 5. 执行增强检索问答
result = qa_chain("什么是QLoRA技术的最新进展？")
print(result["result"])

# QLoRA与知识蒸馏结合示例
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch

# 1. 加载QLoRA微调的教师模型
teacher_model_name = "meta-llama/Llama-2-7b-hf"
teacher_base = AutoModelForCausalLM.from_pretrained(
    teacher_model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)
teacher = PeftModel.from_pretrained(
    teacher_base,
    "qlora-finetuned-adapter"
)
teacher = teacher.merge_and_unload()

# 2. 加载小型学生模型
student_model = AutoModelForCausalLM.from_pretrained(
    "distil-whisper/distil-small.en",  # 示例小型模型
    torch_dtype=torch.float16,
    device_map="auto"
)

# 3. 配置蒸馏训练
from transformers import TrainingArguments, Trainer
from transformers import DataCollatorForLanguageModeling

# 这里需要定义自定义的蒸馏损失函数
# 实际实现需要根据具体任务定制

# 4. 执行蒸馏训练
# ...

# QLoRA与提示工程结合示例
def create_prompt_template(task_type):
    # 根据任务类型返回不同的提示模板
    templates = {
        "summarization": "请简洁地总结以下内容:\n{content}",
        "qa": "基于以下内容回答问题:\n内容: {context}\n问题: {question}",
        "translation": "将以下内容从{source_lang}翻译为{target_lang}:\n{content}"
    }
    return templates.get(task_type, "{content}")

# 在数据预处理中应用提示模板
def preprocess_with_template(examples, task_type="qa"):
    template = create_prompt_template(task_type)
    texts = []
    
    for example in examples:
        # 根据任务类型填充模板
        if task_type == "qa":
            text = template.format(
                context=example["context"],
                question=example["question"]
            ) + f"\n答案: {example['answer']}"
        # 其他任务类型的处理...
        
        texts.append(text)
    
    return tokenizer(texts, padding="max_length", truncation=True, max_length=1024)