**标题：Meta官方开源LLaMA Cookbook，我第一时间做了完整评测，涵盖微调、推理、量化全流程**

# 创建并激活新的conda环境
conda create -n llama-cookbook python=3.10
conda activate llama-cookbook

# 安装PyTorch（CUDA版本请根据你的显卡选择）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 安装transformers和其他核心依赖
pip install transformers accelerate bitsandbytes scipy peft trl

# 安装jupyter以便运行notebook示例
pip install jupyter ipykernel
python -m ipykernel install --user --name=llama-cookbook

# 克隆LLaMA Cookbook仓库
git clone https://github.com/meta-llama/llama-cookbook.git
cd llama-cookbook

# 安装cookbook本身的依赖
pip install -r requirements.txt

import torch
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU设备: {torch.cuda.get_device_name(0)}")
    print(f"GPU显存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")

# 登录Hugging Face并获取访问令牌
# 首先安装huggingface_hub并登录
pip install huggingface_hub
huggingface-cli login

# 或者在代码中直接设置令牌
from huggingface_hub import login
login(token="your_huggingface_token_here")

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

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

# 简单的文本生成
prompt = "人工智能将如何改变未来的教育方式？"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# 生成文本
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)
from datasets import load_dataset
import torch

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

# 设置pad token（重要的细节！）
tokenizer.pad_token = tokenizer.eos_token

# 准备数据集（这里使用alpaca格式的数据）
def format_instruction(sample):
    return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{sample['instruction']}

### Input:
{sample['input']}

