"""
大模型的核心特征：

1. 规模巨大：参数数量从数十亿到数万亿
2. 预训练+微调：先在海量数据上预训练，再针对特定任务微调
3. 涌现能力：模型达到一定规模后出现的新能力
4. 上下文学习：通过提示词而非梯度更新来适应新任务
"""

# 理解模型参数
class SimpleTransformer:
    """简化的 Transformer 模型结构示意"""

    def __init__(self, vocab_size=50000, hidden_size=768, num_layers=12, num_heads=12):
        self.vocab_size = vocab_size      # 词汇表大小
        self.hidden_size = hidden_size    # 隐藏层维度
        self.num_layers = num_layers      # 层数
        self.num_heads = num_heads        # 注意力头数

        # 计算参数量
        self.param_count = self._calculate_params()

    def _calculate_params(self):
        """计算模型参数量"""
        # 嵌入层参数
        embedding_params = self.vocab_size * self.hidden_size

        # 每层 Transformer 参数
        # 多头注意力：4 个投影矩阵 (Q, K, V, O)
        attention_params = 4 * self.hidden_size * self.hidden_size

        # 前馈网络：通常 hidden_size * 4 * hidden_size * 2
        ffn_params = 2 * self.hidden_size * self.hidden_size * 4

        layer_params = attention_params + ffn_params

        # 总参数
        total = embedding_params + self.num_layers * layer_params

        return total

# 示例：GPT-2 规模
gpt2_small = SimpleTransformer(
    vocab_size=50257,
    hidden_size=768,
    num_layers=12,
    num_heads=12
)
print(f"GPT-2 Small 参数量: {gpt2_small.param_count / 1e6:.1f}M")

# 示例：GPT-3 规模
gpt3 = SimpleTransformer(
    vocab_size=50257,
    hidden_size=12288,
    num_layers=96,
    num_heads=96
)
print(f"GPT-3 参数量: {gpt3.param_count / 1e9:.1f}B")

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
    """多头注意力机制"""

    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # Q, K, V 的线性变换
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # 线性变换并分割成多头
        Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        # 应用掩码（用于因果注意力或填充）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # Softmax 获得注意力权重
        attention_weights = F.softmax(scores, dim=-1)

        # 加权求和
        attention_output = torch.matmul(attention_weights, V)

        # 合并多头
        attention_output = attention_output.transpose(1, 2).contiguous().view(
            batch_size, -1, self.d_model
        )

        # 最终线性变换
        output = self.W_o(attention_output)

        return output, attention_weights

class FeedForward(nn.Module):
    """前馈神经网络"""

    def __init__(self, d_model=512, d_ff=2048, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.linear2(self.dropout(F.gelu(self.linear1(x))))

class TransformerBlock(nn.Module):
    """Transformer 块"""

    def __init__(self, d_model=512, num_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()

        self.attention = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = FeedForward(d_model, d_ff, dropout)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 自注意力 + 残差连接 + 层归一化
        attn_output, attn_weights = self.attention(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))

        # 前馈网络 + 残差连接 + 层归一化
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))

        return x, attn_weights

class GPTDecoder(nn.Module):
    """简化的 GPT 解码器"""

    def __init__(self, vocab_size=50257, d_model=768, num_layers=12, 
                 num_heads=12, d_ff=3072, max_seq_len=1024, dropout=0.1):
        super().__init__()

        # 词嵌入和位置嵌入
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(max_seq_len, d_model)

        # Transformer 层
        self.layers = nn.ModuleList([
            TransformerBlock(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])

        # 输出层
        self.ln_final = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

        self.dropout = nn.Dropout(dropout)
        self.max_seq_len = max_seq_len

    def forward(self, input_ids, attention_mask=None):
        batch_size, seq_len = input_ids.shape

        # 获取位置索引
        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)

        # 词嵌入 + 位置嵌入
        token_embeds = self.token_embedding(input_ids)
        position_embeds = self.position_embedding(positions)
        x = self.dropout(token_embeds + position_embeds)

        # 因果注意力掩码（防止看到未来信息）
        causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device))
        causal_mask = causal_mask.view(1, 1, seq_len, seq_len)

        # 通过所有 Transformer 层
        for layer in self.layers:
            x, _ = layer(x, causal_mask)

        # 输出层
        x = self.ln_final(x)
        logits = self.lm_head(x)

        return logits

    def generate(self, input_ids, max_new_tokens=50, temperature=1.0, top_k=50):
        """自回归生成文本"""
        for _ in range(max_new_tokens):
            # 截取最后 max_seq_len 个 token
            input_ids_cond = input_ids[:, -self.max_seq_len:]

            # 获取 logits
            with torch.no_grad():
                logits = self.forward(input_ids_cond)

            # 只取最后一个位置的 logits
            logits = logits[:, -1, :] / temperature

            # Top-K 采样
            if top_k > 0:
                indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
                logits[indices_to_remove] = -float('Inf')

            # 转换为概率
            probs = F.softmax(logits, dim=-1)

            # 采样
            next_token = torch.multinomial(probs, num_samples=1)

            # 拼接
            input_ids = torch.cat([input_ids, next_token], dim=-1)

        return input_ids

"""
Hugging Face Transformers 是大模型开发的标准库
安装：pip install transformers torch accelerate
"""

from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM, 
    pipeline,
    GenerationConfig
)
import torch

# ============== 1. 加载模型和分词器 ==============

# 方法1：加载本地或 Hugging Face Hub 上的模型
model_name = "gpt2"  # 可以替换为 "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"  # 自动分配设备
)

# 设置填充 token（GPT-2 没有 pad_token）
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print(f"模型参数量: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B")

# ============== 2. 基础文本生成 ==============

def generate_text(prompt, max_length=100, temperature=0.7):
    """基础文本生成函数"""

    # 编码输入
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # 生成配置
    generation_config = GenerationConfig(
        max_new_tokens=max_length,
        temperature=temperature,
        do_sample=True,
        top_p=0.95,
        top_k=50,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

    # 生成
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            generation_config=generation_config
        )

    # 解码输出
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return generated_text

# 测试生成
prompt = "人工智能的未来是"
result = generate_text(prompt)
print(f"Prompt: {prompt}")
print(f"Generated: {result}")

# ============== 3. 使用 Pipeline API ==============

# Pipeline 是最简单的使用方式
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto"
)

# 生成文本
outputs = generator(
    "今天天气真好，",
    max_length=50,
    num_return_sequences=3,
    temperature=0.8,
    do_sample=True
)