### Response:
{sample['output']}"""

def tokenize_function(examples):
    formatted = [format_instruction(sample) for sample in zip(
        examples['instruction'],
        examples['input'],
        examples['output']
    )]
    result = tokenizer(
        formatted,
        truncation=True,
        max_length=512,
        padding="max_length"
    )
    result["labels"] = result["input_ids"].copy()
    return result

# 加载数据集
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
dataset = dataset.train_test_split(test_size=0.1)
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=dataset["train"].column_names
)

# 设置训练参数
training_args = TrainingArguments(
    output_dir="./llama2-finetuned",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-5,
    weight_decay=0.001,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    logging_steps=10,
    save_steps=500,
    fp16=True,
    optim="paged_adamw_8bit",
    report_to="tensorboard",
)

# 创建Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

# 开始训练
trainer.train()

from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset
import torch

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

# 配置LoRA参数
lora_config = LoraConfig(
    r=16,                          # LoRA矩阵的秩，值越大表达能力越强
    lora_alpha=32,                  # 缩放因子
    target_modules=[               # 应用LoRA的目标模块
        "q_proj", "v_proj", "k_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

# 将LoRA应用到模型
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# 输出将显示只有约0.1%的参数是可训练的

# 准备数据集和训练器（与全参数微调类似）
# ... 数据处理代码省略 ...

# 训练参数 - LoRA微调可以使用更大的batch size
training_args = TrainingArguments(
    output_dir="./llama2-lora",
    num_train_epochs=3,
    per_device_train_batch_size=8,    # 可以更大因为显存占用更少
    gradient_accumulation_steps=2,
    learning_rate=3e-4,               # LoRA通常使用更高的学习率
    warmup_ratio=0.03,
    logging_steps=10,
    save_steps=100,
    fp16=True,
    optim="paged_adamw_8bit",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

trainer.train()

# 保存LoRA权重
model.save_pretrained("./llama2-lora-checkpoint")

# 合并LoRA权重到基础模型（可选）
from peft import PeftModel
base_model = AutoModelForCausalLM.from_pretrained(model_name)
merged_model = PeftModel.from_pretrained(base_model, "./llama2-lora-checkpoint")
merged_model = merged_model.merge_and_unload()
merged_model.save_pretrained("./llama2-merged")

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# 配置4-bit量化
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",           # nf4比fp4效果更好
    bnb_4bit_use_double_quant=True        # 二次量化，进一步压缩
)

# 加载量化模型
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map="auto"
)

# 验证模型加载成功
print(f"模型参数量: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B")
print(f"量化后模型大小估算: ~4-5GB")

# 推理测试
prompt = "请介绍一下人工智能的发展历史："
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig

# GPTQ配置
quantization_config = GPTQConfig(
    bits=4,
    dataset="c4",                      # 使用什么数据集进行校准
    desc_act=True,                      # 激活顺序量化，效果更好但更慢
)

model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 加载并量化模型（首次运行会花费较长时间）
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map="auto"
)

# 保存量化后的模型
model.save_pretrained("./llama2-4bit-gptq")
tokenizer.save_pretrained("./llama2-4bit-gptq")

# 使用llama.cpp进行GGUF格式推理
# 首先需要将Hugging Face模型转换为GGUF格式

# 安装llama.cpp的Python绑定
pip install llama-cpp-python

# 或者使用完整的llama.cpp
# git clone https://github.com/ggerganov/llama.cpp.git
# cd llama.cpp && mkdir build && cd build && cmake .. && make

# 转换模型（使用Python脚本）
# python scripts/convert_hf_to_gguf.py meta-llama/Llama-2-7b-hf --outfile ./llama-2-7b.gguf

# 使用llama-cpp-python加载和推理
from llama_cpp import Llama

llm = Llama(
    model_path="./llama-2-7b.gguf",    # Q4_K_M量化版本约4GB
    n_ctx=2048,                        # 上下文窗口大小
    n_gpu_layers=0                     # 在CPU上运行，设为0
)

output = llm(
    "人工智能的未来发展趋势是什么？",
    max_tokens=200,
    temperature=0.7,
    stop=["</s>", "User:"]
)
print(output['choices'][0]['text'])

from langchain.llms import HuggingFacePipeline
from langchain.chains import RetrievalQA
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

# 创建本地的Hugging Face pipeline
model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    temperature=0.7,
    top_p=0.9,
    do_sample=True
)

# 包装成LangChain LLM
llm = HuggingFacePipeline(pipeline=pipe)

# 创建检索增强生成（RAG）系统
# 文档嵌入
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# 创建向量数据库（以Chroma为例）
vectorstore = Chroma.from_documents(
    documents=documents,          # 你的文档列表
    embedding=embeddings,
    persist_directory="./chroma_db"
)

# 创建QA链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
    return_source_documents=True
)

# 使用RAG系统回答问题
result = qa_chain({"query": "你关于XXX的问题是什么？"})
print(result["result"])

from langchain.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
import os

# 定义文档加载器
loader = DirectoryLoader(
    "./documents",                 # 你的文档目录
    glob="**/*.md",                # 加载所有markdown文件
    loader_cls=TextLoader
)

# 加载文档
documents = loader.load()
print(f"加载了 {len(documents)} 个文档")

# 文本分块
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,               # 每个块的目标大小
    chunk_overlap=50,              # 块之间的重叠，保证上下文连续性
    length_function=len,
    separators=["\n\n", "\n", "。", "！", "？", "，", " ", ""]
)

chunks = text_splitter.split_documents(documents)
print(f"分块后得到 {len(chunks)} 个文本块")

# 创建嵌入向量
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": "cpu"}
)

# 构建向量数据库
vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    persist_directory="./knowledge_base"    # 持久化存储位置
)

# 重要：持久化保存
vectorstore.persist()
print("知识库构建完成！")

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from langchain.llms import HuggingFacePipeline
import torch

# 加载Chat模型
model_name = "meta-llama/Llama-2-7b-chat-hf"

# 使用量化模型以节省显存
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_4bit=True
)

# 创建text-generation pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=1024,
    temperature=0.3,               # 较低温度保证答案准确性
    top_p=0.9,
    repetition_penalty=1.1,        # 避免重复生成
)

# 包装成LangChain LLM
llm = HuggingFacePipeline(pipeline=pipe)

from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings

# 加载之前构建的知识库
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = Chroma(
    persist_directory="./knowledge_base",
    embedding_function=embeddings
)

# 定义系统提示模板
system_prompt = """你是一个专业的知识库问答助手。请根据提供的上下文信息，准确回答用户的问题。

回答要求：
1. 只基于提供的上下文信息进行回答，不要编造信息
2. 如果上下文中没有相关信息，请明确告知用户
3. 回答要条理清晰，必要时可以列出要点
4. 在回答末尾注明参考的文档来源

上下文信息：
{context}

用户问题：{question}

请给出回答："""

prompt_template = PromptTemplate(
    template=system_prompt,
    input_variables=["context", "question"]
)

# 创建检索增强问答链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",           # 简单链类型，适合较短上下文
    retriever=vectorstore.as_retriever(
        search_kwargs={
            "k": 5                 # 检索最相关的5个文档块
        }
    ),
    return_source_documents=True,
    chain_type_kwargs={"prompt": prompt_template}
)

# 测试问答系统
question = "关于[你的问题]..."
result = qa_chain({"query": question})

print("=" * 50)
print("问题:", question)
print("=" * 50)
print("回答:")
print(result["result"])
print("=" * 50)
print("参考文档:")
for i, doc in enumerate(result["source_documents"]):
    print(f"{i+1}. {doc.page_content[:100]}...")

from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory

# 创建对话记忆
memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=True,
    output_key="answer"
)

# 创建带记忆的检索链
qa_chain_with_memory = ConversationalRetrievalChain.from_llm(
    llm=llm,
    retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
    memory=memory,
    return_source_documents=True,
    combine_docs_chain_kwargs={"prompt": prompt_template}
)

# 多轮对话示例
questions = [
    "第一个问题是什么？",
    "追问/补充问题是什么？",
    "继续深入的问题是什么？"
]

for q in questions:
    result = qa_chain_with_memory({"question": q})
    print(f"问: {q}")
    print(f"答: {result['answer']}\n")

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# 使用CodeLLaMa以获得更好的代码理解能力
model_name = "codellama/CodeLlama-7b-hf"

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
)

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

# 设置特殊的token以区分不同任务
CODE_COMPLETION_PREFIX = "<filepath>\n"
CODE_COMPLETION_SUFFIX = "\n</think>"
CODE_INFILL_PREFIX = " fills below: \n   <PRE>\n"
CODE_INFILL_SUFFIX = "\n   <SUF>\n   <MID>"

def complete_code(code_snippet, max_tokens=200):
    """
    补全给定的代码片段
    """
    prompt = f"""<EOT>
{code_snippet}"""

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=0.2,
            top_p=0.95,
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id
        )

    completion = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # 去除原始输入部分，只返回补全的内容
    return completion[len(prompt):]

def generate_code_from_docstring(docstring, language="python"):
    """
    根据文档字符串生成实现代码
    """
    prompt = f"""根据以下文档字符串实现代码：

{docstring}

实现代码：
"""

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=500,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
        )

    code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return code[len(prompt):]

def explain_code(code_snippet):
    """
    解释给定代码的功能
    """
    prompt = f"""请详细解释以下代码的功能和工作原理：

```{code_snippet}

inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=800,
        temperature=0.3,
        top_p=0.9,
    )

explanation = tokenizer.decode(outputs[0], skip_special_tokens=True)
return explanation.split("代码解释：")[-1]

completion = complete_code(snippet)
print("代码补全结果：")
print(snippet + completion)

# 测试代码解释
code = """

**第三步：添加语法检测和代码审查**

```python
def review_code(code_snippet):
    """
    对代码进行审查，提出改进建议
    """
    prompt = f"""你是一位经验丰富的代码审查员。请审查以下代码，指出潜在的问题并提供改进建议：

{code_snippet}

代码审查报告：
1. 代码质量评分（1-10）：
2. 优点：
3. 需要改进的地方：
4. 潜在bug或安全问题：
5. 改进后的代码示例："""

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=1000,
            temperature=0.3,
            top_p=0.9,
        )

    review = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return review.split("代码审查报告：")[-1]

def translate_code(code, from_lang="python", to_lang="javascript"):
    """
    将代码从一种语言翻译到另一种语言
    """
    prompt = f"""请将以下{from_lang}代码翻译成{to_lang}代码：

原始代码：
{code}

翻译后的{to_lang}代码："""

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=1000,
            temperature=0.3,
            top_p=0.9,
        )

    translated = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated.split(f"翻译后的{to_lang}代码：")[-1]

from datasets import load_dataset
import json