print("\n多序列生成:")
for i, output in enumerate(outputs):
    print(f"序列 {i+1}: {output['generated_text']}")

# ============== 4. 批量推理 ==============

def batch_generate(prompts, **kwargs):
    """批量生成文本"""

    # 批量编码
    inputs = tokenizer(
        prompts, 
        return_tensors="pt", 
        padding=True,
        truncation=True
    ).to(model.device)

    # 批量生成
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=kwargs.get('max_length', 50),
            temperature=kwargs.get('temperature', 0.7),
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id,
        )

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

    return generated_texts

prompts = [
    "Python 是一种",
    "机器学习是",
    "深度学习与神经网络的区别在于"
]

results = batch_generate(prompts, max_length=30)
for prompt, result in zip(prompts, results):
    print(f"\nPrompt: {prompt}")
    print(f"Result: {result}")

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

class ChatModel:
    """对话模型封装"""

    def __init__(self, model_name="Qwen/Qwen2-7B-Instruct"):
        """
        初始化对话模型
        支持 Qwen、Llama、ChatGLM 等对话模型
        """
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            trust_remote_code=True
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16,
            device_map="auto",
            trust_remote_code=True
        )

        self.conversation_history = []

    def _format_messages(self, messages):
        """格式化消息为模型输入格式"""
        # 不同模型的对话格式不同，这里以通用格式为例
        formatted = ""
        for msg in messages:
            if msg["role"] == "system":
                formatted += f"<|system|>\n{msg['content']}\n"
            elif msg["role"] == "user":
                formatted += f"<|user|>\n{msg['content']}\n"
            elif msg["role"] == "assistant":
                formatted += f"<|assistant|>\n{msg['content']}\n"

        formatted += "<|assistant|>\n"
        return formatted

    def chat(self, user_message, system_prompt=None, max_length=512, temperature=0.7):
        """单轮对话"""

        messages = []

        # 添加系统提示
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})

        # 添加历史对话
        messages.extend(self.conversation_history)

        # 添加当前用户消息
        messages.append({"role": "user", "content": user_message})

        # 格式化输入
        prompt = self._format_messages(messages)

        # 编码
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)

        # 生成回复
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_length,
                temperature=temperature,
                do_sample=True,
                top_p=0.9,
                repetition_penalty=1.1,
                pad_token_id=self.tokenizer.eos_token_id,
            )

        # 解码回复（只取新生成的部分）
        response = self.tokenizer.decode(
            outputs[0][inputs['input_ids'].shape[1]:],
            skip_special_tokens=True
        )

        # 更新对话历史
        self.conversation_history.append({"role": "user", "content": user_message})
        self.conversation_history.append({"role": "assistant", "content": response})

        # 保持历史长度合理
        if len(self.conversation_history) > 10:
            self.conversation_history = self.conversation_history[-10:]

        return response

    def clear_history(self):
        """清空对话历史"""
        self.conversation_history = []

    def stream_chat(self, user_message, system_prompt=None):
        """流式对话（逐词生成）"""
        from transformers import TextStreamer

        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.extend(self.conversation_history)
        messages.append({"role": "user", "content": user_message})

        prompt = self._format_messages(messages)
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)

        # 创建流式输出器
        streamer = TextStreamer(
            self.tokenizer,
            skip_prompt=True,
            skip_special_tokens=True
        )

        # 生成（流式输出）
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=512,
                temperature=0.7,
                do_sample=True,
                streamer=streamer,
                pad_token_id=self.tokenizer.eos_token_id,
            )

        # 获取完整回复
        response = self.tokenizer.decode(
            outputs[0][inputs['input_ids'].shape[1]:],
            skip_special_tokens=True
        )

        # 更新历史
        self.conversation_history.append({"role": "user", "content": user_message})
        self.conversation_history.append({"role": "assistant", "content": response})

        return response

# 使用示例
# chat_model = ChatModel("Qwen/Qwen2-7B-Instruct")
# response = chat_model.chat("什么是机器学习？", system_prompt="你是一个专业的AI助手")
# print(response)

class PromptEngineering:
    """提示工程技巧集合"""

    @staticmethod
    def zero_shot(prompt):
        """零样本提示：不给示例，直接提问"""
        return prompt

    @staticmethod
    def few_shot(prompt, examples):
        """少样本提示：提供几个示例"""
        formatted = ""
        for example in examples:
            formatted += f"问题: {example['question']}\n"
            formatted += f"回答: {example['answer']}\n\n"
        formatted += f"问题: {prompt}\n回答:"
        return formatted

    @staticmethod
    def chain_of_thought(prompt):
        """思维链提示：要求模型展示推理过程"""
        return f"""{prompt}

请一步一步地思考，展示你的推理过程，然后给出最终答案。

推理过程：
步骤1:"""

    @staticmethod
    def role_playing(prompt, role, expertise, tone):
        """角色扮演提示"""
        return f"""你是一位{role}，专长于{expertise}。
请用{tone}的语气回答问题。

问题：{prompt}

回答："""

    @staticmethod
    def structured_output(prompt, format_template):
        """结构化输出提示"""
        return f"""{prompt}

请严格按照以下JSON格式输出：
{format_template}

输出："""

# 使用示例
pe = PromptEngineering()

# 1. 零样本
print("=== 零样本提示 ===")
zero_shot_prompt = pe.zero_shot("解释什么是机器学习")
print(zero_shot_prompt)

# 2. 少样本
print("\n=== 少样本提示 ===")
examples = [
    {"question": "苹果是什么颜色的？", "answer": "苹果通常是红色的，但也有绿色和黄色的品种。"},
    {"question": "香蕉是什么颜色的？", "answer": "香蕉通常是黄色的，未成熟时是绿色的。"}
]
few_shot_prompt = pe.few_shot("橙子是什么颜色的？", examples)
print(few_shot_prompt)

# 3. 思维链
print("\n=== 思维链提示 ===")
cot_prompt = pe.chain_of_thought("一个农场有15只鸡和8只兔子，总共有多少条腿？")
print(cot_prompt)

# 4. 角色扮演
print("\n=== 角色扮演提示 ===")
role_prompt = pe.role_playing(
    "如何学习编程？",
    role="资深软件工程师",
    expertise="编程教育和系统架构",
    tone="专业但友善"
)
print(role_prompt)

# 5. 结构化输出
print("\n=== 结构化输出提示 ===")
format_template = """
{
    "summary": "一句话总结",
    "key_points": ["要点1", "要点2", "要点3"],
    "sentiment": "positive/neutral/negative",
    "confidence": 0.0-1.0
}
"""
structured_prompt = pe.structured_output("分析这段话：今天天气很好，我和朋友去了公园。", format_template)
print(structured_prompt)

class AdvancedPrompting:
    """高级提示技术"""

    @staticmethod
    def react_prompt(question, tools_description):
        """ReAct 提示：推理+行动"""
        return f"""你是一个能够使用工具的AI助手。你可以通过思考、行动、观察的循环来解决问题。

可用工具：
{tools_description}

请使用以下格式回答：
Thought: 我需要思考什么
Action: 要使用的工具和输入
Observation: 工具返回的结果
... (可以重复多次)
Thought: 我现在知道答案了
Final Answer: 最终答案

问题：{question}"""

    @staticmethod
    def self_consistency_prompt(prompt, num_paths=3):
        """自洽性提示：生成多个推理路径，投票决定答案"""
        return f"""请用{num_paths}种不同的方式思考和解决以下问题，然后综合给出最可靠的答案。

问题：{prompt}

请按以下格式输出：
推理路径1:
[第一种推理过程]
答案1: [答案]

推理路径2:
[第二种推理过程]
答案2: [答案]

推理路径3:
[第三种推理过程]
答案3: [答案]

综合答案: [基于以上推理的最可靠答案]"""

    @staticmethod
    def tree_of_thoughts_prompt(problem, num_branches=3):
        """思维树提示：探索多个思维分支"""
        return f"""请用思维树的方式解决以下问题。对于每个步骤，生成{num_branches}个可能的方向，评估每个方向，选择最有希望的继续。

问题：{problem}

格式：
步骤1: [当前状态]
- 分支A: [想法] (评分: X/10)
- 分支B: [想法] (评分: X/10)
- 分支C: [想法] (评分: X/10)
选择: [最佳分支]

步骤2: [基于选择继续]
..."""

    @staticmethod
    def automatic_prompt_engineering(task_description, input_examples):
        """自动提示工程：让模型自己生成优化后的提示"""
        return f"""你是一个提示工程专家。请为以下任务设计一个最优的提示模板。

任务描述：{task_description}

输入示例：
{input_examples}

请设计一个提示模板，包含：
1. 系统角色定义
2. 任务说明
3. 输出格式要求
4. 示例（如果需要）
5. 约束条件

优化后的提示模板："""

# 使用示例
ap = AdvancedPrompting()

# ReAct 提示
tools_desc = """
1. search(query: str) -> str: 搜索互联网获取信息
2. calculate(expression: str) -> float: 计算数学表达式
3. get_weather(city: str) -> str: 获取天气信息
"""
react = ap.react_prompt("今天北京适合户外运动吗？", tools_desc)
print("=== ReAct 提示 ===\n", react)

"""
LoRA (Low-Rank Adaptation) 是一种高效的微调方法
只训练少量参数，大幅降低显存需求

安装依赖：
pip install peft transformers datasets accelerate bitsandbytes
"""

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)
from peft import (
    LoraConfig,
    get_peft_model,
    TaskType,
    prepare_model_for_kbit_training
)
from datasets import Dataset
import torch
import json

class LoRATrainer:
    """LoRA 微调器"""

    def __init__(
        self,
        base_model_name="Qwen/Qwen2-7B",
        lora_r=8,
        lora_alpha=16,
        lora_dropout=0.05,
        use_4bit=True
    ):
        self.base_model_name = base_model_name

        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(
            base_model_name,
            trust_remote_code=True
        )

        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        # 加载模型
        load_kwargs = {
            "trust_remote_code": True,
            "device_map": "auto",
        }

        if use_4bit:
            # 4-bit 量化加载
            from transformers import BitsAndBytesConfig
            load_kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4"
            )
        else:
            load_kwargs["torch_dtype"] = torch.float16

        self.model = AutoModelForCausalLM.from_pretrained(
            base_model_name,
            **load_kwargs
        )

        # 准备模型进行 k-bit 训练
        if use_4bit:
            self.model = prepare_model_for_kbit_training(self.model)

        # 配置 LoRA
        self.lora_config = LoraConfig(
            r=lora_r,
            lora_alpha=lora_alpha,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # 根据模型调整
            lora_dropout=lora_dropout,
            bias="none",
            task_type=TaskType.CAUSAL_LM,
        )

        self.model = get_peft_model(self.model, self.lora_config)
        self.model.print_trainable_parameters()

    def format_instruction(self, example):
        """格式化指令数据"""
        # Alpaca 格式
        if "instruction" in example and "output" in example:
            text = f"""### Instruction:
{example['instruction']}

### Input:
{example.get('input', '')}

### Response:
{example['output']}"""
        # 对话格式
        elif "conversations" in example:
            text = ""
            for turn in example["conversations"]:
                if turn["from"] == "human":
                    text += f"### Human:\n{turn['value']}\n\n"
                else:
                    text += f"### Assistant:\n{turn['value']}\n\n"
        else:
            text = str(example)

        return text

    def tokenize_function(self, examples):
        """分词函数"""
        texts = [self.format_instruction(ex) for ex in examples]

        tokenized = self.tokenizer(
            texts,
            truncation=True,
            padding="max_length",
            max_length=512,
            return_tensors="pt"
        )

        # 标签与输入相同（用于语言模型损失）
        tokenized["labels"] = tokenized["input_ids"].clone()

        return tokenized

    def train(
        self,
        train_data,
        output_dir="./lora_output",
        num_epochs=3,
        batch_size=4,
        learning_rate=2e-4,
        gradient_accumulation_steps=4,
        save_steps=100,
        logging_steps=10,
    ):
        """执行训练"""

        # 准备数据集
        if isinstance(train_data, str):
            # 从文件加载
            with open(train_data, 'r', encoding='utf-8') as f:
                data = [json.loads(line) for line in f]
        else:
            data = train_data

        dataset = Dataset.from_list(data)
        tokenized_dataset = dataset.map(
            self.tokenize_function,
            batched=True,
            remove_columns=dataset.column_names
        )

        # 数据整理器
        data_collator = DataCollatorForLanguageModeling(
            tokenizer=self.tokenizer,
            mlm=False,
        )

        # 训练参数
        training_args = TrainingArguments(
            output_dir=output_dir,
            num_train_epochs=num_epochs,
            per_device_train_batch_size=batch_size,
            gradient_accumulation_steps=gradient_accumulation_steps,
            learning_rate=learning_rate,
            logging_steps=logging_steps,
            save_steps=save_steps,
            save_total_limit=2,
            remove_unused_columns=False,
            report_to="none",
            fp16=True,
            optim="paged_adamw_8bit",
            lr_scheduler_type="cosine",
            warmup_ratio=0.03,
        )

        # 训练器
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=tokenized_dataset,
            data_collator=data_collator,
        )

        # 开始训练
        trainer.train()

        # 保存模型
        trainer.save_model(output_dir)
        self.tokenizer.save_pretrained(output_dir)

        print(f"训练完成，模型已保存到 {output_dir}")

    def merge_and_save(self, output_dir="./merged_model"):
        """合并 LoRA 权重并保存完整模型"""
        # 合并权重
        merged_model = self.model.merge_and_unload()

        # 保存
        merged_model.save_pretrained(output_dir)
        self.tokenizer.save_pretrained(output_dir)

        print(f"合并后的模型已保存到 {output_dir}")

        return merged_model