# 方法一：使用开源数据集
def prepare_chinese_conversation_dataset():
    """
    加载并格式化中文对话数据集
    """
    # 加载ShareGPT数据集（中文部分）
    dataset = load_dataset("JackJia/ShareGPT-Chinese-English-Alignment", split="train")

    def format_conversation(example):
        # 将ShareGPT格式转换为训练格式
        conversations = example["conversations"]
        formatted = ""

        for i, turn in enumerate(conversations):
            role = "用户" if turn["from"] == "human" else "助手"
            content = turn["value"]

            # 添加特殊格式标记
            formatted += f"<|im_start|>{role}\n{content}<|im_end|>\n"

        return {"text": formatted}

    formatted_dataset = dataset.map(format_conversation, remove_columns=dataset.column_names)
    return formatted_dataset

# 方法二：使用alpaca格式的数据集
def prepare_alpaca_format_dataset(data_path):
    """
    加载alpaca格式的数据集
    """
    # alpaca格式示例
    # {
    #     "instruction": "给出积极态度的三个理由",
    #     "input": "人工智能将改变教育行业",
    #     "output": "人工智能改变教育行业的三个积极理由：\n1..." 
    # }

    dataset = load_dataset("json", data_files=data_path, split="train")

    def format_alpaca(example):
        prompt = f"""<|im_start|>用户\n{example['instruction']}\n<|im_end|>\n<|im_start|>助手\n{example['output']}<|im_end|>"""
        return {"text": prompt}

    formatted_dataset = dataset.map(format_alpaca)
    return formatted_dataset

# 方法三：创建自己的对话数据集
def create_custom_dataset():
    """
    创建自定义的中文对话数据集
    """
    data = [
        {
            "instruction": "解释什么是机器学习",
            "input": "",
            "output": "机器学习是人工智能的一个分支，它使计算机能够从数据中学习并改进自身的性能，而无需明确编程。机器学习算法通过分析大量数据来识别模式，并基于这些模式做出预测或决策。常见的机器学习类型包括：监督学习、无监督学习和强化学习。"
        },
        {
            "instruction": "写一首关于春天的诗",
            "input": "",
            "output": "春回大地万物醒，\n柳绿桃红鸟争鸣。\n风和日暖花开处，\n一片生机入画屏。"
        },
        # 可以添加更多领域特定的数据
    ]

    # 保存为jsonl格式
    with open("./custom_dataset.jsonl", "w", encoding="utf-8") as f:
        for item in data:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    # 加载数据集
    dataset = load_dataset("json", data_files="./custom_dataset.jsonl", split="train")

    def format_item(example):
        prompt = f"""<|im_start|>用户\n{example['instruction']}\n<|im_end|>\n<|im_start|>助手\n{example['output']}<|im_end|>"""
        return {"text": prompt}

    return dataset.map(format_item)

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollator
)
from peft import LoraConfig, get_peft_model, TaskType
from datasets import load_dataset
import torch

# 定义ChatML格式的特殊token
CHATML_SYSTEM_PROMPT = """你是一个有帮助的AI助手。请用简洁、准确的语言回答用户的问题。"""

def apply_chat_template(example, tokenizer):
    """
    应用ChatML模板格式化对话
    """
    messages = []

    # 添加系统提示
    if example.get("system"):
        messages.append({"role": "system", "content": example["system"]})
    else:
        messages.append({"role": "system", "content": CHATML_SYSTEM_PROMPT})

    # 添加用户输入
    if example.get("instruction"):
        messages.append({"role": "user", "content": example["instruction"]})
        if example.get("input"):
            messages[-1]["content"] += f"\n\n{example['input']}"

    # 添加助手回复
    if example.get("output"):
        messages.append({"role": "assistant", "content": example["output"]})

    # 使用tokenizer的chat_template进行格式化
    formatted = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False
    )

    return {"text": formatted}

# 加载基础模型
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 设置chat template
tokenizer.chat_template = """{% for message in messages %}
<|im_start|>{{ message['role'] }}
{{ message['content'] }}<|im_end|>
{% endfor %}"""

# 加载数据集
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
dataset = dataset.train_test_split(test_size=0.1, seed=42)

# 格式化数据集
formatted_dataset = dataset["train"].map(
    lambda x: apply_chat_template(x, tokenizer),
    remove_columns=dataset["train"].column_names
)

# 分词
def tokenize_function(examples):
    result = tokenizer(
        examples["text"],
        truncation=True,
        max_length=2048,
        padding="max_length",
        return_tensors=None
    )
    result["labels"] = result["input_ids"].copy()
    return result

tokenized_dataset = formatted_dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"]
)

# 配置LoRA微调
lora_config = LoraConfig(
    r=32,
    lora_alpha=64,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_4bit=True
)

# 应用LoRA
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# 配置训练参数
training_args = TrainingArguments(
    output_dir="./llama2-chinese-chat",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=2,
    learning_rate=1e-4,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    logging_steps=10,
    save_steps=200,
    save_total_limit=3,
    bf16=True,
    tf32=True,
    optim="paged_adamw_8bit",
    group_by_length=True,
    report_to="tensorboard",
)

# 创建数据整理器
data_collator = DataCollator(
    tokenizer=tokenizer,
    padding=True,
    return_tensors="pt"
)

# 创建Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

# 开始训练
trainer.train()

# 保存模型
model.save_pretrained("./llama2-chinese-chat-final")

def chat_with_model(model_path, system_prompt="你是一个有帮助的AI助手。"):
    """
    与微调后的模型进行对话
    """
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    # 加载模型
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16,
        device_map="auto"
    )

    # 设置chat template
    tokenizer.chat_template = """<|im_start|>system
{{ system_prompt }}<|im_end|>
{% for message in messages %}
<|im_start|>{{ message['role'] }}
{{ message['content'] }}<|im_end|>
{% endfor %}"""

    messages = [{"role": "system", "content": system_prompt}]

    print("对话开始！输入'quit'退出。\n")

    while True:
        user_input = input("你: ")
        if user_input.lower() == "quit":
            break

        messages.append({"role": "user", "content": user_input})

        # 格式化对话
        prompt = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to("cuda")

        # 生成回复
        with torch.no_grad():
            outputs = model.generate(
                prompt,
                max_new_tokens=512,
                temperature=0.7,
                top_p=0.9,
                do_sample=True,
                repetition_penalty=1.1
            )

        # 解码回复
        response = tokenizer.decode(outputs[0][len(prompt[0]):], skip_special_tokens=True)

        print(f"助手: {response}\n")
        messages.append({"role": "assistant", "content": response})

# 启动对话
if __name__ == "__main__":
    chat_with_model("./llama2-chinese-chat-final")

# 构建企业知识库问答系统的核心代码框架
class EnterpriseKnowledgeBase:
    def __init__(self, documents_path, model_path):
        self.documents_path = documents_path
        self.model_path = model_path
        self.vectorstore = None
        self.qa_chain = None
        self._initialize()

    def _initialize(self):
        # 加载文档
        from langchain.document_loaders import DirectoryLoader, PyPDFLoader
        from langchain.text_splitter import RecursiveCharacterTextSplitter
        from langchain.embeddings import HuggingFaceEmbeddings
        from langchain.vectorstores import Chroma

        # 支持多种文档格式
        loaders = {
            ".pdf": PyPDFLoader,
            ".txt": DirectoryLoader,
            ".md": DirectoryLoader,
        }

        documents = []
        for ext, loader_cls in loaders.items():
            if loader_cls == DirectoryLoader:
                loader = DirectoryLoader(
                    self.documents_path,
                    glob=f"**/*{ext}",
                    loader_cls=TextLoader
                )
            else:
                loader = loader_cls()
            documents.extend(loader.load())

        # 文本分块
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )
        chunks = splitter.split_documents(documents)

        # 创建向量索引
        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/multilingual-MiniLM-L12-v2"
        )

        self.vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=embeddings,
            persist_directory="./enterprise_kb"
        )

    def query(self, question, top_k=5):
        # 检索相关文档
        docs = self.vectorstore.similarity_search(question, k=top_k)

        # 构建上下文
        context = "\n\n".join([doc.page_content for doc in docs])

        # 生成答案
        # ...（与之前的RAG实现类似）

        return answer, source_documents