# 使用示例
def prepare_training_data():
    """准备训练数据示例"""
    data = [
        {
            "instruction": "解释什么是机器学习",
            "input": "",
            "output": "机器学习是人工智能的一个分支，它使计算机系统能够从数据中学习并改进，而无需显式编程。机器学习算法通过分析大量数据来识别模式，并基于这些模式做出预测或决策。"
        },
        {
            "instruction": "用Python写一个计算斐波那契数列的函数",
            "input": "n=10",
            "output": "def fibonacci(n):\n    if n <= 1:\n        return n\n    a, b = 0, 1\n    for _ in range(2, n+1):\n        a, b = b, a + b\n    return b\n\n# 计算第10个斐波那契数\nresult = fibonacci(10)\nprint(result)  # 输出: 55"
        },
        # 添加更多数据...
    ]
    return data

# 训练代码（需要足够的GPU显存）
"""
trainer = LoRATrainer(
    base_model_name="Qwen/Qwen2-7B",
    lora_r=8,
    use_4bit=True
)

train_data = prepare_training_data()
trainer.train(
    train_data,
    output_dir="./my_lora_model",
    num_epochs=3,
    batch_size=2
)

# 合并并保存
trainer.merge_and_save("./my_finetuned_model")
"""

"""
Axolotl 是一个简化的 LLM 微调框架
配置文件示例：config.yml
"""

config_example = """
# Axolotl 配置示例
base_model: Qwen/Qwen2-7B
model_type: AutoModelForCausalLM
tokenizer_type: AutoTokenizer

# LoRA 配置
adapter: lora
lora_r: 8
lora_alpha: 16
lora_dropout: 0.05
lora_target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj

# 数据集配置
datasets:
  - path: ./data/train.jsonl
    type: alpaca
    conversation: alpaca

# 训练配置
sequence_len: 2048
sample_packing: true
pad_to_sequence_len: true

# 优化器配置
optimizer: adamw_bnb_8bit
lr_scheduler: cosine
learning_rate: 2e-4
num_epochs: 3
batch_size: 4
gradient_accumulation_steps: 4

# 保存配置
output_dir: ./outputs
save_steps: 100
save_total_limit: 2

# 其他
bf16: auto
tf32: true
gradient_checkpointing: true
"""

# 数据格式示例
data_format_example = """
# Alpaca 格式 (train.jsonl)
{"instruction": "解释什么是人工智能", "input": "", "output": "人工智能是..."}
{"instruction": "写一个排序算法", "input": "使用Python", "output": "def bubble_sort..."}

# ShareGPT 格式
{
  "conversations": [
    {"from": "human", "value": "你好"},
    {"from": "gpt", "value": "你好！有什么可以帮助你的？"},
    {"from": "human", "value": "解释一下什么是机器学习"},
    {"from": "gpt", "value": "机器学习是..."}
  ]
}
"""

"""
RAG (Retrieval-Augmented Generation) 结合检索和生成
让模型能够访问外部知识库
"""

import numpy as np
from typing import List, Tuple, Optional
from dataclasses import dataclass
import hashlib

@dataclass
class Document:
    """文档结构"""
    content: str
    metadata: dict
    embedding: Optional[np.ndarray] = None

class SimpleVectorStore:
    """简单的向量存储"""

    def __init__(self, embedding_dim=768):
        self.documents: List[Document] = []
        self.embeddings: List[np.ndarray] = []
        self.embedding_dim = embedding_dim

    def add_document(self, doc: Document):
        """添加文档"""
        self.documents.append(doc)
        if doc.embedding is not None:
            self.embeddings.append(doc.embedding)

    def add_documents(self, docs: List[Document]):
        """批量添加文档"""
        for doc in docs:
            self.add_document(doc)

    def search(self, query_embedding: np.ndarray, top_k: int = 5) -> List[Tuple[Document, float]]:
        """搜索最相似的文档"""
        if not self.embeddings:
            return []

        # 计算余弦相似度
        embeddings = np.array(self.embeddings)
        similarities = np.dot(embeddings, query_embedding) / (
            np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query_embedding)
        )

        # 获取 top-k
        top_indices = np.argsort(similarities)[-top_k:][::-1]

        results = []
        for idx in top_indices:
            results.append((self.documents[idx], similarities[idx]))

        return results

class RAGPipeline:
    """RAG 管道"""

    def __init__(
        self,
        embedding_model,
        llm_model,
        tokenizer,
        vector_store: SimpleVectorStore,
        max_context_length=2048
    ):
        self.embedding_model = embedding_model
        self.llm_model = llm_model
        self.tokenizer = tokenizer
        self.vector_store = vector_store
        self.max_context_length = max_context_length

    def embed_text(self, text: str) -> np.ndarray:
        """将文本转换为向量"""
        # 使用嵌入模型编码
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512
        )

        with torch.no_grad():
            outputs = self.embedding_model(**inputs)
            # 取 [CLS] token 或平均池化
            embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy()

        return embedding[0]

    def retrieve(self, query: str, top_k: int = 5) -> List[Document]:
        """检索相关文档"""
        query_embedding = self.embed_text(query)
        results = self.vector_store.search(query_embedding, top_k)
        return [doc for doc, _ in results]

    def build_prompt(self, query: str, retrieved_docs: List[Document]) -> str:
        """构建带上下文的提示"""

        # 构建上下文
        context = "\n\n".join([
            f"文档 {i+1}:\n{doc.content[:500]}"
            for i, doc in enumerate(retrieved_docs)
        ])

        # 构建提示
        prompt = f"""请基于以下参考资料回答问题。如果参考资料中没有相关信息，请说明无法回答。

参考资料：
{context}

问题：{query}

回答："""

        return prompt

    def generate(self, prompt: str, max_length=256, temperature=0.7) -> str:
        """生成回答"""
        inputs = self.tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=self.max_context_length
        ).to(self.llm_model.device)

        with torch.no_grad():
            outputs = self.llm_model.generate(
                **inputs,
                max_new_tokens=max_length,
                temperature=temperature,
                do_sample=True,
                top_p=0.95,
                pad_token_id=self.tokenizer.eos_token_id
            )

        response = self.tokenizer.decode(
            outputs[0][inputs['input_ids'].shape[1]:],
            skip_special_tokens=True
        )

        return response

    def query(self, question: str, top_k: int = 5) -> dict:
        """执行 RAG 查询"""

        # 1. 检索相关文档
        retrieved_docs = self.retrieve(question, top_k)

        # 2. 构建提示
        prompt = self.build_prompt(question, retrieved_docs)

        # 3. 生成回答
        answer = self.generate(prompt)

        return {
            "question": question,
            "answer": answer,
            "retrieved_documents": [
                {"content": doc.content[:200] + "...", "metadata": doc.metadata}
                for doc in retrieved_docs
            ]
        }

# 文档处理工具
class DocumentProcessor:
    """文档处理器"""

    @staticmethod
    def chunk_text(text: str, chunk_size=512, overlap=50) -> List[str]:
        """将长文本分割成块"""
        chunks = []
        start = 0

        while start < len(text):
            end = start + chunk_size

            # 尝试在句子边界分割
            if end < len(text):
                # 寻找最近的句号、问号或感叹号
                for sep in ['。', '！', '？', '.', '!', '?', '\n']:
                    pos = text.rfind(sep, start, end)
                    if pos != -1:
                        end = pos + 1
                        break

            chunk = text[start:end].strip()
            if chunk:
                chunks.append(chunk)

            start = end - overlap if end - overlap > start else end

        return chunks

    @staticmethod
    def process_file(file_path: str) -> List[Document]:
        """处理文件，返回文档列表"""
        import os

        _, ext = os.path.splitext(file_path)

        if ext == '.txt':
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()
        elif ext == '.md':
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()
        else:
            raise ValueError(f"不支持的文件格式: {ext}")

        chunks = DocumentProcessor.chunk_text(text)

        documents = []
        for i, chunk in enumerate(chunks):
            doc = Document(
                content=chunk,
                metadata={
                    "source": file_path,
                    "chunk_id": i,
                    "chunk_hash": hashlib.md5(chunk.encode()).hexdigest()[:8]
                }
            )
            documents.append(doc)

        return documents

# 使用示例
def setup_rag_pipeline():
    """设置 RAG 管道"""
    from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM

    # 1. 加载嵌入模型
    embedding_model_name = "sentence-transformers/all-MiniLM-L6-v2"
    embedding_tokenizer = AutoTokenizer.from_pretrained(embedding_model_name)
    embedding_model = AutoModel.from_pretrained(embedding_model_name)

    # 2. 加载 LLM
    llm_model_name = "Qwen/Qwen2-1.5B-Instruct"
    llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name, trust_remote_code=True)
    llm_model = AutoModelForCausalLM.from_pretrained(
        llm_model_name,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True
    )

    # 3. 创建向量存储
    vector_store = SimpleVectorStore()

    # 4. 处理文档并添加向量
    processor = DocumentProcessor()

    # 假设有一些文档
    sample_docs = [
        "人工智能是计算机科学的一个分支，致力于创建能够执行通常需要人类智能的任务的系统。",
        "机器学习是人工智能的子集，它使计算机能够从数据中学习而无需显式编程。",
        "深度学习是机器学习的子集，使用多层神经网络来学习数据的层次化表示。"
    ]

    for i, text in enumerate(sample_docs):
        doc = Document(
            content=text,
            metadata={"id": i, "source": "sample"}
        )
        # 计算嵌入
        inputs = embedding_tokenizer(text, return_tensors="pt", truncation=True)
        with torch.no_grad():
            outputs = embedding_model(**inputs)
            doc.embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy()[0]
        vector_store.add_document(doc)

    # 5. 创建 RAG 管道
    rag = RAGPipeline(
        embedding_model=embedding_model,
        llm_model=llm_model,
        tokenizer=llm_tokenizer,
        vector_store=vector_store
    )

    return rag

# 测试
# rag = setup_rag_pipeline()
# result = rag.query("什么是深度学习？")
# print(result["answer"])

class AdvancedRAG(RAGPipeline):
    """高级 RAG 技术"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conversation_history = []

    def rerank_documents(
        self, 
        query: str, 
        documents: List[Document], 
        top_k: int = 5
    ) -> List[Document]:
        """
        重排序文档
        使用交叉编码器或 LLM 对检索结果重新排序
        """
        # 简单实现：基于关键词重叠度重排序
        query_words = set(query.lower().split())

        scored_docs = []
        for doc in documents:
            doc_words = set(doc.content.lower().split())
            overlap = len(query_words & doc_words) / len(query_words)
            scored_docs.append((doc, overlap))

        scored_docs.sort(key=lambda x: x[1], reverse=True)
        return [doc for doc, _ in scored_docs[:top_k]]

    def self_query_retrieval(self, query: str) -> List[Document]:
        """
        自查询检索
        让 LLM 先解析查询意图，生成检索条件
        """
        parse_prompt = f"""分析以下查询，提取关键信息和检索条件。

查询：{query}