class EducationAssistant:
    def __init__(self, model_path):
        self.model, self.tokenizer = self._load_model(model_path)

    def generate_quiz(self, topic, num_questions=5, question_type="multiple_choice"):
        """
        根据主题生成测验题目
        """
        prompt = f"""为以下主题生成{num_questions}道{question_type}题目：

主题：{topic}

要求：
1. 题目应覆盖主题的核心概念
2. 难度适中，适合学习者检验理解程度
3. 提供答案和简要解析

输出格式：
【题目1】
A. ...
B. ...
C. ...
D. ...
答案：X
解析：..."""

        return self._generate_text(prompt)

    def explain_concept(self, concept, level="intermediate"):
        """
        用指定难度级别解释概念
        """
        level_prompts = {
            "beginner": "用最简单易懂的语言，适合完全没有基础的读者",
            "intermediate": "用较为专业的语言，需要一定的背景知识",
            "advanced": "用高度专业的学术语言，假设读者已具备深厚背景"
        }

        prompt = f"""{level_prompts[level]}

请解释以下概念：
{concept}

解释应包括：
1. 基本定义
2. 核心原理
3. 实际应用例子
4. 相关概念的联系"""

        return self._generate_text(prompt)

    def grade_assignment(self, assignment_text, rubric):
        """
        根据评分标准批改作业
        """
        prompt = f"""请根据以下评分标准批改这份作业：

评分标准：
{rubric}

学生作业：
{assignment_text}

请提供：
1. 总体评分
2. 各维度得分
3. 具体改进建议"""

        return self._generate_text(prompt)

class ContentCreationAssistant:
    def __init__(self, model_path):
        self.model, self.tokenizer = self._load_model(model_path)

    def generate_blog_post(self, topic, target_length="medium", style="informative"):
        """
        生成博客文章
        """
        length_guide = {
            "short": "约800字",
            "medium": "约1500字",
            "long": "约2500字"
        }

        style_guide = {
            "informative": "以信息传递为主，客观专业",
            "casual": "轻松活泼，像朋友聊天",
            "professional": "严谨正式，适合商务场景"
        }

        prompt = f"""请撰写一篇{length_guide[target_length]}的博客文章。

主题：{topic}
风格：{style_guide[style]}

文章结构要求：
1. 吸引人的开头
2. 核心内容（分3-5个要点）
3. 总结和行动建议"""

        return self._generate_text(prompt)

    def improve_writing(self, text, focus="general"):
        """
        润色和改进文章
        """
        focus_guide = {
            "clarity": "提高清晰度，消除歧义",
            "engagement": "增加趣味性和吸引力",
            "grammar": "修正语法和拼写错误",
            "general": "全面改进"
        }

        prompt = f"""请对以下文章进行润色，专注于{focus_guide[focus]}：

原文：
{text}

改进后的版本："""

        return self._generate_text(prompt)

    def generate_seo_titles(self, article_content, num_titles=5):
        """
        生成SEO友好的标题
        """
        prompt = f"""根据以下文章内容，生成{num_titles}个SEO友好的标题：

{article_content[:500]}

要求：
1. 标题应包含关键词
2. 吸引点击，有吸引力
3. 长度适中（25-60字符）
4. 使用数字或疑问句增加吸引力"""

        return self._generate_text(prompt)

class BatchedInference:
    def __init__(self, model_path, batch_size=8):
        self.model_path = model_path
        self.batch_size = batch_size
        self.model, self.tokenizer = self._load_model()

    def batch_generate(self, prompts, max_length=200):
        """
        批量生成文本
        """
        # 批量编码
        inputs = self.tokenizer(
            prompts,
            return_tensors="pt",
            padding=True,
            truncation=True
        ).to("cuda")

        # 批量生成
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_length,
                do_sample=True,
                temperature=0.7
            )

        # 批量解码
        results = self.tokenizer.batch_decode(
            outputs,
            skip_special_tokens=True
        )

        return results

# 使用示例
inference_engine = BatchedInference("./llama2-7b", batch_size=16)
prompts = [f"第{i}个提示词" for i in range(100)]
results = inference_engine.batch_generate(prompts)

# 在transformers中启用KV cache（默认已启用）
# 但可以通过调整配置来优化

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig

model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 使用更激进的KV cache设置
generation_config = GenerationConfig(
    max_new_tokens=200,
    use_cache=True,              # 启用KV cache
    cache_implementation="static"  # 静态cache，可能更快
)