请输出JSON格式：
{{
    "keywords": ["关键词1", "关键词2"],
    "entities": ["实体1", "实体2"],
    "intent": "查询意图",
    "filters": {{"字段": "值"}}
}}"""

        # 让 LLM 解析查询
        parse_result = self.generate(parse_prompt, max_length=200)

        # 基于解析结果进行检索（简化实现）
        try:
            import json
            parsed = json.loads(parse_result)
            search_query = " ".join(parsed.get("keywords", [query]))
        except:
            search_query = query

        return self.retrieve(search_query)

    def hybrid_search(
        self, 
        query: str, 
        alpha: float = 0.5,
        top_k: int = 5
    ) -> List[Document]:
        """
        混合检索：结合稠密检索和稀疏检索（BM25）
        """
        # 稠密检索
        dense_results = self.retrieve(query, top_k=top_k*2)

        # 稀疏检索（BM25）
        # 这里需要预先建立 BM25 索引
        sparse_results = self.bm25_search(query, top_k=top_k*2)

        # 融合分数
        combined_scores = {}

        for rank, doc in enumerate(dense_results):
            score = 1.0 / (rank + 1)
            combined_scores[doc.metadata.get("id", id(doc))] = {
                "doc": doc,
                "score": alpha * score
            }

        for rank, doc in enumerate(sparse_results):
            doc_id = doc.metadata.get("id", id(doc))
            score = 1.0 / (rank + 1)
            if doc_id in combined_scores:
                combined_scores[doc_id]["score"] += (1 - alpha) * score
            else:
                combined_scores[doc_id] = {
                    "doc": doc,
                    "score": (1 - alpha) * score
                }

        # 按分数排序
        sorted_items = sorted(
            combined_scores.values(),
            key=lambda x: x["score"],
            reverse=True
        )

        return [item["doc"] for item in sorted_items[:top_k]]

    def conversational_rag(self, query: str) -> dict:
        """对话式 RAG：考虑对话历史"""

        # 构建带历史的查询
        history_context = ""
        if self.conversation_history:
            history_context = "对话历史：\n"
            for turn in self.conversation_history[-3:]:  # 最近3轮
                history_context += f"用户: {turn['user']}\n助手: {turn['assistant']}\n"

        # 改写查询，融入历史信息
        if self.conversation_history:
            rewrite_prompt = f"""{history_context}
基于以上对话历史，将当前问题改写为一个独立的问题。

当前问题：{query}
改写后的问题："""

            rewritten_query = self.generate(rewrite_prompt, max_length=100)
        else:
            rewritten_query = query

        # 使用改写后的查询进行检索
        retrieved_docs = self.retrieve(rewritten_query)

        # 构建最终提示
        final_prompt = f"""{history_context}
参考资料：
{self._format_docs(retrieved_docs)}

请基于对话历史和参考资料回答用户问题。

用户问题：{query}
回答："""

        answer = self.generate(final_prompt)

        # 更新历史
        self.conversation_history.append({
            "user": query,
            "assistant": answer
        })

        return {
            "query": query,
            "rewritten_query": rewritten_query,
            "answer": answer,
            "retrieved_docs": retrieved_docs
        }

    def _format_docs(self, docs: List[Document]) -> str:
        """格式化文档"""
        return "\n\n".join([
            f"[{i+1}] {doc.content[:300]}"
            for i, doc in enumerate(docs)
        ])

    def bm25_search(self, query: str, top_k: int = 5) -> List[Document]:
        """BM25 稀疏检索（简化实现）"""
        # 实际项目中应使用 rank_bm25 库
        # from rank_bm25 import BM25Okapi

        # 简化：基于词频的检索
        query_terms = query.lower().split()
        scored_docs = []

        for doc in self.vector_store.documents:
            doc_terms = doc.content.lower().split()
            score = sum(1 for term in query_terms if term in doc_terms)
            if score > 0:
                scored_docs.append((doc, score))

        scored_docs.sort(key=lambda x: x[1], reverse=True)
        return [doc for doc, _ in scored_docs[:top_k]]

"""
大模型评估指标
"""

import numpy as np
from typing import List, Dict
from collections import Counter
import re

class LLMEvaluator:
    """大模型评估器"""

    @staticmethod
    def perplexity(logits: np.ndarray, target_ids: np.ndarray) -> float:
        """
        计算困惑度（Perplexity）
        困惑度越低，模型对数据的拟合越好
        """
        # 获取目标位置的 logits
        batch_size, seq_len, vocab_size = logits.shape

        # 计算交叉熵损失
        log_probs = np.log(np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True))

        # 获取目标 token 的对数概率
        target_log_probs = log_probs[np.arange(batch_size)[:, None], 
                                     np.arange(seq_len), 
                                     target_ids]

        # 计算困惑度
        avg_neg_log_prob = -np.mean(target_log_probs)
        perplexity = np.exp(avg_neg_log_prob)

        return perplexity

    @staticmethod
    def bleu_score(reference: str, candidate: str, n_gram: int = 4) -> float:
        """
        计算 BLEU 分数
        用于评估生成文本与参考文本的相似度
        """
        def get_ngrams(text: str, n: int) -> Counter:
            words = text.lower().split()
            ngrams = [tuple(words[i:i+n]) for i in range(len(words)-n+1)]
            return Counter(ngrams)

        # 计算各阶 n-gram 精度
        precisions = []
        for n in range(1, n_gram + 1):
            ref_ngrams = get_ngrams(reference, n)
            cand_ngrams = get_ngrams(candidate, n)

            if not cand_ngrams:
                precisions.append(0)
                continue

            # 截断计数
            clipped_count = 0
            for ngram, count in cand_ngrams.items():
                clipped_count += min(count, ref_ngrams.get(ngram, 0))

            precision = clipped_count / sum(cand_ngrams.values())
            precisions.append(precision)

        # 几何平均
        if min(precisions) == 0:
            return 0

        bleu = np.exp(np.mean(np.log(precisions)))

        # 简短惩罚
        ref_len = len(reference.split())
        cand_len = len(candidate.split())
        if cand_len < ref_len:
            bp = np.exp(1 - ref_len / cand_len)
            bleu *= bp

        return bleu

    @staticmethod
    def rouge_score(reference: str, candidate: str) -> Dict[str, float]:
        """
        计算 ROUGE 分数
        包括 ROUGE-1, ROUGE-2, ROUGE-L
        """
        def lcs_length(s1: List[str], s2: List[str]) -> int:
            """最长公共子序列长度"""
            m, n = len(s1), len(s2)
            dp = [[0] * (n + 1) for _ in range(m + 1)]

            for i in range(1, m + 1):
                for j in range(1, n + 1):
                    if s1[i-1] == s2[j-1]:
                        dp[i][j] = dp[i-1][j-1] + 1
                    else:
                        dp[i][j] = max(dp[i-1][j], dp[i][j-1])

            return dp[m][n]

        ref_words = reference.lower().split()
        cand_words = candidate.lower().split()

        # ROUGE-1 (unigram)
        ref_unigrams = Counter(ref_words)
        cand_unigrams = Counter(cand_words)

        overlap_1 = sum(min(cand_unigrams[w], ref_unigrams[w]) 
                        for w in cand_unigrams if w in ref_unigrams)
        rouge_1 = overlap_1 / len(cand_words) if cand_words else 0

        # ROUGE-2 (bigram)
        ref_bigrams = Counter(zip(ref_words, ref_words[1:]))
        cand_bigrams = Counter(zip(cand_words, cand_words[1:]))

        overlap_2 = sum(min(cand_bigrams[b], ref_bigrams[b])
                        for b in cand_bigrams if b in ref_bigrams)
        rouge_2 = overlap_2 / len(cand_bigrams) if cand_bigrams else 0

        # ROUGE-L (LCS)
        lcs = lcs_length(ref_words, cand_words)
        rouge_l_f1 = 2 * (lcs / len(ref_words)) * (lcs / len(cand_words)) / \
                     ((lcs / len(ref_words)) + (lcs / len(cand_words))) \
                     if ref_words and cand_words else 0

        return {
            "rouge-1": rouge_1,
            "rouge-2": rouge_2,
            "rouge-l": rouge_l_f1
        }

    @staticmethod
    def exact_match(prediction: str, ground_truth: str) -> float:
        """精确匹配"""
        return 1.0 if prediction.strip() == ground_truth.strip() else 0.0

    @staticmethod
    def f1_score(prediction: str, ground_truth: str) -> float:
        """F1 分数"""
        pred_tokens = set(prediction.lower().split())
        truth_tokens = set(ground_truth.lower().split())

        if not pred_tokens or not truth_tokens:
            return 0.0

        intersection = pred_tokens & truth_tokens
        precision = len(intersection) / len(pred_tokens)
        recall = len(intersection) / len(truth_tokens)

        if precision + recall == 0:
            return 0.0

        return 2 * precision * recall / (precision + recall)

# 使用示例
evaluator = LLMEvaluator()

reference = "机器学习是人工智能的一个分支"
candidate = "机器学习属于人工智能的一个子领域"

bleu = evaluator.bleu_score(reference, candidate)
rouge = evaluator.rouge_score(reference, candidate)
f1 = evaluator.f1_score(candidate, reference)

print(f"BLEU: {bleu:.4f}")
print(f"ROUGE: {rouge}")
print(f"F1: {f1:.4f}")

"""
vLLM 是一个高吞吐量的 LLM 推理引擎
安装：pip install vllm
"""

from vllm import LLM, SamplingParams
from vllm.outputs import RequestOutput
from typing import List, Dict, Any
import asyncio

class VLLMInference:
    """vLLM 推理封装"""

    def __init__(
        self,
        model_name: str = "Qwen/Qwen2-7B-Instruct",
        tensor_parallel_size: int = 1,
        max_model_len: int = 4096,
        trust_remote_code: bool = True,
    ):
        """
        初始化 vLLM

        Args:
            model_name: 模型名称或路径
            tensor_parallel_size: 张量并行数（多GPU）
            max_model_len: 最大模型长度
        """
        self.llm = LLM(
            model=model_name,
            tensor_parallel_size=tensor_parallel_size,
            max_model_len=max_model_len,
            trust_remote_code=trust_remote_code,
        )

        self.tokenizer = self.llm.get_tokenizer()

    def generate(
        self,
        prompts: List[str],
        temperature: float = 0.7,
        top_p: float = 0.95,
        max_tokens: int = 512,
        stop: List[str] = None,
    ) -> List[str]:
        """
        批量生成

        Args:
            prompts: 提示词列表
            temperature: 温度参数
            top_p: nucleus sampling 参数
            max_tokens: 最大生成 token 数
            stop: 停止词列表

        Returns:
            生成的文本列表
        """
        sampling_params = SamplingParams(
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            stop=stop,
        )

        outputs = self.llm.generate(prompts, sampling_params)

        return [output.outputs[0].text for output in outputs]

    def chat(
        self,
        messages: List[List[Dict[str, str]]],
        **kwargs
    ) -> List[str]:
        """
        对话生成

        Args:
            messages: 对话消息列表，每个元素是一个对话
                     [{"role": "system", "content": "..."},
                      {"role": "user", "content": "..."}]
        """
        # 应用聊天模板
        prompts = [
            self.tokenizer.apply_chat_template(
                msg,
                tokenize=False,
                add_generation_prompt=True
            )
            for msg in messages
        ]

        return self.generate(prompts, **kwargs)

    async def async_generate(
        self,
        prompts: List[str],
        **kwargs
    ) -> List[str]:
        """异步批量生成"""
        sampling_params = SamplingParams(**kwargs)

        # 创建异步请求
        async def process_prompt(prompt):
            return await self.llm.add_request(prompt, sampling_params)

        tasks = [process_prompt(p) for p in prompts]
        outputs = await asyncio.gather(*tasks)

        return [output.outputs[0].text for output in outputs]

# 使用示例
"""
# 初始化
vllm = VLLMInference("Qwen/Qwen2-7B-Instruct", tensor_parallel_size=1)

# 批量生成
prompts = [
    "解释什么是人工智能",
    "用Python写一个快速排序",
    "介绍一下深度学习"
]

responses = vllm.generate(prompts, max_tokens=256)
for prompt, response in zip(prompts, responses):
    print(f"Q: {prompt}")
    print(f"A: {response}\n")

# 对话模式
messages = [
    [
        {"role": "system", "content": "你是一个Python专家"},
        {"role": "user", "content": "如何读取CSV文件？"}
    ]
]