# 生成时使用配置
inputs = tokenizer("提示词", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, generation_config=generation_config)

# 安装flash-attn（需要编译）
# pip install flash-attn --no-build-isolation

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 使用Flash Attention加载模型
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    attn_implementation="flash_attention_2"    # 启用Flash Attention 2
)

# 技巧一：使用梯度检查点（Gradient Checkpointing）
# 在训练时用计算换显存

model.gradient_checkpointing_enable()

# 技巧二：优化器分页管理
# 将优化器状态移到CPU内存

training_args = TrainingArguments(
    optim="paged_adamw_32bit",    # 使用分页优化器
    optim_args={"page_size": 16}  # 分页大小
)

# 技巧三：CPU卸载
# 不活跃的层卸载到CPU

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    max_memory={    # 配置每层存放位置
        0: "10GB",       # GPU 0: 10GB
        "cpu": "100GB"  # CPU: 100GB
    }
)

# 技巧四：使用更激进的量化
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True
)

# 技巧一：使用结构化输出
def structured_output_prompt(task, schema):
    """
    生成结构化输出的提示词
    """
    return f"""{task}

请按以下JSON格式输出：
{schema}

重要：
1. 只输出JSON，不要包含其他文字
2. 确保JSON格式正确
3. 所有字段都必须有值"""

# 技巧二：Few-shot学习
def few_shot_prompt(examples, query):
    """
    Few-shot提示示例
    """
    prompt = ""
    for ex in examples:
        prompt += f"输入：{ex['input']}\n输出：{ex['output']}\n\n"
    prompt += f"输入：{query}\n输出："
    return prompt

# 技巧三：思维链提示
def cot_prompt(question):
    """
    思维链提示，引导模型逐步思考
    """
    return f"""请逐步推理以下问题：

问题：{question}

推理过程：
1. """

# 技巧四：系统提示优化
SYSTEM_PROMPT_TEMPLATE = """你是一个{role}。

背景信息：
{background}

回答要求：
{requirements}

约束条件：
{constraints}"""

def create_system_prompt(role, background, requirements, constraints):
    """
    创建结构化的系统提示
    """
    return SYSTEM_PROMPT_TEMPLATE.format(
        role=role,
        background=background,
        requirements=requirements,
        constraints=constraints
    )

class SafetyChecker:
    def __init__(self):
        self.banned_topics = [...]  # 定义禁止讨论的话题
        self.safety_keywords = [...]  # 安全关键词

    def check_input(self, text):
        """
        检查用户输入是否安全
        """
        for keyword in self.safety_keywords:
            if keyword in text.lower():
                return False, f"输入包含敏感内容：{keyword}"
        return True, "安全"

    def check_output(self, text):
        """
        检查模型输出是否安全
        """
        # 检查有害内容
        if self._contains_harmful_content(text):
            return False, "输出包含有害内容"

        # 检查事实准确性
        if self._has_confidence_issues(text):
            return False, "输出可能包含不准确信息"

        return True, "安全"

    def _contains_harmful_content(self, text):
        """
        检测有害内容的辅助方法
        """
        harmful_patterns = [
            "如何制造", "怎么攻击", "非法获取",
            # 添加更多有害模式
        ]
        return any(pattern in text for pattern in harmful_patterns)

    def _has_confidence_issues(self, text):
        """
        检测模型过度自信的问题
        """
        # 如果模型不确定，应该表达不确定性
        confidence_indicators = [
            "我不确定", "可能是", "我不清楚",
            "据我所知", "一般来说"
        ]
        # 如果是很长且确定的回复但没有不确定性表达，可能有问题
        return False  # 简化实现

# 安全检查使用示例
safety_checker = SafetyChecker()

def safe_generate(prompt, model, tokenizer):
    # 检查输入
    is_safe, msg = safety_checker.check_input(prompt)
    if not is_safe:
        return f"抱歉，无法处理此请求：{msg}"

    # 生成回复
    response = generate_text(prompt, model, tokenizer)

    # 检查输出
    is_safe, msg = safety_checker.check_output(response)
    if not is_safe:
        return f"抱歉，生成的回复存在问题：{msg}"

    return response