chat_responses = vllm.chat(messages)
print(chat_responses[0])
"""

"""
使用 FastAPI 部署大模型服务
"""

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
import asyncio
import json
import time
import uuid

# 请求模型
class GenerateRequest(BaseModel):
    prompt: str = Field(..., description="输入提示")
    max_tokens: int = Field(256, description="最大生成 token 数")
    temperature: float = Field(0.7, ge=0, le=2, description="温度参数")
    top_p: float = Field(0.95, ge=0, le=1, description="核采样参数")
    stop: Optional[List[str]] = Field(None, description="停止词")

class ChatMessage(BaseModel):
    role: str = Field(..., description="角色：system/user/assistant")
    content: str = Field(..., description="消息内容")

class ChatRequest(BaseModel):
    messages: List[ChatMessage] = Field(..., description="对话消息")
    max_tokens: int = Field(512, description="最大生成 token 数")
    temperature: float = Field(0.7, description="温度参数")
    stream: bool = Field(False, description="是否流式输出")

class GenerateResponse(BaseModel):
    id: str = Field(..., description="请求ID")
    text: str = Field(..., description="生成文本")
    usage: Dict[str, int] = Field(..., description="使用统计")

# 创建 FastAPI 应用
app = FastAPI(title="LLM Inference API", version="1.0.0")

# 全局模型实例
llm_engine = None

@app.on_event("startup")
async def load_model():
    """启动时加载模型"""
    global llm_engine
    from vllm import AsyncLLMEngine, AsyncEngineArgs

    engine_args = AsyncEngineArgs(
        model="Qwen/Qwen2-7B-Instruct",
        tensor_parallel_size=1,
        max_model_len=4096,
        trust_remote_code=True,
    )

    llm_engine = AsyncLLMEngine.from_engine_args(engine_args)
    print("模型加载完成")

@app.post("/v1/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
    """文本生成接口"""
    global llm_engine

    request_id = str(uuid.uuid4())

    try:
        from vllm import SamplingParams

        sampling_params = SamplingParams(
            temperature=request.temperature,
            top_p=request.top_p,
            max_tokens=request.max_tokens,
            stop=request.stop,
        )

        # 异步生成
        results_generator = llm_engine.generate(
            request.prompt,
            sampling_params,
            request_id
        )

        final_output = None
        async for request_output in results_generator:
            final_output = request_output

        if final_output is None:
            raise HTTPException(status_code=500, detail="生成失败")

        output_text = final_output.outputs[0].text

        return GenerateResponse(
            id=request_id,
            text=output_text,
            usage={
                "prompt_tokens": len(final_output.prompt_token_ids),
                "completion_tokens": len(final_output.outputs[0].token_ids),
                "total_tokens": len(final_output.prompt_token_ids) + 
                               len(final_output.outputs[0].token_ids)
            }
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/chat")
async def chat(request: ChatRequest):
    """对话接口"""
    global llm_engine

    # 应用聊天模板
    tokenizer = await llm_engine.get_tokenizer()

    prompt = tokenizer.apply_chat_template(
        [{"role": m.role, "content": m.content} for m in request.messages],
        tokenize=False,
        add_generation_prompt=True
    )

    if request.stream:
        # 流式输出
        return StreamingResponse(
            stream_generate(prompt, request),
            media_type="text/event-stream"
        )
    else:
        # 非流式输出
        gen_request = GenerateRequest(
            prompt=prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature
        )
        return await generate(gen_request)

async def stream_generate(prompt: str, request: ChatRequest):
    """流式生成"""
    global llm_engine

    from vllm import SamplingParams

    sampling_params = SamplingParams(
        temperature=request.temperature,
        max_tokens=request.max_tokens,
    )

    request_id = str(uuid.uuid4())
    results_generator = llm_engine.generate(prompt, sampling_params, request_id)

    async for request_output in results_generator:
        if request_output.outputs:
            text = request_output.outputs[0].text
            yield f"data: {json.dumps({'text': text})}\n\n"

    yield "data: [DONE]\n\n"

@app.get("/v1/models")
async def list_models():
    """列出可用模型"""
    return {
        "object": "list",
        "data": [
            {
                "id": "Qwen2-7B-Instruct",
                "object": "model",
                "created": int(time.time()),
                "owned_by": "local"
            }
        ]
    }

@app.get("/health")
async def health_check():
    """健康检查"""
    return {"status": "healthy", "model_loaded": llm_engine is not None}

# 运行服务
"""
uvicorn app:app --host 0.0.0.0 --port 8000
"""

"""
大模型学习路线图：

第一阶段：基础入门（1-2个月）
├── Python 基础
├── PyTorch/TensorFlow 基础
├── Transformer 架构理解
└── 使用 Hugging Face 调用模型

第二阶段：深入理解（2-3个月）
├── 注意力机制详解
├── 预训练目标（MLM, CLM）
├── 分词器原理（BPE, WordPiece）
├── 模型架构对比（GPT, LLaMA, Qwen）
└── 提示工程实践

第三阶段：微调与应用（2-3个月）
├── LoRA/QLoRA 微调
├── 指令微调（Instruction Tuning）
├── RLHF 原理
├── RAG 系统构建
└── Agent 开发

第四阶段：优化与部署（2-3个月）
├── 量化技术（GPTQ, AWQ）
├── 推理优化（vLLM, TensorRT-LLM）
├── 分布式训练（DeepSpeed, FSDP）
├── 服务部署（FastAPI, Ray Serve）
└── 监控与评估

第五阶段：前沿探索（持续）
├── MoE 架构
├── 多模态模型
├── 长上下文技术
├── 模型压缩与蒸馏
└── 新架构研究（Mamba, RWKV）
"""

"""
推荐学习资源：

📚 论文必读：
1. Attention Is All You Need (Transformer)
2. GPT 系列论文
3. LLaMA: Open and Efficient Foundation Language Models
4. LoRA: Low-Rank Adaptation of Large Language Models
5. RAG: Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

🛠️ 开源框架：
1. Hugging Face Transformers - 模型库
2. LangChain - LLM 应用框架
3. vLLM - 高效推理
4. Axolotl / LLaMA-Factory - 微调框架
5. Ollama - 本地运行

📖 在线课程：
1. DeepLearning.AI - LangChain & LLM 课程
2. Fast.ai - Practical Deep Learning
3. Hugging Face NLP Course

💻 实践项目建议：
1. 个人知识库问答系统
2. 代码助手
3. 文档摘要工具
4. 多轮对话机器人
5. Agent 任务执行器
"""

"""
大模型开发核心要点：

1. 理解模型能力边界
   - 知道模型能做什么、不能做什么
   - 合理设置期望值

2. 提示工程是关键
   - 好的提示能大幅提升效果
   - 不断迭代优化提示模板

3. RAG 解决知识更新问题
   - 结合外部知识库
   - 减少幻觉问题

4. 微调适应特定领域
   - LoRA 降低微调成本
   - 高质量数据最重要

5. 评估不可或缺
   - 建立评估数据集
   - 自动化评估 + 人工评估

6. 部署考虑成本
   - 量化降低显存需求
   - 批处理提高吞吐量
   - 选择合适的模型规模

7. 安全与伦理
   - 内容审核
   - 避免偏见和歧视
   - 保护用户隐私
"""