fix(chapter5): align labels/attention_mask semantics and add padding-aware batch generation (#170)

fix: typo

.gitignore

@@ -0,0 +1,164 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+.idea/
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+.history
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+.DS_Store
+.obsidian
+
+.claude/

Extra-Chapter/CDDRS/readme.md

@@ -0,0 +1,879 @@
+# 建筑文档智能RAG审查系统
+
+一个从零开始实现的建筑文档智能审查系统，旨在帮助开发者理解知识引导检索在专业领域文档审查中的核心原理和实现细节。
+
+## 项目动机
+
+建筑施工交底文档的合规性审查是保障施工项目安全性、经济性的关键环节。在施工项目全周期中，各项操作必须符合相关规范条文要求，才能确保建设项目的安全性与可持续性。然而，相关查询参考往往分散在各个项目文件中，传统基于人工的审查方法难以处理庞大复杂的建筑条文，其审查过程需要基于审查人员的经验与专业知识，具有主观性强，耗时长且易出错等弊端。
+
+随着大语言模型技术的发展，LLM为自动化建筑文档审查带来了新的希望。然而，大语言模型通常使用通用语料进行训练，缺乏建筑相关背景知识，在处理建造背景下的复杂推理问题中会产生严重的幻觉现象。通过使用基于向量相似匹配的RAG方法，可以为LLMs提供初步的相似参考知识，从而减轻基于人工或规则的审查方法难以处理庞大建筑文本所带来的错误率高的问题。
+
+然而，传统RAG方法在建筑专业文档审查中存在关键局限：由于固定的分块设计，使得文本块之间面临知识信息缺失问题；在检索过程中，使用整句问询嵌入的方法进行相似性匹配，缺少对问询细粒度特征的识别与考量，检索效率低下。在建筑施工交底文档中，这类文档详细阐述了施工工艺特点和方法、质量规格、操作程序以及安全协议，包含大量知识细节且专业性极强。因此需要一个能够精准理解和检索建筑领域专业知识的智能系统。
+
+因此，本项目提出了一个生成式知识引导的建筑文档审查系统，旨在提升审查的可靠性和准确性。系统具有两大核心创新：首先提出动态语义知识分块策略，构建具有更优语义连贯性和完整性的知识库；其次基于增强的知识表示，提出生成式知识引导检索框架，在语义嵌入检索过程中增强对细粒度信息的关注，从而提高知识参考检索的准确性和建筑文档审查任务中修正的可靠性。
+
+需要注意的是，由于篇幅限制，我们无法展示完整的整个实现过程，但是，我们将在文档中讲解每个必要的实现步骤以及背后的思考，您可以通过这些内容快速理解如何实现一个建筑文档智能审查系统。
+
+## 前置实现
+
+接下来，我们将带领大家，从0开始，实现一个建筑文档智能审查系统。首先，我们将完成一些基本的准备过程。
+
+### 1. 实现 LLM 模块
+
+首先我们需要实现 LLM 模块，这是系统中最基本的模块，我们将利用大模型完成文档的清洗，信息提取等工作，可以说本系统的一部分精髓即为使用大模型预先处理文档信息，方便后续进行检索，这里我们使用 DeepSeek 的 api 来实现。
+
+```python
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+class BaseLLM(ABC):
+    """Interface for large language models."""
+
+    def __init__(
+        self,
+        model_name: str,
+        model_params: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        self.model_name = model_name
+        self.model_params = model_params or {}
+
+    @abstractmethod
+    def predict(self, input: str) -> str:
+        """Sends a text input to the LLM and retrieves a response."""
+```
+
+如上是一个调用大模型的抽象接口，这可以帮助我们统一调用大模型的格式，我们继承这个基类，实现调用大模型的接口。
+
+```python
+from openai import OpenAI
+from typing import Any, Optional
+from .base import BaseLLM
+
+class DeepSeekLLM(BaseLLM):
+    """Implementation of the BaseLLM interface using DeepSeek API."""
+
+    def __init__(
+        self,
+        model_name: str,
+        api_key: str,
+        base_url: str = "https://api.deepseek.com/v1",
+        model_params: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name, model_params, **kwargs)
+        self.client = OpenAI(api_key=api_key, base_url=base_url)
+
+    def predict(self, input: str) -> str:
+        response = self.client.chat.completions.create(
+            model=self.model_name,
+            messages=[{"role": "user", "content": input}],
+        )
+        return response.choices[0].message.content
+```
+
+完成搭建后，我们可以通过尝试调用 predict 方法来测试是否成功。
+
+```python
+llm = DeepSeekLLM(
+    model_name="deepseek-chat", 
+    api_key="your-api-key-here",
+    base_url="https://api.deepseek.com/v1"
+)
+print(llm.predict("你好，你能帮助我进行建筑文档审查吗？"))
+```
+
+当观察到 LLM 正确回复后，我们这一模块的构建就完成了。
+
+### 2. 实现 Embedding 模块
+
+除了调用大模型，我们还需要实现 Embedding 模块，Embedding 模块用于将文本转换为向量，我们将使用向量来表示文档中的信息，这样的好处是，我们可以通过向量的相似度来衡量文档与查询之间的相似度，从而召回对回复用户问题最有帮助的文档。
+
+构建 Embedding 模块的方法与构建 LLM 模块类似。
+
+```python
+from abc import ABC, abstractmethod
+from typing import List, Any, Optional
+
+class BaseEmb(ABC):
+    def __init__(
+        self,
+        model_name: str,
+        model_params: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        self.model_name = model_name
+        self.model_params = model_params or {}
+
+    @abstractmethod
+    def get_emb(self, input: str) -> List[float]:
+        """Sends a text input to the embedding model and retrieves the embedding."""
+        pass
+
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from .base import BaseEmb
+
+class BGEEmbedding(BaseEmb):
+    def __init__(self, model_name: str = "BAAI/bge-m3", **kwargs):
+        super().__init__(model_name=model_name, **kwargs)
+        self.embed_model = HuggingFaceEmbedding(
+            model_name=model_name,
+            trust_remote_code=True,
+            cache_folder="./model_cache"
+        )
+
+    def get_emb(self, text: str) -> List[float]:
+        embedding = self.embed_model.get_text_embedding(text)
+        return embedding
+```
+
+完成搭建后，我们可以通过尝试调用 get_emb 方法来测试是否成功。
+
+```python
+emb = BGEEmbedding(model_name="BAAI/bge-m3")
+print(emb.get_emb("建筑结构的安全性检查包括哪些方面？"))
+```
+
+当观察到 Embedding 正确给出了编码后的向量，我们这一模块的构建就完成了。
+
+### 3. 实现文档预处理模块
+
+为了处理建筑文档，我们需要预先准备好文档读取模块。本系统假设所有建筑规范和标准已经转换为Markdown格式，便于后续的文本处理和分析。
+
+```python
+import os
+from pathlib import Path
+from typing import Dict, List
+
+class DocumentProcessor:
+    def __init__(self):
+        pass
+    
+    def load_documents(self, directory_path: str) -> List[str]:
+        documents = []
+        
+        for file_path in Path(directory_path).rglob('*.md'):
+            try:
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                    documents.append(content)
+            except Exception as e:
+                print(f"Error reading {file_path}: {e}")
+                    
+        return documents
+```
+
+完成文档预处理模块的设置后，我们就可以采用下面的方法来加载建筑规范文档了。
+
+```python
+processor = DocumentProcessor()
+documents = processor.load_documents("./construction_standards")
+print(f"加载了 {len(documents)} 个建筑规范文档")
+```
+
+## 核心实现
+
+建筑文档审查系统的主要流程如下。首先，让我们来梳理一下建筑文档审查的工作流程，系统的一个核心思想在于，我们需要把用户提供的文档内容通过智能化的问询生成和知识引导检索来识别潜在的合规性问题。与传统RAG方法不同，我们的系统专门针对建筑领域的专业特点进行了优化，能够更准确地理解建筑规范要求，提供更可靠的审查建议。
+
+### 动态语义知识分块
+
+在传统RAG流程中，文本通过设置固定的token数量划分文本区块。然而，固定token数量会在句子中间截断，导致信息缺失。为此，本系统使用基于建筑文本语义动态划分的方式，通过双重语义聚类的方式，完成考虑建筑语义连贯性的知识chunk划分。
+
+首先，将整个文档内容处理成单独句子序列 $S = \{s_0, s_1, \ldots, s_a\}$。通过计算相邻句子间的语义差异度来识别潜在的语义边界：
+
+$$\gamma_i = 1 - \frac{s_{i-1} \cdot s_i}{\|s_{i-1}\| \|s_i\|}$$
+
+基于语义差异度分布自动确定动态阈值：
+
+$$\psi = \text{Quantile}(\Gamma, \frac{a-p}{a})$$
+
+确保最终的分块既保持语义连贯性又满足长度约束：
+
+$$\mathbb{E}[\gamma_{\text{intra}}] < \mathbb{E}[\gamma_{\text{inter}}]$$
+
+```python
+import re
+import numpy as np
+from typing import List, Dict, Tuple
+from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
+
+class DynamicSemanticChunker:
+    def __init__(self, 
+                 embedding_model: str = "BAAI/bge-m3",
+                 max_chunk_length: int = 512,
+                 min_chunk_length: int = 50):
+        self.embedding_model = SentenceTransformer(embedding_model)
+        self.max_chunk_length = max_chunk_length
+        self.min_chunk_length = min_chunk_length
+    
+    def split_text(self, text: str) -> Dict[str, str]:
+        sentences = self._split_into_sentences(text)
+        if len(sentences) == 0:
+            return {}
+        
+        sentence_embeddings = self.embedding_model.encode(sentences)
+        gamma_values = self._compute_semantic_discrepancy(sentence_embeddings)
+        
+        total_tokens = sum(len(s.split()) for s in sentences)
+        baseline_chunks = max(1, total_tokens // self.max_chunk_length)
+        alpha = max(0.1, (len(sentences) - baseline_chunks) / len(sentences))
+        threshold = np.quantile(gamma_values, alpha) if len(gamma_values) > 0 else 0.5
+        
+        boundaries = self._identify_boundaries(gamma_values, threshold)
+        initial_chunks = self._create_initial_chunks(sentences, boundaries)
+        final_chunks = self._enforce_length_constraints(initial_chunks)
+        
+        chunks_dict = {}
+        for i, chunk in enumerate(final_chunks):
+            chunk_id = f"chunk-{i+1:03d}"
+            chunks_dict[chunk_id] = chunk
+        
+        return chunks_dict
+    
+    def _split_into_sentences(self, text: str) -> List[str]:
+        sentence_pattern = r'[。！？；\n]+'
+        sentences = re.split(sentence_pattern, text)
+        
+        cleaned_sentences = []
+        for sentence in sentences:
+            sentence = sentence.strip()
+            if len(sentence) > 5:
+                cleaned_sentences.append(sentence)
+        
+        return cleaned_sentences
+    
+    def _compute_semantic_discrepancy(self, embeddings: np.ndarray) -> List[float]:
+        gamma_values = []
+        
+        for i in range(1, len(embeddings)):
+            similarity = cosine_similarity(
+                embeddings[i-1].reshape(1, -1),
+                embeddings[i].reshape(1, -1)
+            )[0][0]
+            
+            gamma = 1 - similarity
+            gamma_values.append(gamma)
+        
+        return gamma_values
+    
+    def _identify_boundaries(self, gamma_values: List[float], threshold: float) -> List[int]:
+        boundaries = [0]
+        
+        for i, gamma in enumerate(gamma_values):
+            if gamma > threshold:
+                boundaries.append(i + 1)
+        
+        boundaries.append(len(gamma_values) + 1)
+        return sorted(set(boundaries))
+    
+    def _create_initial_chunks(self, sentences: List[str], boundaries: List[int]) -> List[str]:
+        chunks = []
+        
+        for i in range(len(boundaries) - 1):
+            start = boundaries[i]
+            end = boundaries[i + 1]
+            
+            chunk_sentences = sentences[start:end]
+            chunk_text = ' '.join(chunk_sentences)
+            chunks.append(chunk_text)
+        
+        return chunks
+    
+    def _enforce_length_constraints(self, chunks: List[str]) -> List[str]:
+        final_chunks = []
+        
+        for chunk in chunks:
+            chunk_length = len(chunk.split())
+            
+            if chunk_length <= self.max_chunk_length:
+                if chunk_length >= self.min_chunk_length:
+                    final_chunks.append(chunk)
+            else:
+                sub_chunks = self._split_long_chunk(chunk)
+                final_chunks.extend(sub_chunks)
+        
+        return final_chunks
+    
+    def _split_long_chunk(self, chunk: str) -> List[str]:
+        sentences = chunk.split('。')
+        sub_chunks = []
+        current_chunk = ""
+        
+        for sentence in sentences:
+            if sentence.strip():
+                test_chunk = current_chunk + sentence + "。"
+                if len(test_chunk.split()) <= self.max_chunk_length:
+                    current_chunk = test_chunk
+                else:
+                    if current_chunk:
+                        sub_chunks.append(current_chunk.strip())
+                    current_chunk = sentence + "。"
+        
+        if current_chunk:
+            sub_chunks.append(current_chunk.strip())
+        
+        return sub_chunks
+```
+
+### 建筑文档审查系统
+
+整体的审查过程如下图所示。系统获取需要审查的区域后，依据提示生成审查问题推荐，此部分也可供工程师进行相关问题输入或推荐问题选择，生成待审查问题。随后，系统通过生成式知识引导检索框架，依据审查问题在所建文本知识库中检索出相应的知识参考。最终，依据检索的部分与审查原文，进行问题分析与审查修正，完成最终的审查流程。
+
+![picture](images/pic1.png)
+
+#### 审查问题生成
+
+在文档审查流程中，系统引入了双阶段Prompt工程驱动的智能化问询生成机制，旨在对建筑施工交底文档进行预见性分析与风险挖掘，实现对文档潜在问题的高效、精准定位。
+
+阶段1为待查文档主旨目标解构，模型被指示从文本中提炼核心事件、关键技术、工艺流程等要素，结构化地总结文档的核心内容，由此界定本次审查的靶向目标，为后续的精细化问询奠定基础。阶段2为多维度风险探测与定制化问询生成，基于第一阶段提炼的核心要素，通过few-shot等方式引导 LLM 从合规性、安全性、可操作性等多维度对文档进行风险探测。Prompt 指示模型围绕潜在的限制条件、操作流程、以及可能存在的合规性隐患等方面，进行细粒度、多角度的审查提问。
+
+```python
+import re
+
+CORE_COMPONENTS_PROMPT = """
+Task: Your task involves the extraction of crucial information components from a designated text segment. The purpose of this extraction is to assist in uncovering hidden descriptions indicative of regulatory non-compliance. Key information components encompass, but are not limited to, core descriptive events, essential construction techniques, technologies, and associated limitations and restrictions.
+
+Input: {document_chunk}
+Answer:
+"""
+
+REVIEW_QUERIES_PROMPT = """
+Task: Your task is to generate relevant search queries based on the text under review and provided core descriptive references. These queries should target potential areas of non-compliance within the text, facilitating the subsequent retrieval of original regulatory documents for detailed examination.
+
+Input: {document_chunk}
+Core components: {core_components}
+Queries:
+"""
+
+def generate_review_queries(llm, document_chunk: str) -> List[str]:
+    core_prompt = CORE_COMPONENTS_PROMPT.format(document_chunk=document_chunk)
+    core_response = llm.predict(core_prompt)
+    
+    # 生成审查查询
+    queries_prompt = REVIEW_QUERIES_PROMPT.format(
+        document_chunk=document_chunk,
+        core_components=core_response
+    )
+    queries_response = llm.predict(queries_prompt)
+    
+    # 从响应中提取查询列表
+    queries = re.findall(r"'([^']*)'", queries_response)
+    
+    return queries[:5]
+```
+
+#### 知识引导生成式检索
+
+系统的核心创新在于知识引导的检索框架，整个过程分为三个关键步骤。步骤1为句子级编码，主要负责输入查询句子的初始表示学习，计算查询与知识库chunks间的句子级相似度分数。步骤2为知识引导检索，进一步从查询中提取关键信息，利用这些信息结合文档长度自适应加权等机制，对每个知识库chunk进行更详细的评分。步骤3为重排序与增强，使用大语言模型对步骤2检索的结果进行进一步重排序，并利用精炼的知识来增强原始查询。
+![picture](images/pic2.png)
+
+首先建立专门针对建筑领域文本分析的深度提取模块，集成领域预训练BERT进行上下文编码，结合双向LSTM进行建筑法规依赖建模。建立三级重要性分类层次：max（最高）、mid（中等）、lit（字面）优先级。本项目直接通过大语言模型进行关键信息提取，如果需要更精准的效果，可以自行训练BERT模型进行专门的关键信息提取。
+![picture](images/pic3.png)
+
+```python
+import re
+from typing import Dict, Tuple, List
+
+KEY_INFO_EXTRACTION_PROMPT = """
+Your task is to extract key information from the query with three different priority levels:
+
+Maximum priority (max): The most important core concepts or entities
+Medium priority (mid): Important modifiers or qualifying conditions  
+Literal priority (lit): Specific values, standards or specifications
+
+Query: {query}
+max:
+mid:
+lit:
+"""
+
+class KeyInfoExtractor:
+    def __init__(self, llm):
+        self.llm = llm
+
+    def extract_key_info(self, query: str) -> Dict[str, Tuple[str, float]]:
+        prompt = KEY_INFO_EXTRACTION_PROMPT.format(query=query)
+        response = self.llm.predict(prompt)
+        
+        lines = response.strip().split('\n')
+        key_info = {}
+        weights = {'max': 0.5, 'mid': 0.3, 'lit': 0.2}
+        
+        for line in lines:
+            if line.startswith('max:'):
+                key_info['max'] = (line[4:].strip(), weights['max'])
+            elif line.startswith('mid:'):
+                key_info['mid'] = (line[4:].strip(), weights['mid'])
+            elif line.startswith('lit:'):
+                key_info['lit'] = (line[4:].strip(), weights['lit'])
+        
+        return key_info
+```
+
+#### 文档长度自适应因子
+
+在知识引导检索过程中，文档长度自适应因子用于调整不同长度文档的权重分配，确保长短文档都能得到公平的评分机会。该因子的计算考虑了当前文档chunk的长度与平均文档长度的关系。
+
+$$\Lambda_{\text{DL}} = \frac{\overline{|k|} + |k_j|}{2\overline{|k|}}$$
+
+其中 $|k_j|$ 表示当前文档chunk的长度，$\overline{|k|}$ 表示平均文档长度。通过这种归一化处理，可以避免因文档长度差异导致的评分偏差。
+
+```python
+def compute_document_length_factor(chunk_length: int, avg_length: int = 100) -> float:
+    lambda_dl = (avg_length + chunk_length) / (2 * avg_length)
+    return lambda_dl
+```
+
+#### 术语重要性计算
+
+术语重要性指标衡量术语在文档中的显著程度，结合术语频率和文档长度自适应因子，能够更准确地评估术语在当前文档中的重要性。计算公式考虑了术语频率的非线性增长特性。
+
+$$\text{Sign}(t_{e_i}^\tau, k_j) = \frac{2 \cdot f(t_{e_i}^\tau, k_j) \cdot \Lambda_{\text{DL}}}{f(t_{e_i}^\tau, k_j) + 1}$$
+
+其中 $f(t_{e_i}^\tau, k_j)$ 表示术语在文档chunk中的出现频率，$\Lambda_{\text{DL}}$ 为文档长度自适应因子。这种计算方式能够防止高频术语过度影响评分。
+
+```python
+def compute_term_significance(term_freq: int, doc_length_factor: float) -> float:
+    significance = (2 * term_freq * doc_length_factor) / (term_freq + 1)
+    return significance
+```
+
+#### 术语稀有度计算
+
+术语稀有度用于衡量术语在整个知识库中的稀缺程度，稀有度越高的术语在检索中的权重越大。计算采用了改进的IDF公式，增加了平滑处理以避免零除问题。
+
+$\text{Rarity}(t_{e_i}^\tau) = \log\left(\frac{D - \text{df}(t_{e_i}^\tau) + 0.5}{\text{df}(t_{e_i}^\tau) + 0.5} + 1\right)$
+
+其中 $D$ 表示文档总数，$\text{df}(t_{e_i}^\tau)$ 表示包含该术语的文档数量。加一操作确保了对数值始终为正数。
+
+```python
+import numpy as np
+
+def compute_term_rarity(doc_freq: int, total_docs: int) -> float:
+    rarity = np.log((total_docs - doc_freq + 0.5) / (doc_freq + 0.5) + 1)
+    return rarity
+```
+
+#### 连贯性指数评估
+
+连贯性指数反映术语在文档中的分布连贯性，通过滑动窗口技术分析术语在文档中的局部分布情况。连贯性高的术语往往在文档的特定区域集中出现，表明其与文档主题的强相关性。
+
+$$\text{CI}(t_{e_i}^\tau, k_j) = \max_{w \in W, \, t \in w} \frac{\sum I(t = t_{e_i}^\tau) \cdot |w|}{|k_j|}$$
+
+其中 $W$ 表示文档中的滑动窗口集合，$I(t = t_{e_i}^\tau)$ 为指示函数，当窗口中包含该术语时为1，否则为0。
+
+```python
+def compute_coherence_index(term: str, chunk: str, window_size: int = 50) -> float:
+    chunk_tokens = chunk.lower().split()
+    chunk_length = len(chunk_tokens)
+    
+    if chunk_length == 0:
+        return 0.0
+    
+    max_coherence = 0.0
+    
+    for i in range(0, chunk_length - window_size + 1, 10):
+        window = chunk_tokens[i:i + window_size]
+        term_count = window.count(term.lower())
+        
+        if term_count > 0:
+            coherence = (term_count * window_size) / chunk_length
+            max_coherence = max(max_coherence, coherence)
+    
+    return max_coherence
+```
+
+#### 评分融合与检索
+
+将句子级相似度评分与知识级评分进行融合，形成最终的文档相关性评分。融合过程采用加权平均的方式，平衡参数λ控制两种评分方式的重要性。
+
+$\Phi = \lambda \Phi(\mathcal{K}) + (1 - \lambda) \Phi(\mathcal{S})$
+
+其中 $\lambda$ 为平衡参数，$\Phi(\mathcal{K})$ 为知识级评分，$\Phi(\mathcal{S})$ 为句子级评分。通过调整λ值，可以控制系统更偏向语义相似还是知识匹配。当λ=0时，系统完全依赖句子级语义相似度；当λ=1时，系统完全依赖知识匹配评分；λ=0.5时，两种评分方式权重相等。在建筑文档审查场景中，通常设置λ=0.5以平衡专业知识匹配和语义理解。
+
+```python
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+from typing import List, Tuple, Dict, Any
+
+class GKGRRetriever:
+    def __init__(self, 
+                 knowledge_base: List[str],
+                 embedding_model,
+                 key_info_extractor: KeyInfoExtractor,
+                 llm,
+                 config: Dict[str, Any] = None):
+        self.knowledge_base = knowledge_base
+        self.embedding_model = embedding_model
+        self.key_info_extractor = key_info_extractor
+        self.llm = llm
+        
+        default_config = {
+            "lambda_param": 0.5,
+            "top_k": 5,
+            "rerank_enabled": True,
+            "query_expansion": True,
+            "similarity_threshold": 0.1
+        }
+        self.config = {**default_config, **(config or {})}
+        
+        self.kb_embeddings = self._precompute_embeddings()
+    
+    def _precompute_embeddings(self) -> np.ndarray:
+        embeddings = self.embedding_model.encode(self.knowledge_base, show_progress_bar=True)
+        return embeddings
+    
+    def retrieve_with_scores(self, query: str) -> List[Tuple[str, float, Dict[str, float]]]:
+        query_embedding = self.embedding_model.encode([query])[0]
+        sentence_scores = cosine_similarity(
+            query_embedding.reshape(1, -1), 
+            self.kb_embeddings
+        )[0]
+        
+        key_info = self.key_info_extractor.extract_key_info(query)
+        knowledge_scores = self._compute_knowledge_scores(key_info)
+        
+        final_scores = []
+        for i in range(len(self.knowledge_base)):
+            norm_sent = sentence_scores[i]
+            norm_know = knowledge_scores[i] / max(knowledge_scores) if max(knowledge_scores) > 0 else 0
+            
+            final_score = (self.config["lambda_param"] * norm_know + 
+                          (1 - self.config["lambda_param"]) * norm_sent)
+            final_scores.append(final_score)
+        
+        results_with_scores = []
+        for i, final_score in enumerate(final_scores):
+            if final_score > self.config["similarity_threshold"]:
+                score_details = {
+                    "sentence_score": float(sentence_scores[i]),
+                    "knowledge_score": float(knowledge_scores[i]),
+                    "final_score": float(final_score)
+                }
+                results_with_scores.append((self.knowledge_base[i], final_score, score_details))
+        
+        results_with_scores.sort(key=lambda x: x[1], reverse=True)
+        return results_with_scores[:self.config["top_k"]]
+    
+    def _compute_knowledge_scores(self, key_info: Dict[str, Tuple[str, float]]) -> List[float]:
+        scores = []
+        avg_length = sum(len(chunk.split()) for chunk in self.knowledge_base) / len(self.knowledge_base)
+        
+        for chunk in self.knowledge_base:
+            chunk_score = 0.0
+            chunk_tokens = chunk.lower().split()
+            chunk_length = len(chunk_tokens)
+            
+            lambda_dl = compute_document_length_factor(chunk_length, avg_length)
+            
+            for priority, (info_text, weight) in key_info.items():
+                if not info_text.strip():
+                    continue
+                
+                terms = info_text.lower().split()
+                for term in terms:
+                    if term in chunk_tokens:
+                        tf = chunk_tokens.count(term)
+                        
+                        significance = compute_term_significance(tf, lambda_dl)
+                        
+                        segments_with_term = sum(1 for kb_chunk in self.knowledge_base 
+                                                if term in kb_chunk.lower())
+                        rarity = compute_term_rarity(segments_with_term, len(self.knowledge_base))
+                        
+                        coherence = compute_coherence_index(term, chunk)
+                        
+                        term_score = significance * rarity * (1 + coherence) * weight
+                        chunk_score += term_score
+            
+            scores.append(chunk_score)
+        
+        return scores
+    
+    def retrieve(self, query: str) -> Tuple[List[str], str]:
+        results_with_scores = self.retrieve_with_scores(query)
+        
+        documents = [doc for doc, _, _ in results_with_scores]
+        
+        if self.config["rerank_enabled"] and len(documents) > 1:
+            documents = self._llm_rerank(query, documents)
+        
+        augmented_query = query
+        if self.config["query_expansion"]:
+            augmented_query = self._augment_query(query, documents[:3])
+        
+        return documents, augmented_query
+```
+
+#### 重排序优化
+
+系统使用大语言模型对检索结果进行进一步重排序，通过LLM的语义理解能力优化文档的相关性排序。重排序过程中，系统会构造包含查询和候选文档的提示，要求LLM根据相关性对文档进行重新排序。
+
+```python
+def _llm_rerank(self, query: str, documents: List[str]) -> List[str]:
+    if len(documents) <= 1:
+        return documents
+    
+    rerank_prompt = f"""
+Task: A list of documents is shown below. Each document has a number next to it. A question is also provided. Your task is to return the numbers of ALL documents in order of relevance from MOST to LEAST relevant. MUST include EVERY document number exactly once.
+
+Example format:
+    Document 1: <document 1>
+    Document 2: <document 2>
+    Document 3: <document 3>
+    Question: <question>
+    Answer: 3,1,2
+
+Now here are the actual documents and question.
+
+"""
+    for i, doc in enumerate(documents):
+        rerank_prompt += f"Document {i+1}: {doc[:150]}...\n"
+    
+    rerank_prompt += f"Question: {query}\nAnswer:"
+    
+    try:
+        response = self.llm.predict(rerank_prompt)
+        order_nums = [int(x.strip()) - 1 for x in response.split(',') 
+                     if x.strip().isdigit() and 0 <= int(x.strip()) - 1 < len(documents)]
+        
+        reranked = [documents[i] for i in order_nums if i < len(documents)]
+        
+        # 添加遗漏的文档
+        used_indices = set(order_nums)
+        for i, doc in enumerate(documents):
+            if i not in used_indices:
+                reranked.append(doc)
+        
+        return reranked[:len(documents)]
+    except:
+        return documents
+```
+
+#### 查询增强
+
+同时系统还会利用检索到的知识来增强原始查询，生成更具体、更详细的查询用于进一步检索。查询增强通过分析检索结果的上下文信息，识别查询中可能遗漏的关键概念和术语。
+
+```python
+def _augment_query(self, original_query: str, top_results: List[str]) -> str:
+    if not top_results:
+        return original_query
+    
+    document_list = ""
+    for i, doc in enumerate(top_results):
+        document_list += f"Document {i+1}: {doc[:100]}...\n"
+    
+    augment_prompt = f"""
+Task: Your task is to generate a detailed answer to the question by synthesizing information from ALL provided documents. Prioritize relevance, cite document numbers, and structure your response as follows:
+
+Question: {original_query}
+{document_list}
+Answer:
+"""
+    
+    try:
+        augmented = self.llm.predict(augment_prompt)
+        return augmented.strip()
+    except:
+        return original_query
+```
+
+#### 偏差检测分析
+
+在先期知识增强检索阶段获取领域知识后，系统随即进入误差辨析模块。该模块基于检索得到的知识参考，并结合预设的审阅问题，对原文进行细致的偏差检测与评估。
+
+```python
+class ErrorAnalyzer:
+    def __init__(self, llm):
+        self.llm = llm
+    
+    def analyze_errors(self, document_chunk: str, query: str, retrieved_knowledge: List[str]) -> Dict[str, Any]:
+        
+        analysis_prompt = f"""
+Task: Your task is to conduct an error analysis on a given review document, based on a provided review query and relevant reference specifications. This analysis MUST strictly adhere to the provided reference and focus specifically on reviewing and analyzing the original descriptive sections within the review document.
+
+Review document: {document_chunk}
+Query: {query}
+Reference: {chr(10).join([f"{i+1}. {ref}" for i, ref in enumerate(retrieved_knowledge)])}
+Analysis:
+"""
+        
+        analysis = self.llm.predict(analysis_prompt)
+        
+        return {
+            "analysis": analysis,
+            "reference_support": retrieved_knowledge
+        }
+```
+
+#### 修订建议生成
+
+误差辨析模块完成后，系统将输出标记偏差区域以及相关知识佐证。随后，系统进入修订策略生成模块。该模块依据误差分析结果和知识参考，对标记区域进行针对性的修订建议生成，最终实现对原文的知识驱动型自动修正。
+
+```python
+class RevisionGenerator:
+    def __init__(self, llm):
+        self.llm = llm
+    
+    def generate_revisions(self, document_chunk: str, analysis: Dict[str, Any]) -> Dict[str, str]:     
+        revision_prompt = f"""
+Task: Your task is to review and revise the provided document based on the given analysis and corresponding reference specifications. STRICT adherence to the provided reference specifications is required. If the review document aligns with the analysis and reference specifications WITHOUT discrepancies, revision is not necessary.
+
+Review document: {document_chunk}
+Analysis: {analysis['analysis']}
+Reference: {chr(10).join([f"- {ref}" for ref in analysis['reference_support']])}
+Revision:
+"""
+        
+        revision = self.llm.predict(revision_prompt)
+        
+        return {
+            "original_text": document_chunk,
+            "revision_suggestions": revision,
+            "modified_regions": analysis.get("error_regions", []),
+            "confidence": self._calculate_confidence(analysis)
+        }
+    
+    def _calculate_confidence(self, analysis: Dict[str, Any]) -> float:
+        ref_count = len(analysis.get("reference_support", []))
+        error_count = len(analysis.get("error_regions", []))
+        
+        confidence = min(0.9, 0.5 + (ref_count * 0.1) + (error_count * 0.05))
+        return confidence
+```
+
+#### 完整审查流程
+
+将上述所有模块整合，形成完整的文档审查流程。系统首先生成审查问题，然后进行知识引导检索，接着执行错误分析，最后生成修订建议。
+
+```python
+def complete_review_process(document_chunk: str, 
+                          gkgr_framework: GKGRRetriever, 
+                          error_analyzer: ErrorAnalyzer,
+                          revision_generator: RevisionGenerator) -> Dict[str, Any]:    
+    review_queries = generate_review_queries(gkgr_framework.llm, document_chunk)
+    
+    results = {}
+    for query in review_queries[:3]:
+        retrieved_docs, augmented_query = gkgr_framework.retrieve(query)
+        
+        knowledge_refs = retrieved_docs
+        analysis = error_analyzer.analyze_errors(document_chunk, query, knowledge_refs)
+        
+        revision = revision_generator.generate_revisions(document_chunk, analysis)
+        
+        results[query] = {
+            "retrieved_knowledge": retrieved_docs,
+            "augmented_query": augmented_query,
+            "analysis": analysis,
+            "revision": revision
+        }
+    
+    return results
+```
+
+至此，我们就完成了建筑文档智能审查系统的核心实现。
+
+## 实际应用示例
+
+让我们通过一个完整的示例来展示系统的使用：
+
+```python
+# 初始化系统组件
+llm = DeepSeekLLM(
+    model_name='deepseek-chat',
+    api_key='your-api-key',
+    base_url='https://api.deepseek.com/v1'
+)
+
+embedding = BGEEmbedding(model_name="BAAI/bge-m3")
+key_extractor = KeyInfoExtractor(llm)
+
+# 从markdown文档构建知识库
+processor = DocumentProcessor()
+documents = processor.load_documents("./construction_standards")
+
+# 对文档进行动态语义分块
+chunker = DynamicSemanticChunker()
+knowledge_base = []
+for doc in documents:
+    chunks = chunker.split_text(doc)
+    knowledge_base.extend(chunks.values())
+
+# 初始化检索器
+gkgr_retriever = GKGRRetriever(
+    knowledge_base=knowledge_base,
+    embedding_model=embedding,
+    key_info_extractor=key_extractor,
+    llm=llm
+)
+
+# 初始化分析器
+error_analyzer = ErrorAnalyzer(llm)
+revision_generator = RevisionGenerator(llm)
+
+# 待审查的文档内容
+sample_document = """
+钢筋混凝土柱的施工应符合以下要求：
+1. 混凝土强度等级不低于C25
+2. 钢筋保护层厚度为25mm
+3. 混凝土浇筑应连续进行，间歇时间不超过1小时
+4. 养护期间应保持混凝土表面湿润
+"""
+
+# 执行审查
+result = complete_review_process(
+    sample_document, 
+    gkgr_retriever, 
+    error_analyzer, 
+    revision_generator
+)
+
+# 查看审查结果
+for query, analysis in result.items():
+    print(f"审查问题: {query}")
+    print(f"修订建议: {analysis['revision']['revision_suggestions']}")
+    print("-" * 50)
+```
+
+## 扩展性说明
+
+系统可以通过更换知识库轻松适应其他领域。对于特定企业或项目，可以通过微调关键信息提取模型来提升准确性。在性能优化方面，使用动态语义分块可以提升检索质量，预计算并缓存知识库嵌入以提升检索速度，对于大量文档可使用批量处理模式，根据具体应用场景调整λ参数和top-k值。
+
+## 写在最后
+
+恭喜你阅读完此文，你已经充分了解了如何实现一个建筑文档智能审查系统以及其背后的思考。这个系统展示了如何将动态语义分块、知识引导检索和大语言模型有机结合，为建筑行业的文档审查提供了一个实用的解决方案。
+
+虽然当前系统已经取得了不错的效果，但仍有改进空间。全局关联增强方面，当前基于文本块的检索可以进一步结合知识图谱等技术。多模态支持方面，未来可以扩展支持CAD图纸、施工图等视觉信息。实时更新方面，支持知识库的增量更新和动态维护。个性化定制方面，根据不同企业和项目特点进行系统定制。
+
+读者们可以运行项目中的示例代码，体验完整的建筑文档智能审查流程。我们相信这个系统不仅能够提升审查效率，更能为建筑行业的数字化转型贡献力量。
+
+## 致谢
+
+本项目的开发过程中，我们深入研究了建筑工程领域的专业知识和最新的自然语言处理技术。特别感谢建筑行业专家提供的宝贵建议，以及开源社区在技术实现方面的支持。项目代码实现参考了LlamaIndex、Transformers等优秀开源项目的设计理念。
+
+需要说明的是，本项目专门针对建筑施工领域的文档审查场景进行了深度优化。如果您需要处理其他领域的文档，建议根据具体需求对系统进行相应调整。
+
+## 源码获取
+
+本项目的源码以及实例数据存放在 [GitHub 仓库](https://github.com/Hongru0306/CDDRS)。
+
+## 引用
+
+如果您在研究中使用了本项目的成果，请按如下方式引用：
+
+```bibtex
+@article{XIAO2025103618,
+    title = {Generative knowledge-guided review system for construction disclosure documents},
+    journal = {Advanced Engineering Informatics},
+    volume = {68},
+    pages = {103618},
+    year = {2025},
+    issn = {1474-0346},
+    doi = {https://doi.org/10.1016/j.aei.2025.103618},
+    url = {https://www.sciencedirect.com/science/article/pii/S1474034625005117},
+    author = {Hongru Xiao and Jiankun Zhuang and Bin Yang and Jiale Han and Yantao Yu and Songning Lai},
+    keywords = {Construction documents review, Large language model (LLM), Knowledge-guided retrieval, Natural Language Processing (NLP)}
+}
+```

Extra-Chapter/Readme.md

@@ -0,0 +1,120 @@
+<div align="center">
+  <h2>🚀 Happy-LLM 扩展内容</h2>
+  <p><em>社区驱动的大语言模型学习资源</em></p>
+</div>
+
+---
+
+## 📖 为什么会有 Extra Chapter？
+
+&emsp;&emsp;在 Happy-LLM 主教程的基础上，我们发现社区中有许多优秀的学习者和实践者，他们在学习和使用大语言模型的过程中积累了宝贵的经验、独到的见解和实用的技巧。这些内容虽然不属于主教程的核心知识体系，但对于深入理解和应用大语言模型具有重要价值。
+
+**Extra Chapter 的设立目的：**
+
+- 🌟 **汇聚智慧**：收集社区成员的优秀学习笔记、实践经验和技术博客
+- 🔄 **持续更新**：保持内容的时效性，跟上大语言模型领域的快速发展
+- 🤝 **促进交流**：为社区成员提供分享和交流的平台
+- 📚 **补充完善**：对主教程内容进行有益的补充和扩展
+- 💡 **启发思考**：通过不同视角和实践案例，启发读者的深度思考
+
+**Extra Chapter 包含的内容类型：**
+
+- 📝 **学习笔记**：深度学习心得和知识总结
+- 🛠️ **实践案例**：真实项目中的应用经验
+- 🔬 **技术探索**：前沿技术的研究和实验
+- 💭 **思考感悟**：对大语言模型发展的独特见解
+- 🎯 **专题研究**：特定领域或问题的深入分析
+
+---
+
+## 📋 PR 贡献规范
+
+&emsp;&emsp;我们热烈欢迎社区成员为 Extra Chapter 贡献优质内容！为了保证内容质量和项目的整体性，请遵循以下规范：
+
+### 🗂️ 目录结构规范
+
+每个贡献的内容应按照以下目录结构组织：
+
+```
+Extra-Chapter/
+├── your-topic-name/                    # 你的主题文件夹
+│   ├── readme.md                       # 主要内容文件（必需）
+│   ├── images/                         # 图片资源文件夹（可选）
+│   │   ├── figure1.png
+│   │   └── figure2.jpg
+│   ├── code/                           # 代码文件夹（可选）
+│   │   ├── example.py
+│   │   └── requirements.txt
+│   ├── data/                           # 数据文件夹（可选）
+│   │   └── sample_data.json
+│   └── notebook.ipynb                  # Jupyter Notebook（如涉及代码必选）
+└── Readme.md                           # 本文件
+```
+
+### 📝 文件命名规范
+
+1. **主题文件夹命名**：
+   - 使用小写字母和连字符
+   - 名称要简洁明了，能够概括主题内容
+   - 例如：`why-fine-tune-small-large-language-models`、`rag-optimization-techniques`
+
+2. **主要内容文件**：
+   - 必须命名为 `readme.md`
+   - 使用 Markdown 格式编写
+
+3. **图片文件**：
+   - 统一放在 `images/` 文件夹下
+   - 使用描述性的文件名
+   - 支持格式：`.png`、`.jpg`、`.jpeg`、`.gif`、`.svg`
+
+4. **代码文件**：
+   - 如涉及代码，请尽量提供可直接运行的 Jupyter Notebook 文件
+   - 统一放在 `code/` 文件夹下
+   - 使用标准的文件扩展名
+   - 如有依赖，请提供 `requirements.txt`
+   - 如有 Jupyter Notebook 文件，请放在主文件夹下
+
+### ✍️ 内容质量要求
+
+1. **原创性**：
+   - 内容必须是原创或经过授权的
+   - 如引用他人内容，请注明出处
+
+2. **技术准确性**：
+   - 确保技术内容的准确性
+   - 代码示例应能正常运行
+   - 提供必要的环境说明
+
+3. **结构清晰**：
+   - 使用清晰的标题层次
+   - 合理使用列表、表格等格式
+   - 重要内容使用适当的强调
+
+4. **语言规范**：
+   - 使用规范的中文表达
+   - 技术术语使用准确
+   - 避免错别字和语法错误
+
+### PR commit messgae 内容
+
+请在 PR commit message 中 包含以下内容：
+
+- 新增的主题文件夹名称
+- 贡献内容的概述
+- 贡献内容的详细描述
+- 你的 Github 个人主页链接，及你的个人介绍
+- 个人 title 及工作经历 or 学校 or 研究方向
+
+如以下所示：
+
+```
+Extra Chapter: 你的主题名称
+
+详细描述你的贡献内容，包括新增的主题文件夹、文件内容和目录结构。
+
+- 新增的主题文件夹名称：your-topic-name
+- 贡献内容的概述：详细介绍你的贡献内容
+- 贡献内容的详细描述：详细描述你的贡献内容，包括新增的主题文件夹、文件内容和目录结构。
+- 你的 Github 个人主页链接及个人介绍：[你的个人主页链接](https://example.com)，介绍你的研究方向、技术专长等。
+- 个人 title 及工作经历 or 学校 or 研究方向：内容贡献者-xxxx学校，研究方向为自然语言处理。
+```

Extra-Chapter/generation-method/llm_generation.py

@@ -0,0 +1,152 @@
+import torch
+from modelscope import AutoModelForCausalLM, AutoTokenizer
+
+def test_decoding_strategies():
+    """
+    测试三种解码策略：贪婪解码、随机采样、束搜索
+    """
+    model_id = "../model/kmno4zx/happy-llm-215M-sft/"
+
+    print("正在加载模型和tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, device_map="cpu").eval()
+
+    # 测试prompt
+    test_prompt = "请介绍一下自己"
+    messages = [
+        {"role": "system", "content": "你是一个AI助手"},
+        {"role": "user", "content": test_prompt}
+    ]
+
+    # 准备输入
+    input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    input_ids = tokenizer(input_ids).data['input_ids']
+    x = (torch.tensor(input_ids, dtype=torch.long)[None, ...]).to(model.device)
+
+    print(f"测试prompt: {test_prompt}")
+    print(f"输入token数量: {len(input_ids)}")
+    print("=" * 60)
+
+    # 测试1: 贪婪解码 (Greedy Search)
+    print("🔍 测试1: 贪婪解码 (Greedy Search)")
+    print("参数: do_sample=False, num_beams=1, temperature=0.0")
+    print("特点: 每步选择概率最大的token，结果确定，速度快")
+
+    with torch.no_grad():
+        greedy_output = model.generate_super(
+            x,
+            stop_id=tokenizer.eos_token_id,
+            max_new_tokens=50,
+            temperature=0.0,
+            do_sample=False,
+            num_beams=1
+        )
+        greedy_response = tokenizer.decode(greedy_output[0].tolist(), skip_special_tokens=True)
+
+    print(f"贪婪解码结果: {greedy_response}")
+    print()
+
+    # 测试2: 随机采样 (Random Sampling)
+    print("🎲 测试2: 随机采样 (Random Sampling)")
+    print("参数: do_sample=True, num_beams=1, temperature=0.8, top_k=50")
+    print("特点: 基于概率分布随机采样，结果多样，创造性高")
+
+    with torch.no_grad():
+        # 运行多次以展示随机性
+        for i in range(3):
+            sampling_output = model.generate_super(
+                x,
+                stop_id=tokenizer.eos_token_id,
+                max_new_tokens=50,
+                temperature=0.8,
+                top_k=50,
+                do_sample=True,
+                num_beams=1
+            )
+            sampling_response = tokenizer.decode(sampling_output[0].tolist(), skip_special_tokens=True)
+            print(f"随机采样结果 {i+1}: {sampling_response}")
+
+    print()
+
+    # 测试3: 束搜索 (Beam Search)
+    print("🔦 测试3: 束搜索 (Beam Search)")
+    print("参数: do_sample=False, num_beams=3, temperature=1.0")
+    print("特点: 维护多条候选路径，选择总概率最高的序列，质量更高")
+
+    with torch.no_grad():
+        beam_output = model.generate_super(
+            x,
+            stop_id=tokenizer.eos_token_id,
+            max_new_tokens=50,
+            temperature=1.0,
+            do_sample=False,
+            num_beams=3
+        )
+        beam_response = tokenizer.decode(beam_output[0].tolist(), skip_special_tokens=True)
+
+    print(f"束搜索结果: {beam_response}")
+    print()
+
+    # 测试4: 不同的温度参数对随机采样的影响
+    print("🌡️ 测试4: 不同温度参数对随机采样的影响")
+    print("参数: do_sample=True, num_beams=1, 测试不同temperature值")
+
+    temperatures = [0.2, 0.8, 1.5]
+    for temp in temperatures:
+        with torch.no_grad():
+            temp_output = model.generate_super(
+                x,
+                stop_id=tokenizer.eos_token_id,
+                max_new_tokens=30,
+                temperature=temp,
+                do_sample=True,
+                num_beams=1
+            )
+            temp_response = tokenizer.decode(temp_output[0].tolist(), skip_special_tokens=True)
+            print(f"温度 {temp}: {temp_response}")
+
+    print()
+    print("=" * 60)
+    print("✅ 三种解码策略测试完成！")
+    print()
+    print("📊 总结对比:")
+    print("• 贪婪解码: 速度快，结果确定，适合确定性任务")
+    print("• 随机采样: 创造性强，结果多样，适合创意生成")
+    print("• 束搜索: 质量较高，平衡速度和质量，适合一般对话")
+
+def test_original_generation():
+    """
+    原始的生成代码作为对比
+    """
+    model_id = "../model/kmno4zx/happy-llm-215M-sft/"
+
+    print("运行原始生成代码...")
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, device_map="cpu").eval()
+
+    messages = [
+        {"role": "system", "content": "你是一个AI助手"},
+        {"role": "user", "content": "你好，请介绍一下自己。"}
+    ]
+
+    input_ids = tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True)
+    input_ids = tokenizer(input_ids).data['input_ids']
+
+    x = (torch.tensor(input_ids, dtype=torch.long)[None, ...]).to(model.device)
+
+    with torch.no_grad():
+        y = model.generate_super(x, stop_id=tokenizer.eos_token_id, max_new_tokens=512, temperature=0.6)
+        response = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
+
+    print(f"Assistant: {response}")
+
+if __name__ == "__main__":
+    print("开始测试三种解码策略...")
+    print()
+
+    try:
+        test_decoding_strategies()
+    except Exception as e:
+        print(f"测试过程中出现错误: {e}")
+        print("运行原始生成代码...")
+        test_original_generation()

Extra-Chapter/generation-method/model_down.py

@@ -0,0 +1,3 @@
+from modelscope import snapshot_download
+
+model_dir = snapshot_download('kmno4zx/happy-llm-215M-sft', cache_dir='your/cache/dir', revision='master')

Extra-Chapter/generation-method/readme.md

@@ -0,0 +1,511 @@
+# 大模型生成Token的方式
+
+> 代码已更新到 Happy-LLM 仓库第五章的代码中。
+
+## 贪婪解码（Greedy Decoding）
+
+### 原理说明
+贪婪解码是最简单直接的文本生成策略。在每一步生成时，它总是选择概率最大的那个token作为下一个token，然后继续生成，直到遇到停止条件或达到最大长度。
+
+**核心思想**：局部最优选择 → 希望全局最优
+
+**数学表达**：
+```
+token_t = argmax P(token_t | token_1, token_2, ..., token_{t-1})
+```
+
+### 代码实现
+基于我们实现的 `_greedy_decode` 方法：
+
+```python
+def _greedy_decode(self, logits: torch.Tensor) -> torch.Tensor:
+    """
+    贪婪解码：选择概率最大的token
+
+    Args:
+        logits: 模型输出的logits，形状为 (batch_size, vocab_size)
+
+    Returns:
+        选择的token索引，形状为 (batch_size, 1)
+    """
+    _, idx_next = torch.topk(logits, k=1, dim=-1)
+    return idx_next
+```
+
+**关键步骤解析**：
+1. `torch.topk(logits, k=1, dim=-1)`：找到logits中最大值的位置
+2. 返回最大概率token的索引
+3. 该token被添加到序列中，继续下一轮生成
+
+### 使用示例
+```python
+# 在 generate_super 函数中调用贪婪解码
+output = model.generate_super(
+    input_ids,
+    do_sample=False,      # 不使用采样
+    num_beams=1,          # 不使用束搜索
+    temperature=0.0,      # 温度为0确保确定性
+    max_new_tokens=100
+)
+```
+
+### 优缺点分析
+
+**优点**：
+- ✅ **速度快**：每步只需要一次前向传播和简单的argmax操作
+- ✅ **结果确定**：相同的输入总是产生相同的输出
+- ✅ **内存效率高**：不需要维护多个候选序列
+- ✅ **实现简单**：算法逻辑直观易懂
+
+**缺点**：
+- ❌ **容易陷入局部最优**：每步的局部最优不一定等于全局最优
+- ❌ **缺乏多样性**：总是产生相同的序列，缺乏创造性
+- ❌ **可能产生重复内容**：容易陷入重复循环
+- ❌ **忽略长程依赖**：不考虑序列的整体连贯性
+
+### 典型例子
+假设模型生成了以下概率分布：
+
+```
+输入: "今天天气"
+下一token概率:
+- "很" (0.4)
+- "不错" (0.3)
+- "真好" (0.2)
+- "不太好" (0.1)
+```
+
+贪婪解码会选择"很"，生成"今天天气很"，然后继续这个过程。
+
+### 使用场景
+- **确定性任务**：如数学计算、代码生成
+- **需要一致性的应用**：如API服务、自动化脚本
+- **计算资源受限的环境**：需要快速生成结果
+- **基准测试**：作为其他算法的对比基准
+
+## 采样解码（Sampling Decoding）
+
+### 原理说明
+采样解码不是选择概率最大的token，而是基于模型的概率分布进行随机采样。这样可以在每次生成时产生不同的结果，增加文本的多样性和创造性。
+
+**核心思想**：基于概率分布随机选择 → 增加多样性
+
+**数学表达**：
+
+```
+token_t ~ P(token_t | token_1, token_2, ..., token_{t-1})
+```
+
+### 关键参数
+
+#### 1. Temperature（温度）
+- **作用**：控制概率分布的平滑程度
+- **原理**：将logits除以temperature，然后进行softmax
+- **效果**：
+  - `temperature > 1`：分布更平滑，增加随机性
+  - `temperature < 1`：分布更尖锐，更接近贪婪解码
+  - `temperature → 0`：等价于贪婪解码
+
+#### 2. Top-k Sampling
+- **作用**：限制候选token的范围
+- **原理**：只考虑概率最高的k个token，其他token概率设为0
+- **效果**：避免选择概率很低的"奇怪"token，提高质量
+
+### 代码实现
+基于我们实现的 `_random_sample` 方法：
+
+```python
+def _random_sample(self, logits: torch.Tensor, temperature: float = 1.0, top_k: int = None) -> torch.Tensor:
+    """
+    随机采样：基于概率分布随机选择token
+
+    Args:
+        logits: 模型输出的logits，形状为 (batch_size, vocab_size)
+        temperature: 温度参数，控制随机性
+        top_k: 只考虑概率最高的k个token
+
+    Returns:
+        选择的token索引，形状为 (batch_size, 1)
+    """
+    # 1. 温度缩放
+    logits = logits / temperature
+
+    # 2. Top-k过滤
+    if top_k is not None:
+        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+        logits[logits < v[:, [-1]]] = -float('Inf')
+
+    # 3. 计算概率并采样
+    probs = F.softmax(logits, dim=-1)
+    idx_next = torch.multinomial(probs, num_samples=1)
+    return idx_next
+```
+
+**关键步骤解析**：
+1. **温度缩放**：调整概率分布的平滑程度
+2. **Top-k过滤**：移除低概率候选，提高质量
+3. **概率归一化**：使用softmax得到概率分布
+4. **随机采样**：根据概率分布随机选择token
+
+### 使用示例
+```python
+# 基本采样
+output = model.generate_super(
+    input_ids,
+    do_sample=True,         # 启用采样
+    num_beams=1,           # 不使用束搜索
+    temperature=0.8,       # 中等温度
+    max_new_tokens=100
+)
+
+# 带top-k的采样
+output = model.generate_super(
+    input_ids,
+    do_sample=True,
+    num_beams=1,
+    temperature=1.0,       # 较高温度增加随机性
+    top_k=50,             # 只考虑前50个候选
+    max_new_tokens=100
+)
+```
+
+### 温度参数详解
+
+**不同温度的效果对比**：
+
+```python
+# 示例概率分布
+original_probs = [0.6, 0.2, 0.1, 0.05, 0.05]
+
+# Temperature = 0.1 (低温度，接近贪婪)
+scaled_probs = [0.85, 0.08, 0.04, 0.015, 0.015]
+# 结果：很可能选择第一个token
+
+# Temperature = 1.0 (标准温度)
+scaled_probs = [0.6, 0.2, 0.1, 0.05, 0.05]
+# 结果：按原始概率采样
+
+# Temperature = 2.0 (高温度，增加随机性)
+scaled_probs = [0.35, 0.25, 0.18, 0.11, 0.11]
+# 结果：各个token都有机会被选中
+```
+
+### Top-k机制详解
+
+**Top-k过滤过程**：
+
+```python
+# 假设词汇表大小为1000，top_k=50
+logits = [0.1, 2.3, 0.5, 1.8, 0.3, 3.2, 0.9, 0.2, 1.5, 0.7, ...]  # 1000个值
+
+# 步骤1：找到前50个最大值
+v, _ = torch.topk(logits, 50)
+threshold = v[-1]  # 第50大的值
+
+# 步骤2：过滤
+logits[logits < threshold] = -float('Inf')
+# 结果：只有50个token有非零概率，其他950个token概率为0
+```
+
+### 优缺点分析
+
+**优点**：
+- ✅ **多样性好**：每次生成可能产生不同的结果
+- ✅ **创造性高**：能产生意想不到的内容
+- ✅ **避免重复**：不容易陷入重复循环
+- ✅ **可调性强**：通过参数控制随机程度
+
+**缺点**：
+- ❌ **结果不确定**：相同输入可能产生不同输出
+- ❌ **质量不稳定**：可能产生低质量或不连贯的内容
+- ❌ **需要调参**：temperature和top_k需要仔细调节
+- ❌ **计算开销**：需要计算完整的概率分布
+
+### 使用场景
+- **创意写作**：故事生成、诗歌创作
+- **对话系统**：让对话更加自然和有趣
+- **数据增强**：生成多样化的训练数据
+- **探索性任务**：需要探索多种可能性的场景
+
+## 束搜索（Beam Search）
+
+### 原理说明
+束搜索是一种启发式搜索算法，它在每一步生成时保留多个候选序列（束），而不是只选择一个最佳序列。通过维护多条路径，它能够在计算效率和生成质量之间取得平衡。
+
+**核心思想**：维护多条候选路径 → 选择累积概率最高的序列
+
+**算法流程**：
+1. **初始化**：从输入序列开始
+2. **扩展**：为每个候选序列生成多个扩展
+3. **评分**：计算每个新序列的累积概率
+4. **筛选**：保留分数最高的N个候选
+5. **重复**：继续扩展直到结束条件
+
+### 关键概念
+
+#### 束宽度（Beam Width）
+- **定义**：每步保留的候选序列数量
+- **权衡**：
+  - 宽度=1：等价于贪婪解码
+  - 宽度越大：搜索空间越大，质量越高，但计算成本也越大
+
+#### 累积概率
+- **计算方式**：序列概率 = 各个token概率的乘积
+- **数值稳定性**：通常使用对数概率求和
+- **公式**：`log P(sequence) = Σ log P(token_i | context)`
+
+### 代码实现
+基于我们实现的 `_beam_search` 方法：
+
+```python
+def _beam_search(self, idx: torch.Tensor, max_new_tokens: int, num_beams: int,
+                 temperature: float = 1.0, top_k: int = None, stop_id: int = None) -> torch.Tensor:
+    """
+    束搜索：维护多个候选序列，选择最优路径
+
+    Args:
+        idx: 输入序列，形状为 (batch_size, seq_len)
+        max_new_tokens: 最大生成token数量
+        num_beams: 束宽度，表示保留的候选路径数量
+        temperature: 温度参数，控制分布的平滑程度
+        top_k: top-k过滤参数，限制候选token范围
+        stop_id: 停止生成的token ID，遇到则停止
+
+    Returns:
+        生成的token序列，形状为 (batch_size, generated_length)
+    """
+    # 1. 初始化束
+    beams = [idx.clone() for _ in range(num_beams)]
+    beam_scores = torch.zeros(num_beams, device=idx.device)
+    beam_scores[0] = 0.0  # 第一个候选是原始序列
+    beam_scores[1:] = float('-inf')  # 其他候选初始分数为负无穷
+
+    # 2. 主循环：逐步生成token
+    for step in range(max_new_tokens):
+        new_beams = []
+        new_scores = []
+
+        # 3. 扩展每个候选序列
+        for beam_idx, beam in enumerate(beams):
+            if beam_scores[beam_idx] == float('-inf'):
+                continue  # 跳过无效候选
+
+            # 前向传播获取logits
+            output = self(beam)
+            logits = output.logits[:, -1, :]
+
+            # 应用温度和top-k
+            if temperature != 1.0:
+                logits = logits / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+
+            # 计算对数概率
+            log_probs = F.log_softmax(logits, dim=-1)
+
+            # 获取前num_beams个候选token
+            top_log_probs, top_indices = torch.topk(log_probs, k=num_beams, dim=-1)
+
+            # 4. 为当前候选生成多个扩展
+            for k in range(num_beams):
+                token = top_indices[:, k:k+1]
+                log_prob = top_log_probs[:, k]
+
+                new_beam = torch.cat([beam, token], dim=1)
+                new_score = beam_scores[beam_idx] + log_prob.item()
+
+                new_beams.append(new_beam)
+                new_scores.append(new_score)
+
+        # 5. 筛选最佳候选
+        if not new_beams:
+            break
+
+        # 按分数排序，选择前num_beams个
+        sorted_indices = sorted(range(len(new_scores)), key=lambda i: new_scores[i], reverse=True)
+        beams = [new_beams[i] for i in sorted_indices[:num_beams]]
+        beam_scores = [new_scores[i] for i in sorted_indices[:num_beams]]
+
+        # 检查停止条件
+        if stop_id is not None and beams[0][0, -1] == stop_id:
+            break
+
+    # 6. 返回最佳序列
+    return beams[0][:, idx.shape[1]:]  # 只返回生成部分
+```
+
+### 束搜索过程示例
+
+假设束宽度=3，输入="今天天气"：
+
+**第1步扩展**：
+```
+候选1: "今天天气很好" (分数: 0.4)
+候选2: "今天天气不错" (分数: 0.3)
+候选3: "今天天气真好" (分数: 0.2)
+```
+
+**第2步扩展**（每个候选再扩展3个）：
+```
+候选1.1: "今天天气很好啊" (分数: 0.4 + 0.1 = 0.5)
+候选1.2: "今天天气很好。" (分数: 0.4 + 0.2 = 0.6) ← 保留
+候选1.3: "今天天气很好，" (分数: 0.4 + 0.05 = 0.45)
+
+候选2.1: "今天天气不错啊" (分数: 0.3 + 0.15 = 0.45)
+候选2.2: "今天天气不错。" (分数: 0.3 + 0.1 = 0.4) ← 保留
+候选2.3: "今天天气不错，" (分数: 0.3 + 0.08 = 0.38)
+
+候选3.1: "今天天气真好啊" (分数: 0.2 + 0.12 = 0.32)
+候选3.2: "今天天气真好。" (分数: 0.2 + 0.25 = 0.45) ← 保留
+候选3.3: "今天天气真好，" (分数: 0.2 + 0.1 = 0.3)
+```
+
+**筛选结果**（保留分数最高的3个）：
+```
+最佳候选: "今天天气很好。" (分数: 0.6)
+次佳候选: "今天天气不错。" (分数: 0.4)
+第三候选: "今天天气真好。" (分数: 0.45)
+```
+
+### 使用示例
+```python
+# 基本束搜索
+output = model.generate_super(
+    input_ids,
+    do_sample=False,        # 不使用采样
+    num_beams=3,           # 束宽度为3
+    temperature=1.0,       # 标准温度
+    max_new_tokens=100
+)
+
+# 带top-k的束搜索
+output = model.generate_super(
+    input_ids,
+    do_sample=False,
+    num_beams=5,           # 更大的束宽度
+    temperature=0.8,       # 稍微降低温度
+    top_k=50,             # 限制候选范围
+    max_new_tokens=100
+)
+```
+
+### 优缺点分析
+
+**优点**：
+- ✅ **质量较高**：比贪婪解码质量更好
+- ✅ **确定性**：结果相对稳定（相同输入产生相同输出）
+- ✅ **平衡性好**：在质量和效率之间取得平衡
+- ✅ **避免明显错误**：不容易选择明显不合适的token
+
+**缺点**：
+- ❌ **计算开销大**：需要维护多个候选序列
+- ❌ **内存占用高**：存储多个候选序列和分数
+- ❌ **仍可能局部最优**：虽然比贪婪好，但仍可能错过全局最优
+- ❌ **多样性有限**：仍然偏向高概率路径，创造性不如采样
+
+### 束宽度选择建议
+
+| 束宽度 | 适用场景 | 优点 | 缺点 |
+|--------|----------|------|------|
+| 1-2 | 实时应用、计算资源有限 | 速度快、资源占用少 | 质量相对较低 |
+| 3-5 | 一般对话、文本生成 | 质量较好、速度适中 | 资源占用中等 |
+| 6-10 | 高质量生成、翻译 | 质量很高 | 计算开销大 |
+| 10+ | 专业应用、研究 | 最高质量 | 开销很大 |
+
+### 使用场景
+- **机器翻译**：需要准确性和流畅性的平衡
+- **文本摘要**：生成连贯的摘要内容
+- **对话系统**：生成有逻辑的回复
+- **代码生成**：需要语法正确和逻辑合理
+- **长文本生成**：如文章写作、报告生成
+
+## 辅助模型投机解码（Assisted Decoding）
+
+### 原理说明
+投机解码是一种**用小模型加速大模型推理**的技术。它通过"草稿-验证"的方式，让小先生成候选token，然后大家模型快速验证，减少大模型的前向传播次数。
+
+**核心思想**：小模型投机生成 → 大模型批量验证 → 减少大模型计算负担
+
+### 工作流程
+
+#### 1. 草稿生成阶段
+```
+输入: "今天天气"
+小模型快速生成草稿: "今天天气很好，适合出门散步"
+```
+
+#### 2. 验证阶段
+大模型一次性验证整个草稿序列：
+- ✅ 接受的token："今天天气很好，"
+- ❌ 拒绝的token：从"适合"开始拒绝
+- 🔧 大模型重新生成："适合在家休息"
+
+#### 3. 最终结果
+```
+输出: "今天天气很好，适合在家休息"
+```
+
+### 关键优势
+
+**速度提升**：
+- 小模型推理快 → 生成多个候选token
+- 大模型批量验证 → 一次处理多个token
+- 减少90%+的大模型前向传播
+
+**质量保证**：
+- 大模型有最终否决权
+- 只有大模型认可的token才会被保留
+- 不会降低生成质量
+
+### 具体例子对比
+
+**传统方式**（大模型逐个生成）：
+```
+第1步: 大模型 → "今天"
+第2步: 大模型 → "今天天气"
+第3步: 大模型 → "今天天气很"
+第4步: 大模型 → "今天天气很好"
+第5步: 大模型 → "今天天气很好，"
+第6步: 大模型 → "今天天气很好，适合"
+... (每步都需要大模型前向传播)
+```
+
+**投机解码**：
+```
+第1步: 小模型快速草稿 → "今天天气很好，适合出门散步"
+第2步: 大模型批量验证 → 接受"今天天气很好，"，拒绝"适合出门散步"
+第3步: 大模型重新生成 → "适合在家休息"
+```
+
+这样原本需要6次大模型推理的过程，现在只需要2次！
+
+### 技术实现要点
+
+#### 1. 草稿长度控制
+- **草稿不宜过长**：通常2-10个token
+- **接受率平衡**：太长接受率低，太短加速效果不明显
+- **动态调整**：根据接受率调整草稿长度
+
+#### 2. 验证机制
+```python
+# 伪代码
+def assisted_decoding(input_ids, assistant_model, main_model):
+    # 小模型生成草稿
+    draft_tokens = assistant_model.generate_draft(input_ids, max_draft_len=5)
+
+    # 大模型验证
+    accepted_count = main_model.verify_draft(input_ids, draft_tokens)
+
+    # 构建最终结果
+    if accepted_count == len(draft_tokens):
+        return draft_tokens  # 全部接受
+    else:
+        # 部分接受，大模型重新生成剩余部分
+        accepted_part = draft_tokens[:accepted_count]
+        remaining_part = main_model.generate_remaining(input_ids + accepted_part)
+        return accepted_part + remaining_part
+```
+
+### 总结
+投机解码本质上是用**计算资源换时间**，通过小模型的"投机"来减少大模型的计算负担。它是一种聪明的工程优化，在不牺牲质量的前提下显著提升推理速度。

Extra-Chapter/s1-vllm-thinking-budget/readme.md

@@ -0,0 +1,179 @@
+# S1: Thinking Budget with vLLM
+
+首先，我们来了解一下AI教母李飞飞教授关于 Test-time scaling 的论文：[*《s1: Simple test-time scaling》*](http://arxiv.org/abs/2501.19393)
+
+<div align='center'>
+    <img src="./images/image-1.png" alt="alt text" width="50%">
+</div>
+
+论文大致讲了个什么事情呢？简单来说，提出了一种新的测试时间缩放方法，旨在提高模型在推理阶段的效率和准确性。通过调整模型的思考预算，可以在不同的任务和数据集上实现更好的性能。
+
+就是说对于一些复杂问题，需要用推理链来解决的问题，我们可以通过调整模型的思考预算来提高推理效率和准确性。上图也可以看到当思考预算增加时，模型的性能会有明显提升。
+
+<div align='center'>
+    <img src="./images/image-2.png" alt="alt text" width="50%">
+</div>
+
+插一句题外话，论文中判断问题难易程度的方式是通过让 Qwen2.5-32B-Instruct 模型回答问题，答对的问题就是简单问题，答错的就是复杂问题。
+
+<div align='center'>
+    <img src="./images/image-3.png" alt="alt text" width="50%">
+</div>
+
+
+论文也做了消融实验来探讨，在未满足思考预算时插入一些不同的特定词（如：Wait！）对模型最终性能的影响。结果表明，插入特定词可以有效地引导模型进行更深入的思考，并且“Wait，Wait”的效果最好。
+
+## 代码实现
+
+我们使用 vLLM 来实现模型的思考预算。vLLM 是一个高性能的推理引擎，支持大规模语言模型的高效推理。以下为代码实现的步骤：
+
+> 考虑到部分同学配置环境可能会遇到一些问题，我们在 ucloud 平台准备了环境镜像，点击下方链接并直接创建 ucloud 示例即可。 https://www.compshare.cn/images/8gfTTB5y0ql6?referral_code=ELukJdQS3vvCwYIfgsQf2C
+
+<div align='center'>
+    <img src="./images/thinking-budget.png" alt="alt text" width="80%">
+</div>
+
+左侧为不使用思考预算的推理过程，右侧为使用思考预算的推理过程。可以看到，使用思考预算后，模型会在推理过程中插入特定词来引导模型进行更深入的思考。
+
+以下为核心代码实现，完整代码请参考 [*s1.py*](./s1.py)
+
+```python
+def run_thinking_budget_sample(llm_model, tokenizer, user_input, thinking_budget):
+    input_text = build_input(user_input, tokenizer)
+    input_token_count = count_token(input_text, tokenizer)
+
+    iteration_count= 0
+    max_token = input_token_count + thinking_budget
+
+    sampling_params = SamplingParams(
+        temperature=0.7,
+        max_tokens=4096,
+        skip_special_tokens=False
+    )
+
+    think_token_count = 0
+
+    while True:
+
+        wait_sampling_params = SamplingParams(
+            temperature=0.7,
+            max_tokens=thinking_budget - think_token_count,
+            stop='</think>',
+            skip_special_tokens=False
+        )
+
+        outputs = llm_model.generate(
+            input_text,
+            wait_sampling_params
+        )
+        total_token, think_token_count = count_thinking_token(outputs, tokenizer)
+
+        print(f'第{iteration_count}次迭代，思考token数：{think_token_count}')
+
+        if think_token_count > thinking_budget:
+            break
+        input_text = total_token + "\nWait!\n"
+
+        # \nWait a moment. Was there any loophole in my thought just now?!\n
+        # \nWait!\n
+
+        iteration_count += 1
+
+    final_outputs = llm_model.generate(
+        outputs[0].prompt + outputs[0].outputs[0].text + "\n</think>\n",
+        sampling_params
+    )   
+    
+    total_content = final_outputs[0].prompt + final_outputs[0].outputs[0].text
+    thinking_content = total_content.split("<think>")[-1].split("</think>")[0]
+
+    print(total_content)
+
+    print(f"迭代次数：{iteration_count}, 输入token数：{input_token_count}, 思考token数：{count_token(thinking_content, tokenizer)}, 总token数：{count_token(total_content, tokenizer)}")
+```
+
+首先是要定义一个函数 `run_thinking_budget_sample`，该函数接收模型、tokenizer、用户输入和思考预算作为参数。然后构建输入文本并计算输入的 token 数量。
+
+因为`max_tokens` 参数是指生成的最大 token 数量，所以我们需要计算输入文本的 token 数量，并将其与思考预算相加，得到 `max_token = thinking_budget - think_token_count`。如果思考 token 数量超过了思考预算，就停止迭代。另外还需要在 `SamplingParams` 中设置 `stop` 参数为 `</think>`，这样模型在生成文本时会在遇到 `</think>` 时停止生成。
+
+```python
+wait_sampling_params = SamplingParams(
+    temperature=0.7,
+    max_tokens=thinking_budget - think_token_count,
+    stop='</think>',
+    skip_special_tokens=False
+)
+```
+
+另外还需要在每次迭代中，使用 `llm_model.generate` 方法生成文本，并计算思考 token 数量。如果思考 token 数量超过了思考预算，就停止迭代。否则，将生成的文本添加到输入文本中，并在文本末尾添加 `\nWait!\n`，以引导模型进行更深入的思考。
+
+```python
+while True:
+    wait_sampling_params = SamplingParams(
+        temperature=0.7,
+        max_tokens=thinking_budget - think_token_count,
+        stop='</think>',
+        skip_special_tokens=False
+    )
+
+    outputs = llm_model.generate(
+        input_text,
+        wait_sampling_params
+    )
+    total_token, think_token_count = count_thinking_token(outputs, tokenizer)
+
+    print(f'第{iteration_count}次迭代，思考token数：{think_token_count}')
+
+    if think_token_count > thinking_budget:
+        break
+    input_text = total_token + "\nWait!\n"
+
+    # \nWait a moment. Was there any loophole in my thought just now?!\n
+    # \nWait!\n
+
+    iteration_count += 1
+```
+
+当达到思考预算后，使用 `llm_model.generate` 方法生成最终的输出文本，并将其打印出来。最后输出迭代次数、输入 token 数量、思考 token 数量和总 token 数量。
+
+```python
+final_outputs = llm_model.generate(
+    outputs[0].prompt + outputs[0].outputs[0].text + "\n</think>\n",
+    sampling_params
+)   
+
+total_content = final_outputs[0].prompt + final_outputs[0].outputs[0].text
+thinking_content = total_content.split("<think>")[-1].split("</think>")[0]
+
+print(total_content)
+
+print(f"迭代次数：{iteration_count}, 输入token数：{input_token_count}, 思考token数：{count_token(thinking_content, tokenizer)}, 总token数：{count_token(total_content, tokenizer)}")
+```
+
+此时我们还需要另外一个 `SamplingParams` 对象来设置最终生成文本的参数。`max_tokens` 参数设置为 4096，模型根据思考过程进行总结最后得出答案，这个过程也需要很多tokn，此时设置为多少都可以，通常设置为一个较大的值即可。
+
+```python
+sampling_params = SamplingParams(
+    temperature=0.7,
+    max_tokens=4096,
+    skip_special_tokens=False
+)
+```
+
+以上为核心代码实现，完整代码请参考 [*s1.py*](./s1.py)。在实际使用中，可以根据具体的任务和数据集调整思考预算和其他参数，以获得更好的性能。
+
+## 结果分析
+
+使用思考预算后，模型在推理过程中能够更深入地思考问题，从而提高了推理效率和准确性。但是也发现了一些有趣的现象。
+
+例如，在某些情况下，就算插入了`Wait!`，模型并不会按照论文中所示进行多种不同方式尝试解答，或是反思之前的思考过程是否正确。而且还会出现模型在思考过程中重复生成相同的内容，导致思考 token 数量超过思考预算的情况。
+
+<div align='center'>
+    <img src="./images/image-4.png" alt="alt text" width="70%">
+</div>
+
+当然，也有可能本身测试的模型只有14B参数，导致其在思考过程中的能力受到限制。
+
+经过测试下来，有可能强行使用特定词（如：Wait!）来引导模型进行更深入的思考，可能会促使模型产生 “一条道走到黑” 的想法。
+
+部分实验测试记录在 [*output*](./output/) 中。

Extra-Chapter/s1-vllm-thinking-budget/s1.py

@@ -0,0 +1,131 @@
+from vllm import LLM, SamplingParams
+from transformers import AutoTokenizer
+import time
+
+def build_input(prompt, tokenizer):
+    messages = [
+        {"role": "system", "content": "Please reason step by step, and put your final answer within \\boxed{{}}."},
+        {"role": "user", "content": prompt}
+    ]
+    input_text = tokenizer.apply_chat_template(
+        messages, 
+        tokenize=False, 
+        add_generation_prompt=True,
+        enable_thinking=True
+    )
+    return input_text
+
+def count_thinking_token(outputs, tokenizer):
+    total_token = outputs[0].prompt + outputs[0].outputs[0].text
+    thinking_token = total_token.split("<think>\n")[-1]
+    thinking_token_id = tokenizer(thinking_token)["input_ids"]
+    return total_token, len(thinking_token_id)
+
+def count_token(string, tokenizer):
+    return len(tokenizer(string)["input_ids"])
+
+
+def run_thinking_budget_sample(llm_model, tokenizer, user_input, thinking_budget):
+    input_text = build_input(user_input, tokenizer)
+    input_token_count = count_token(input_text, tokenizer)
+
+    iteration_count= 0
+    max_token = input_token_count + thinking_budget
+
+    sampling_params = SamplingParams(
+        temperature=0.7,
+        max_tokens=4096,
+        skip_special_tokens=False
+    )
+
+    think_token_count = 0
+
+    while True:
+
+        wait_sampling_params = SamplingParams(
+            temperature=0.7,
+            max_tokens=thinking_budget - think_token_count,
+            stop='</think>',
+            skip_special_tokens=False
+        )
+
+        outputs = llm_model.generate(
+            input_text,
+            wait_sampling_params
+        )
+        total_token, think_token_count = count_thinking_token(outputs, tokenizer)
+
+        print(f'第{iteration_count}次迭代，思考token数：{think_token_count}')
+
+        if think_token_count > thinking_budget:
+            break
+        input_text = total_token + "\nWait!\n"
+
+        # \nWait a moment. Was there any loophole in my thought just now?!\n
+        # \nWait!\n
+
+        iteration_count += 1
+
+    final_outputs = llm_model.generate(
+        outputs[0].prompt + outputs[0].outputs[0].text + "\n</think>\n",
+        sampling_params
+    )   
+    
+    total_content = final_outputs[0].prompt + final_outputs[0].outputs[0].text
+    thinking_content = total_content.split("<think>")[-1].split("</think>")[0]
+
+    print(total_content)
+
+    print(f"迭代次数：{iteration_count}, 输入token数：{input_token_count}, 思考token数：{count_token(thinking_content, tokenizer)}, 总token数：{count_token(total_content, tokenizer)}")
+
+    # 保存输出到文件
+    with open(f"output_{int(time.time())}.txt", "w") as f:
+        f.write(total_content)
+        f.write(f"\n迭代次数：{iteration_count}, 输入token数：{input_token_count}, 思考token数：{count_token(thinking_content, tokenizer)}, 总token数：{count_token(total_content, tokenizer)}")
+
+
+def run_sample(llm_model, tokenizer, user_input):
+    input_text = build_input(user_input, tokenizer)
+    input_token_count = count_token(input_text, tokenizer)
+
+    sampling_params = SamplingParams(
+        temperature=0.7,
+        max_tokens=32768,
+        skip_special_tokens=False
+    )
+
+    final_outputs = llm_model.generate(
+        input_text,
+        sampling_params
+    )
+
+    total_content = final_outputs[0].prompt + final_outputs[0].outputs[0].text
+    thinking_content = total_content.split("<think>")[-1].split("</think>")[0]
+    print(total_content)
+
+    print(f"输入token数：{input_token_count}, 思考token数：{count_token(thinking_content, tokenizer)}, 总token数：{count_token(total_content, tokenizer)}")
+
+
+if __name__ == "__main__":
+    model_path = "/model/ModelScope/Qwen/Qwen3-14B"
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    llm = LLM(
+        model=model_path, 
+        gpu_memory_utilization=0.9,
+        trust_remote_code=True
+    )
+
+    print("=================================== 思考预算采样 ===================================")
+    run_thinking_budget_sample(
+        llm_model=llm,
+        tokenizer=tokenizer,
+        user_input="There are exactly three positive real numbers $ k $ such that the function\n$ f(x) = \\frac{(x - 18)(x - 72)(x - 98)(x - k)}{x} $\ndefined over the positive real numbers achieves its minimum value at exactly two positive real numbers $ x $. Find the sum of these three values of $ k $.",
+        thinking_budget=32768
+    )
+
+    # print("=================================== 无思考预算采样 ===================================")
+    # run_sample(
+    #     llm_model=llm,
+    #     tokenizer=tokenizer,
+    #     user_input="There are exactly three positive real numbers $ k $ such that the function\n$ f(x) = \\frac{(x - 18)(x - 72)(x - 98)(x - k)}{x} $\ndefined over the positive real numbers achieves its minimum value at exactly two positive real numbers $ x $. Find the sum of these three values of $ k $."
+    # )

Extra-Chapter/transformer-architecture/readme.md

@@ -0,0 +1,344 @@
+# transformer-architecture
+
+当本节内容作为第二章 Transformer架构中2.2章节Encoder-Decoder的补充，结合Pytorch的API源码，从Transformer的整体设计上来解释Transformer的各个模块。
+
+## 从经典架构开始：Encoder-Decoder结构
+让我们先从经典架构图理解Transformer的整体设计思路。Transformer分为两个主要部分：左侧的编码器（Encoder）和右侧的解码器（Decoder）。那么这两块结构的输入和输出分别是什么？
+
+<div align='center'>
+    <img src="./images/trans-img-1.png" alt="alt text" width="90%">
+    <p>1.jpg</p>
+</div>
+
+Encoder的职责是接受完整的源序列输入，将其转换为一个富含语义信息的表示序列。想象一下，如果我们要做机器翻译，Encoder就像是一个深度理解原文的专家，它需要充分理解整个句子的含义、语法结构和上下文关系。
+
+Decoder则承担着更复杂的任务：它需要接受目标序列和编码器输出的表示序列，然后输出词汇/字符的概率分布。这就像是一个翻译专家，既要理解原文的含义（通过Encoder的输出），又要根据已经翻译的部分来决定下一个词应该是什么。
+
+## Positional Encoding：位置编码设计
+
+但这里有一个关键问题需要解决：Transformer本身对位置信息不敏感。比如"我爱你"和"你爱我"这两个句子，在没有位置信息的情况下，模型无法感知到这是语义完全不同的句子。这就像是一个人失去了对词语顺序的感知能力，显然无法正确理解语言。
+
+因此，我们需要一个带有位置信息的向量，将其添加到每个input embedding上，来对不同位置得到不同的表征。这个模块就是图中的**Positional Encoding**。
+
+### 位置编码的设计原则
+
+在设计编码模块时，有**三个重要的前提假设**，这些假设直接影响了最终的实现方案：
+
+**1. 确定性原则**：每个位置的编码应该是确定的数字，不同序列中相同位置的编码应该相同。
+
+为什么这个原则如此重要？让我们考虑一个反例：如果用等分的设计方法，将一个序列从0~1之间做均匀划分，那么序列长度不一样时，每个位置上的编码也就不一样。当序列长度为5时，位置编码可能是0、0.2、0.4、0.6、0.8，但如果序列长度为10，就变成了0、0.1、0.2...。同样对于第二个位置上的字符，在第一个序列中是0.2，在第二个序列中又是0.1，这样的编码就失去了确定性。
+
+**2. 相对关系一致性**：不同句子中，对于任意两个位置之间的相对距离，相对关系应该保持一致。
+
+这个目的是为了学习通用的语言关系，比如："修饰词在被修饰词前1个位置"是通用模式。以下面的长短句举例：
+
+```markdown
+- 长句子（10个词）
+  位置: 0  1  2  3  4  5  6  7  8  9
+  词汇: I  am learning about transformers today in class now
+
+- 短句子（6个词）
+  位置: 0  1  2  3  4  5
+  词汇: I  like deep learning models
+```
+
+在长句子中，位置1和位置4之间的编码关系，应该与短句子中位置2和位置5之间的编码关系完全相同，因为模型需要学会的是通用的相对位置关系。
+
+**3. 泛化能力**：位置序列应该能推广到没见过的更长序列。
+
+第三个假设希望位置编码可以推广到更长的测试句子。假如训练集中序列长度都是10以内的，但测试集中可能会有长度为15的句子，我们希望即使测试集中句子长度更长，在训练中没有见过，我们也能通过这样的position encoding推广过去。
+
+### 三角函数编码
+
+基于这些假设，Transformer采用了sin和cos的组合来表征绝对位置信息：
+
+```markdown
+- 向量维度为偶数：PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
+- 向量维度为奇数：PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
+```
+
+通过sin和cos的组合来表征绝对位置的好处是：`pe(pos+k)`可以写成`pe(k)`的线性组合（利用三角函数公式`sin(A + B) = sin(A)cos(B) + cos(A)sin(B)`）。这样做的意义是：即使测试集中出现了pos+k这种未见过的位置，我们也可以把它写成训练集中见过的位置的线性组合，而不用担心测试集中遇到更长的句子无法推广。
+
+### 位置信息的传递机制
+
+但位置编码是在最底层添加，会不会在深层网络中丢失？这个担心是多余的。通过**残差连接**，位置信息能够充分传递到上层网络。
+
+假设有一个N层的神经网络，输入为x₀（包含位置编码），那么：
+
+```markdown
+- 第1层: x₁ = x₀ + F₁(x₀)
+- 第2层: x₂ = x₁ + F₂(x₁) = x₀ + F₁(x₀) + F₂(x₁)
+- 第3层: x₃ = x₂ + F₃(x₂) = x₀ + F₁(x₀) + F₂(x₁) + F₃(x₂)
+- ...
+- 第N层: xₙ = x₀ + Σᵢ₌₁ⁿ Fᵢ(xᵢ₋₁)
+```
+
+可以看到，初始的位置信息x₀始终存在于每一层的输出中，这确保了位置信息不会随着网络层数的加深而消失。
+
+## 从PyTorch源码API来理解transformer的架构设计
+
+通过查看PyTorch的源码，来了解Encoder和Decoder中的架构实现。源码位于`/pytorch/torch/nn/modules/transformer.py`，此次借助的版本是v2.5.1。
+
+### 顶层Transformer类的设计
+
+首先，PyTorch定义了一个顶层的`Transformer`类，我们可以通过`torch.nn.Transformer`来调用它：
+
+```python
+# 使用示例
+transformer_model = nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6)
+```
+
+<div align='center'>
+    <img src="./images/trans-img-2.png" alt="alt text" width="90%">
+    <p>2.jpg</p>
+</div>
+
+
+在`Transformer`的`__init__`函数中，主要有5个核心参数：
+
+1. **d_model**：整个Transformer的特征维度，在原论文中设置的是512。这个维度需要足够大以承载丰富的语义信息，但太大会导致计算复杂度过高。
+
+2. **nhead**：Multi-head attention的头数目。多头设计的目的是让模型可以捕捉到更多位置与位置之间的关系；
+
+3. **num_encoder_layers**：编码器encoder的block数目，encoder的每个block包含多头自注意力机制和前馈神经网络，这里默认block是6个；
+
+4. **num_decoder_layers**：decoder解码器中block数目，decoder的每个block包含多头自注意力机制、交叉注意力机制以及前馈神经网络；
+
+5. **dim_feedforward**：前馈神经网络层中间的特征维度，默认是2048。Multihead attention输出时，会**首先**映射到2048这个大的特征空间，然后再把它映射回来到512这样的空间。必须要保证输出的维度仍然是512，这样就可以进行残差连接；
+
+<div align='center'>
+    <img src="./images/trans-img-3.png" alt="alt text" width="90%">
+    <p>3.jpg</p>
+</div>
+
+
+### 模块化设计
+
+init函数的作用是实例化模块，第一个要实例化的模块就是encoder。
+
+<div align='center'>
+    <img src="./images/trans-img-4.png" alt="alt text" width="90%">
+    <p>4.jpg</p>
+</div>
+
+encoder通过TransformerEncoder的class去实现实例，在这个class中需要传入encoder_layer，在TransformerEncoderLayer的class中实现了Multihead self attention的调用、残差连接、层归一化、全连接层网络，主要是这些来构成一个encoder_layer。
+
+对于decode部分也是一样，传入decodeLayer参数，这个layer包含了自注意力机制、交叉注意力机制以及前馈神经网络。
+
+<div align='center'>
+    <img src="./images/trans-img-5.png" alt="alt text" width="90%">
+    <p>5.jpg</p>
+</div>
+
+总体上Transformer源码，就是由四个class所构成：
+- **TransformerEncoderLayer**：每一个编码层的实现
+- **TransformerEncoder**：负责把这些编码层串起来
+- **TransformerDecoderLayer**：每一个解码层的实现
+- **TransformerDecoder**：把这些解码层串起来
+
+这种模块化的设计体现了软件工程的最佳实践：单一职责原则和组合优于继承的思想。
+
+### Forward函数：编排计算流程
+
+在`forward`函数中，Transformer的计算流程非常清晰：
+
+<div align='center'>
+    <img src="./images/trans-img-6.png" alt="alt text" width="90%">
+    <p>6.jpg</p>
+</div>
+
+首先encoder输入是source句子以及padding_mask，encoder中的注意力机制不需要掩码，因此mask及is_causal参数不需要传入掩码。但需要对样本长度做掩盖，即**padding_mask**，这个mask表示每一个样本的长度。当我们做训练时，序列长度是不一样的，有些短的样本，在后面的一些位置上就是无效的，通过在softmax中把无效位置上的值转成负无穷，这样经过归一化后概率就变成0，使得在这些位置上，这些没有值的位置变得无效。
+
+decode输入第一个是target，也就是目标句子。第二个是memory，表示encoder输出，因为memory会输送到交叉注意力中。第三个是target mask，这个target mask是一个考虑因果的mask，在数学上是一个上三角矩阵。
+
+### 因果掩码mask
+
+每次预测时，decoder都会有一个输入句子，即：output embedding。但这个output embedding不能全部给它，如果全部给它的话，那它就变成了identity映射，相当于从x到x的关系我给你了，那预测出来的答案肯定是从x到x。
+
+我们需要保证output每次只根据当前要预测的单词的左边的所有单词，去预测这个单词。这个单词的本身和其他的右边的单词都不要输送到Output Embedding。这样的操作需要通过mask来实现，随着预测的字符往右进行，我们给到这个decoder中的output会越来越多，所以它就是一个上三角的矩阵。
+
+第四个memory mask，和source sentence的长度有关，在batch训练中source sentence每一个样本都不太一样，memory mask就是每一个输入源序列这个样本的一个长度。
+
+### Transformer框架
+
+这就是Transformer class总体上的构成框架，init函数去实例化encoder和decoder。在forward函数中基于source、target、source mask、target mask，分别入参到encoder和decoder中，最终得到output，即要预测字符的概率。
+
+<div align='center'>
+    <img src="./images/trans-img-7.png" alt="alt text" width="90%">
+    <p>7.jpg</p>
+</div>
+
+Transformer本质上是一个自回归的解码过程，不是并行的预测输出，而是每次只会预测一个，输出一个，然后不断的进行解码去预测出整体的target sentence。
+
+## Encoder编码器层
+
+接下来我们分别来看下init函数中的4个class，首先是**单个编码器**的实现：在`TransformerEncoderLayer`的`__init__`函数中，需要实例化四个关键组件：
+
+<div align='center'>
+    <img src="./images/trans-img-8.png" alt="alt text" width="90%">
+    <p>8.jpg</p>
+</div>
+
+先是init函数中的参数，和transformer中传入的一致，d_model是整个Transformer的特征维度512。nhead是Multi-head attention中多头自注意力机制中头的数目。
+
+### 多头注意力的设计
+
+为什么是多头？目的是让**模型可以捕捉到更多的位置与位置之间的关系**。多头会分为多组的query、key和value，每一组会单独地去计算attention的上下文的向量，最后把这个向量拼起来，再通过FFN得到最终的一个向量。
+
+这样做的话，对embedding的特征向量的维度会降低，比如说原来的特征向量维度是512，如果我们分为8个头，这时向量每一个头它向量的维度就会变成64（这里不是通过压缩，而是线性变换重组）。每个头独立计算注意力后，会得到8个64维的输出向量。然后通过拼接操作（concatenation），将这8个64维向量首尾相接，重新组合成一个512维的向量。最后再通过一个输出线性变换层，得到最终的512维输出。
+
+<div align='center'>
+    <img src="./images/trans-img-9.png" alt="alt text" width="90%">
+    <p>9.jpg</p>
+</div>
+
+### 前馈神经网络
+
+dimension feed forward是前馈神经网络FFN的维度，因为需要先从512到2048，再从2048到512，所以设定了两个全连接层。前馈神经网络它是对每个单独位置进行一个建模，并且不同位置的参数是共享的。类比1×1的pointwise卷积，对图像中每个像素位置的特征向量独立进行变换。参数共享就是为序列中的每个位置都设计相同的参数，目的是希望模型学会"如何处理特征"的通用能力，而不是"如何处理第x个位置的特征"的特定能力。
+
+FFN实际上做的是embedding相同位置不同维度间的融合，在每个位置内部对该位置的不同特征维度进行融合，注意力机制负责位置间的信息交流。
+
+### Encoder编码器层的组件实例化
+
+init函数中需要去实例化一些实例：
+1. 首先是Multi-head attention，本节着重整个框架；
+2. 实例FFN前馈神经网络中的两个Linear层，第一个Linear比较大，第二个Linear是重新这个投射到d_model的尺度；
+3. 实例layer norm，在self attention之后会经过层归一化，以及在前馈神经网络之后也会经过一个层归一；
+4. 实例两个dropout，dropout是为了使得这个网络具备集成学习的特点，即使我们在训练多个模提高泛化能力；
+
+### Encoder编码器前向传播forward: 编排计算流程
+
+forward函数中，encoder层的调用很简单。Transformer encoder layer的第一部分：通过self attention block得到一个表征(self._sa_block)，然后再加一个这个残差连接(就是和x加起来)，最后再经过一个层归一化。self attention的输入是序列x和pendding-mask，这里的序列x既充当了query，又充当了key和value。
+
+<div align='center'>
+    <img src="./images/trans-img-10.png" alt="alt text" width="90%">
+    <p>10.jpg</p>
+</div>
+
+第二部分是**feed forward block**，把第一部分输出经过每个位置独立的一个全连接层，再进行一个残差连接，输送到层归一化中就得到x。这个就是Transformer encoder中的每个layer的输出。
+
+原始论文的设计是层归一化在后，即else的设计。
+
+### Encoder编码器的串联
+
+`TransformerEncoder` class的作用是将多个编码器层串联起来：将上一层的输出作为下一层的输入，经过多层处理得到最终的编码器输出。
+
+<div align='center'>
+    <img src="./images/trans-img-11.png" alt="alt text" width="90%">
+    <p>11.jpg</p>
+</div>
+
+init主要是传入两个参数，一个是encoder_layer，表示TransformerEncoderLayer class的一个实例。第二个参数是num_layers，表示transformer encoder有多少层，层的含义就是block。encoder中自注意力机制+前馈神经网络这两块是一个block，也就是一层。
+
+## Decoder解码器
+
+解码器的实现比编码器更复杂，因为它包含三个子模块，需要处理更多的交互。
+
+在`TransformerDecoderLayer`中，我们需要实例化三套组件（**自注意力+交叉注意力+前馈神经网络**）：
+
+<div align='center'>
+    <img src="./images/trans-img-12.png" alt="alt text" width="90%">
+    <p>12.jpg</p>
+</div>
+
+在init参数中，第一个是d_model，表示transformer模型特征大小，默认512。第二个参数是nhead，是Transformer decoder的多头自注意力机制的头数。第三个参数是dimension feed forward，表示decoder中FFN前馈神经网络的维度。
+
+<div align='center'>
+    <img src="./images/trans-img-13.png" alt="alt text" width="90%">
+    <p>13.jpg</p>
+</div>
+
+
+### Decoder解码器两种注意力机制的区别
+
+init参数中，decoder和encoder不同的地方，就是需要实例化两个Multihead attention。
+
+第一个Multi-head attention，是**自注意力机制**，它是对decoder这个输入序列的，即：target sentence embedding作为输入序列的自身表征。
+
+第二个Multi-head attention是**交叉注意力机制**，我们想知道decoder multihead attention的输出和encoder输出状态的一个关联性，用该注意力机制跨越了encoder和decoder两个不同序列，不是decoder内部的自我关注，而是让decoder去关注encoder的信息。于是我们通过用decoder MHA（多头注意力）的一个输出作为query，然后用encoder的输出作为key和value来去算出一个上下文表征。
+
+同样Decoder要实现两个Linear层，第一个Linear层是比较大的，把我们交叉自注意力机制的输出，投射到一个更高维的空间，就是2048。然后再把它投射到低维的空间，就是从2048降成512。由于Decoder有三个模块（自注意力+交叉注意力+前馈神经网络），所以这里要实现3个norm和3个dropout。
+
+### Decoder解码器forward：编排流程
+
+解码器的`forward`函数体现了三个模块的协同工作：
+
+1. 第一个模块会把target sentence，也就是序列x和target mask输入到self._sa_block中，对target句子做自注意力机制的计算，结果放入到残差网络中，并且经过层归一化得到输出。
+
+2. 第二个模块依赖于第一个模块输出的x，再和encoder输出的memory做交叉注意力的计算，得到新的表征后，经过残差网络和归一化的norm输出；
+
+3. 把第二个模块的输出输送到FFN前馈神经网络，再次进行残差网络和归一化的norm，得到decoder的输出；
+
+<div align='center'>
+    <img src="./images/trans-img-14.png" alt="alt text" width="90%">
+    <p>14.jpg</p>
+</div>
+
+可以看下_sa_block和_mha_block各自的调用，它们都是调用的是Multihead attention，只不过它们的query、key、value是不一样的。
+
+<div align='center'>
+    <img src="./images/trans-img-15.png" alt="alt text" width="90%">
+    <p>15.jpg</p>
+</div>
+
+<div align='center'>
+    <img src="./images/trans-img-16.png" alt="alt text" width="90%">
+    <p>16.jpg</p>
+</div>
+
+self attention中query、key、value都是目标序列，自身对自身的求相关性的计算。但在交叉注意力机制中，query是decoder的一个输出，key和value是encoder的输出（始终是memory）。
+
+通过多个TransformerDecoderLayer，构成了TransformerDecoder，和TransformerEncoder实现类似，这里不重复赘述。
+
+## 注意力机制的核心计算
+
+最后看下注意力机制的核心计算，PyTorch的实际实现更加复杂和优化，但核心思想可以用论文版本来理解：
+
+<div align='center'>
+    <img src="./images/trans-img-17.png" alt="alt text" width="90%">
+    <p>17.jpg</p>
+</div>
+
+### 注意力机制的直观理解
+
+attention函数就是将一个query和一个由key和value形成的一对元素，建立一个连接，最终得到一个输出。比如：我们去百度搜索一个词条，这个词条就是query，然后百度的数据库里有很多词条信息，每个信息自身都有个key，value就是该词条的具体内容。我们通过这个query，百度就会给我们返回一个搜索结果。这个结果就可以理解为一个注意力机制--基于query和key+value计算出来的一个上下文。
+
+注意力机制的计算结果是Value的一个加权求和的结果，权重是基于Query和Key的相似度计算出来的。先算Query和每个Key的相似度，基于这个相似度进行Softmax归一化得到权重，再把这个权重与每个Key所对应的Value进行加权求和。
+
+### Scaled Dot-Product Attention
+
+在Transformer模型中，用的是"Scaled Dot-Product Attention"，这里有个scaled，可以看到公式中就是QK，会除以一个根号d_k，这个目的就是为了使得Softmax的输入分布会更加稳定一点，也就是使得它的方差会更小一点。
+
+这个Attention由三部分构成，分别是Q、K、V，它们都是三个向量。首先我们会把query和key进行一个矩阵相乘，如果我们单个样本来看就是向量内积，批量来看就是矩阵相乘。
+
+内积过后再除以一个根号d_K，把每个位置上的这个内积放到一起去做一个归一化。这样就可以得到每个位置上的一个概率的表示，因为Softmax它出来的结果就是它总和为1，然后每一个值都是在0到1之间，得到这样一个概率，然后我们把概率和每个位置上value进行一个加权求和，最终能得到attention的一个输出，这就是scaled dot product attention的计算逻辑。
+
+论文中讲的Multi-head self attention其实就是有很多个这样的自注意力机制同时计算，算完之后我们把每一个得到的结果给拼起来，得到了Multi-head self attention最终的输出。
+
+### 注意力计算的代码实现
+
+attention代码如下，这里我们用的是论文实现的简单版本。输入由query、key和value构成。
+
+<div align='center'>
+    <img src="./images/trans-img-18.png" alt="alt text" width="90%">
+    <p>18.jpg</p>
+</div>
+
+首先会把这个q跟k的转置进行一个矩阵相乘，那这样的话就能得到一个一✖️t的向量，把这个向量做一个mask。这里的mask就是把等于0的位置，填充一个非常非常小的一个数，负无穷的数，因为负无穷的数经过Softmax这个归一化之后，它就会变成0的概率，目的是希望那些不重要的位置上的概率赋为0。这里的mask，这里只有一个mask，所以你可以理解为那它这里是一个自注意力机制的一个实现，如果有两个mask，那就是交叉注意力机制的实现。
+
+得到这个p attention就概率分布之后，再把这个p attention跟value进行一个加权求和，得到最终自注意力机制的输出，这是单个自注意力机制的一个计算逻辑，如果是多头的话，最终把单个的输出拼起来就好。
+
+在Transformer模型中，不同的注意力机制有着不同的QKV来源和映射方式：
+
+1. **encoder层的query key和value**：在编码器中都是由word embedding加上position encoding后，通过三个独立的线性映射得到QKV；
+
+2. **在decoder中self attention层**：同样也是通过target sentence embedding+position encoding，通过三个独立的线性映射得到QKV；
+
+3. **在交叉attention中**：query是由decoder的输出经过一个线性映射得到的，key和value是编码器的输出memory分别经过两个映射得到；
+
+
+**参考资料：**
+
+1. [DataWhale HapplyLLM](https://datawhalechina.github.io/happy-llm/#/)
+2. [deep_thought](https://www.bilibili.com/video/BV1o44y1Y7cp/?vd_source=a957a54256d2af7f2a1778751c3855cb&spm_id_from=333.788.videopod.sections)
+3. [HandleNLP](https://nlp.seas.harvard.edu/2018/04/03/attention.html)

Extra-Chapter/vlm-concatenation-finetune/README.md

@@ -0,0 +1,742 @@
+# Qwen3-"VL"——超小中文多模态模型的“拼接微调”之路1（附代码和SwanLab记录）
+
+* 作者：情感机器实验室——陈少宏
+
+* 邮箱：<shaohon_chen@115lab.club>
+
+* GitHub：[https://github.com/ShaohonChen/Qwen3-SmVL](https://github.com/ShaohonChen/Qwen3-SmVL)
+* SwanLab：[https://swanlab.cn/@ShaohonChen/Qwen3-SmVL/overview](https://swanlab.cn/@ShaohonChen/Qwen3-SmVL/overview)
+* 数据集：[https://huggingface.co/datasets/HuggingFaceM4/the_cauldron](https://huggingface.co/datasets/HuggingFaceM4/the_cauldron)
+
+> 💚 **特别感谢**  
+> 感谢 [@zhihuazhao-bit](https://github.com/zhihuazhao-bit) 帮我审阅和修复了提交代码中众多的小 bug，并在 NV 上完成了测试。
+> 感谢 [@KMnO4-zx](https://github.com/KMnO4-zx) 对教程文章内容的审核与修正。
+
+## 摘要
+
+最近Huggingface团队发布了超小多模态模型SmolVLM2，可以做到端侧1GB显存推理。在怀着惊喜试用后发现，虽然模型有极其强大的视觉文本理解能力，但是模型却无法理解中文。这对一个“四六级压线过”的笔者来说十分不友好。刚好前段时间做SwanLab硬件检测适配时有一台未到期的沐曦曦云C500服务器，因此萌生了使用**沐曦GPU芯片**微调、把当前中文小模型扛把子Qwen3与SmolVLM2直接微调拼接的想法。
+
+本教程将介绍一种模型拼接的思路，将SmolVLM2的视觉模块（0.09B）与Qwen3最小的模型（0.6B）进行对齐微调，最终使得Qwen模型具备一定的视觉理解能力。由于笔者时间有限且考虑到文章篇幅的原因，因此该系列预计将以系列的方式放出。篇幅规划如下：
+
+* **第一篇**：如何构建和微调一个拼接模型（**本篇博客**）
+* **第二篇**：模型测评、数据集优化、回答人类对齐
+* **第三篇**：微调技巧介绍、视觉位置编码改动与模型结构优化
+
+<div align="center">
+  <figure>
+  <img src="./images/PPAP.png" alt="PPAP" width="400" />
+  <figcaption>I have a Qwen, I have a SmolVLM...</figcaption>
+  </figure>
+</div>
+
+<div style="background-color:#fff3cd; color:black; padding:10px; border-radius:4px; border:1px solid #fbe5b0; width: 90%; max-width: 100%; margin: auto;">
+  ⚠️关于算力的注意：本教程涉及VLM微调训练，对算力要求较高，需要40G及以上的GPU显存才能运行本教程的训练代码。
+</div>
+
+## 目录
+
+* [SmolVLM2的背景知识](#SmolVLM2的背景知识)
+* [模型拼接和微调思路简介](#模型拼接和微调思路简介)
+* [模型拼接实现和关键代码讲解](#模型拼接实现和关键代码讲解)
+* [微调数据集构建](#微调数据集构建)
+* [微调方法与代码实现](#微调方法与代码实现)
+* [微调训练&结果展示](#微调训练&结果展示)
+* [代码及数据集链接汇总](#代码及数据集链接汇总)
+
+## SmolVLM2的背景知识
+
+首先，我们先回顾一下SmolVLM2模型的构建方案，SmolVLM2模型的整体包括三大块：视觉模型层，特征映射层和大语言模型层，见下图：
+
+<div align="center">
+  <figure>
+  <img src="./images/smolvlm2.png" alt="smolvlm2" width="400" />
+  <figcaption>SmolVLM2的架构图</figcaption>
+  </figure>
+</div>
+
+这个设计是现在比较常见的VLM方案。核心设计思想就是让视觉模型的输出特征与经过embedding的文本特征直接拼接后输入到语言模型（LLM）当中，没有交叉注意力等模块。相比于早期LLaVA等架构，这种最大的优点就是可以最大程度复用已有的语言模型。以Qwen2.5-VL为例，其3B、7B、72B模型大小指的只是LLM部分，并没有包含Vision模块，实际上3B模型的参数量接近4B，视觉模块大概0.4B左右，三个不同大小的VLM使用的是统一的视觉模型。对于一些较大的VLM来说，构建视觉模型时绝大多数的训练都集中在特征映射模块和视觉模块，只在最后阶段为了最终效果进行整体微调时才会调整语言模块。保证了VLM的语言能力。
+
+下面简述一下各个模块的细节：
+
+* 视觉模型层：SmolVLM2-256M版本用的是Google的SigLip模型，一个基于ViT的视觉模型，选用的是最小的SigLip-93M的版本，HF论文里没具体写是直接用的SigLip的参数还是他们从零构建的（有注意到的读者可以评论留言下）。在SmolVLM2代码中对应的是`SmolVLMVisionTransformer`类
+
+* 特征映射层：就是一个简单的MLP，不过SmolVLM2中为了降低图像分辨率还做了一个Pixel shuffle来降低图像分辨率，进一步减少视觉的Token占用，减少了文本长度。HF团队在论文里提到对于参数量较小的VLM来说使用Pixel shuffle还能提升性能。但可训练参数其实就是一个单层的神经网络，这个模块的核心作用就是做特征对齐，将视觉特征从768维（SigLip的维度）映射到576维（SmolLLM2的维度）
+
+* 大语言模型：SmolVLM2-256M模型使用的文本模型是SmolLM-135M版本。可能是由于模型较小，HF团队在论文中说到训练时仅采用两阶段训练：大规模图文训练+针对视频任务的专门微调。为了保障模型的文本能力HF团队在训练数据中参杂了大概14%的纯文本微调数据。不过考虑到视觉模块本身参数量（93M）大小接近于文本模型（135M），因此笔者推测相比于冻结文本模型，数据平衡在这之中会起到更关键的作用。
+
+HF团队在原文中还提到了许多影像小模型VLM性能的trick，感兴趣的读者可以进一步参考SmolVLM2的论文
+
+## 模型拼接和微调思路简介
+
+正所谓顶级食材（模型）只需要最简单的烹饪。模型拼接的思路非常简单直接，基本就三步：
+
+1. 调整SmolVLM2的“上下文控制格式”，使得其与Qwen3兼容。
+
+2. 将模型的文本部分直接从SmolLM2换成Qwen3-0.6B，包括其文本tokenizer和词嵌入、文本模型、以及模型最后输出的语言模型头（LM Head）。
+
+3. 需要重新初始化特征映射层的MLP，从768->576的单层神经网络改成768->1024的单层神经网络即可。
+
+整体架构和对图文对前后处理依旧保持SmolVLM2的流程不变，具体改动见下图：
+
+<div align="center">
+  <figure>
+  <img src="./images/concatation.png" alt="concatation" width="400" />
+  <figcaption>将Qwen3-0.6B替换SmolVLM2的语言模型部分</figcaption>
+  </figure>
+</div>
+
+笔者接下来详细介绍下为了实现“拼接”，具体改动的地方，供之后有类似的任务的读者参考。
+
+## 模型拼接实现和关键代码讲解
+
+### 第一处改动：SmolVLM2的Tokenizers部分
+
+首先需要改动的就是需要改动的是SmolVLM2的Tokenizers部分，这里面主要是涉及两个问题：
+
+* 第一个问题是要将SmolVLM2用于指示图像位置的特殊令牌（Special Token）加入到Qwen3的Tokenizer当中，这么做的目的是防止SmolVLM2的图像Token`<image>`被切分为`<`、`image`、`>`三块。幸运的是，Qwen3本身在Tokenizers中预留了未来用于多模态的特殊特殊令牌`<|image_pad|>`。因此读者直接使用了`<|image_pad|>`代替了`<image>`。用于在文本中预留图像特征的插入点。
+
+* 第二个问题是：SmolVLM2的chat_template和Qwen3的chat_template差别极大。chat_template的作用是通过格式化文本让模型清楚知道不同Token所代表的背景信息。用最近比较流行的话来说就是“上下文工程”（Context Engineering）。
+
+这里我列举了一下Qwen3、SmolVLM2、Qwen2.5-VL在聊天场景下的上下文，供读者参考。
+
+**Qwen3聊天上下文格式**
+
+以给一张图片，问题是“你的名字是什么?”，模型回答是“我的名字是Qwen”为例子。模型的上下文如下：
+
+```txt
+<|im_start|>user
+你的名字是什么?<|im_end|>
+<|im_start|>assistant
+<think>
+
+</think>
+
+我的名字是Qwen<|im_end|>
+
+```
+
+注意Qwen3上下文是没有预留图像位置的，但相比于一般的LLM和VLM多了一个用于插入模型思考过程的`<think><\think>`，以及包含额外的函数调用控制文本。为了便于读者理解，读者在在下面举了一个函数调用的例子。这些函数调用上下文用于控制模型调用外部函数、API或者MCP接口和接收其返回的信息。
+
+考虑到篇幅限制，本文就不粘贴带函数调用、推理、思考等一系列上下文的信息了（笔者打印了下发现实在太长了）。感兴趣的读者可以在Qwen3的官方文处了解详细设计
+
+* [Qwen3函数调用案例](https://qwen.readthedocs.io/zh-cn/latest/framework/function_call.html#the-example-case)
+
+可以说正是这些复杂的上下文信息让模型有可能实现推理、调用函数等多样化的能力。包括多模态理解任务也需要先对上下文进行设计。
+
+**SmdwadwdoVLM2聊天上下文格式：**
+
+以给一张图片，问题是“How many dog in there.”，模型回答是“There are Three dogs.”为例子。三种不同模型的上下文如下：
+
+```txt
+<|im_start|>User:<fake_token_around_image><row_1_col_1><image>...<image><fake_token_around_image><row_1_col_2><image>...<image><fake_token_around_image><row_1_col_3><image>...<image>...<fake_token_around_image><row_4_col_4><image>...<image>
+
+<fake_token_around_image><global-img><image>...<image><fake_token_around_image>How many dog in there.<end_of_utterance>
+Assistant: There are Three dogs.<end_of_utterance>
+Assistant:
+```
+
+看起来非常乱，是因为有大量的`<image>`占位符。`<image>...<image>`之间是许多的`<image>`，笔者为了文章观感删掉了大量的占位符。注意模型的回车、空格均为上下文的一部分，在进行推理时需要严格遵守缩进关系。
+
+但是我们仍能找到熟悉的内容，如`User:`，`Assistant:`等用于提示模型用户的输入与模型应当输出的位置。这些关键词和Qwen类似。
+
+读者注意到了除了`<fake_token_around_image>`，`<image>`等用于指示图像的词，还出现了<row_1_col_1>这种位置指示符，这是因为SmolVLM2为了防止降采样对图像分辨率影响，专门使用了`image splitting`技术，简单来说就是将全局图和高清的局部图共同输入到模型当中（见下图`image splitting`模块），感兴趣的读者可在文末找到HF的技术报告了解详细技术。
+
+<div align="center">
+  <figure>
+  <img src="./images/image-split.png" alt="image-split" width="400" />
+  <figcaption>SmolVLM2的完整推理流程，可以看到在图像输入前使用`image splitting`进行了预切分</figcaption>
+  </figure>
+</div>
+
+**本博文的拼接模型Qwen3-SmVL模型**
+
+相比于Qwen3，SmolVLM2少了很多上下控制的
+
+为了尽可能保存或者说预留Qwen3的思考、函数调用等能力，笔者最终选择将SmolVLM2对于图像特征的排列插入到Qwen3的上下文格式当中。最终上下文格式如下：
+
+```txt
+<|im_start|>user
+<vision_start><row_1_col_1><|image_pad|>（图像插入的地方）<|image_pad|><vision_start>
+（用户提问的地方）
+<|im_end|>
+<|im_start|>assistant
+<think>
+
+</think>
+
+（模型回答的地方）<|im_end|>
+<|endoftext|>
+```
+
+可以看到读者尽量保持了与Qwen3的风格和复用特殊令牌。这样能够使得后续拼接的Qwen3-0.6B模型不至于受到上下文差异过大带来的性能损耗。实际上在设计微调上下文时应尽量与模型先前训练的任务接近，以减少微调带来的性能损失。
+
+transformers实现模型上下文格式控制的代码并非python语言，而是一种前端文本格式控制的语言Jinja。这个语言的变量作用域设计简直可以说是有魔法在里面。配合上Qwen3功能丰富且复杂的上下文策略，让笔者花了2个小时用于修改chat_teamplate。这里笔者不赘述如何修改chat_template，感兴趣的读者可以去文末代码链接寻找`chat_template.jinja`文件，笔者专门将chat_template模版拿出来，并且做了格式化方便读者阅读。未来有时间了笔者专门写一篇模型上下文控制与jinja语言的博客。
+
+### 第二处改动：替换SmolVLM2的SmolLM2模型为Qwen3-0.6B
+
+替换模型这块没什么复杂的，主要是需要处理Transformers比较复杂的嵌套逻辑。Tranformers通常建议模型将预训练模型backbone和下游任务分开来。改动逻辑图如下：
+
+<div align="center">
+  <figure>
+  <img src="./images/change_model.png" alt="change_model" width="400" />
+  <figcaption>替换smolvlm2的文本模块和语言模型头</figcaption>
+  </figure>
+</div>
+
+以Qwen3为例，预训练Backbone模型为`Qwen3Model`，仅仅包含embedding层、各个Decoder层，最后输出的是所有输入token的hidden state。负责下游任务的Qwen3提供了包括：用于因果语言序列生成的`Qwen3ForCausalLM`，也就是大家常用的语言生成。负责句子分类`Qwen3ForSequenceClassification`，使用最后一个生成的token输入到一个单层MLP做序列级分类，做句子情绪分类等可以用这个下游模型；`Qwen3ForTokenClassification`用于做Token级分类，比如语言实体抽取任务可以使用这个下游模型。`Qwen3ForQuestionAnswering`则是专门做抽取式问答任务的模型，核心思想是输入（问题，参考文本）让模型从参考文本中找到与问题最相关的一段，这类任务由于RAG系统的出现没那么流行了，未来笔者专门出一个系列的教程阐述除了因果语言序列生成以外的任务则怎么微调。
+
+**关键代码如下**
+
+```python
+from transformers import (
+    AutoProcessor,
+    AutoModelForImageTextToText,
+    AutoTokenizer,
+    AutoModelForCausalLM
+)
+
+# 替换text模型和head
+smolvlm2_02B_model = AutoModelForImageTextToText.from_pretrained(
+    "model/SmolVLM2-256M-Video-Instruct",
+    torch_dtype=torch.bfloat16,
+    _attn_implementation="eager",
+).to(device)
+
+qwen3_06b_model = AutoModelForCausalLM.from_pretrained(
+    "model/Qwen3-0.6B", torch_dtype=torch.bfloat16
+).to(device)
+
+smolvlm2_02B_model.model.text_model = qwen3_06b_model.model
+smolvlm2_02B_model.lm_head = qwen3_06b_model.lm_head
+...
+```
+
+接下来比较复杂的是替换所有的关键变量，比如模型内用于在文本序列中为图像特征预留的占位符`image_token_id`，用于指示停止生成的`eos_token_id`，和计算loss值会用到的`vocab_size`，Qwen的词表大小为151936，远远大过SmolVLM2的词表49280。具体代码如下：
+
+```python
+...
+# 替换词表大小
+smolvlm2_02B_model.vocab_size = qwen3_06b_model.vocab_size
+smolvlm2_02B_model.model.vocab_size = qwen3_06b_model.vocab_size
+smolvlm2_02B_model.config.vocab_size = qwen3_06b_model.vocab_size
+smolvlm2_02B_model.config.text_config.vocab_size = qwen3_06b_model.vocab_size
+smolvlm2_02B_model.model.config.vocab_siz = qwen3_06b_model.vocab_size
+smolvlm2_02B_model.model.config.text_config.vocab_size = qwen3_06b_model.vocab_size
+# 替换图像token
+smolvlm2_02B_model.image_token_id = 151655
+smolvlm2_02B_model.model.image_token_id = 151655
+smolvlm2_02B_model.config.image_token_id = 151655
+smolvlm2_02B_model.model.config.image_token_id = 151655
+# 替换模型生成停止符
+smolvlm2_02B_model.generation_config.eos_token_id = 151645
+···
+```
+
+上面的代码可以看到在替换各个变量时需要将嵌套模型的变量一起替换掉，笔者之前训练时就因为仅仅替换了`SmolVLMForConditionalGeneration`而忘记替换`SmolVLMModel`中的`image_token_id`，导致语言模型接收不到图像特征，最后表现出来就是loss下降的极快且低，grad_norm看起来也学到位了，一推理效果特别差，附上错误训练的损失图：
+
+<div align="center">
+  <figure>
+  <img src="./images/fail_train.png" alt="fail_train" width="800" />
+  <figcaption>SwanLab记录训练结果展示：蓝色为错误训练的完整微调loss图，可以看到损失下降很快，然而实际推理会发现模型并没有图像理解能力。冻结语言模型头（红色）后发现grad_norm为零且loss不收敛，正确的应该是黄色</figcaption>
+  </figure>
+</div>
+
+笔者最早没发现改动错误，先做完整微调（蓝色曲线）后发现损失下降很快达到了0.1以下，结果实际一推理发现模型完全没有图像理解能力，就补了一个冻结语言模型只微调视觉模型的实验（红色曲线），结果发现损失完全没下降，才定位到了视觉特征传入有问题。后续修复后正确的损失下降过程见黄色图像。
+
+### 第三处改动：构建和替换特征映射层
+
+这个相对较简单，只需要重新构建一个维度对齐的`SmolVLMConnector`即可。Qwen3的hidden_dim是1024，SigLip的hidden_dim是768，因此构建一个768➡️1024映射的`SmolVLMConnector`即可。代码如下：
+
+```python
+···
+# 构建配置并且创建连接器
+@dataclass
+class VisionConfig:
+    hidden_size: int = 768
+
+@dataclass
+class TextConfig:
+    hidden_size: int = 1024
+
+@dataclass
+class ConnectConfig:
+    scale_factor: int = 4
+    vision_config: VisionConfig = VisionConfig()
+    text_config: TextConfig = TextConfig()
+
+new_connector_config = ConnectConfig()
+
+# 替换 SigLit 到 LLM 的 connector 层
+new_connector = SmolVLMConnector(new_connector_config).to(device).to(torch.bfloat16)
+smolvlm2_02B_model.model.connector = new_connector
+···
+```
+
+## 微调数据集构建
+
+笔者最初计划寻找中文多模态数据集，但发现相关的资料比较少。因此决定先用英文的多模态数据集凑合一下。之后再考虑通过数据合成的方式将部分数据翻译为中文。关于数据合成和配比的问题将在之后的博客讨论。
+
+<div align="center">
+  <figure>
+  <img src="./images/the_cauldron.png" alt="the_cauldron" width="400" />
+  <figcaption>the_cauldron数据集logo</figcaption>
+  </figure>
+</div>
+
+这里为了方便本项目直接使用HuggingFace团队整合的多模态数据集the Cauldron数据集，Cauldron翻译成中文类似于煮东西的“釜”，不知道HF团队是不是玩“炼丹”的梗。这个数据集整合了50个视觉微调任务数据集的训练集，用于微调Huggingface发布的多模态模型Idefics2模型。这50多个数据集都被处理成了一致的格式（见下图），共有1,880,992条数据，完整下载约169G，非常方便使用。
+
+<div align="center">
+  <figure>
+  <img src="./images/data_show.png" alt="data_show" width="800" />
+  <figcaption>数据集样本展示</figcaption>
+  </figure>
+</div>
+
+不过可惜数据集的文本都是英文内容，且绝大多数数据集的回复非常短，只有一个词，这也给后面模型训练带来了麻烦。本篇博客暂时不讨论关于数据构建和配比的问题，后续有时间了专门做相关的实验。本博客先以为Qwen3模型带来视觉能力为核心目标。
+
+数据集的下载链接如下，国内推荐用modelscope下载：
+
+* [HuggingFace Hub](https://huggingface.co/datasets/HuggingFaceM4/the_cauldron)
+* [ModelScope](https://modelscope.cn/datasets/AI-ModelScope/the_cauldron)
+
+笔者在实际测试时发现"mimic_cgd"，"localized_narratives"，"okvqa"，"ocrvqa"，"clevr_math"这几个子数据集加载有点异常，建议使用此数据集训练的读者手动处理下，社区也有用户反馈这几个数据可以在原始来源处额外下载，未来笔者将会补全这几个数据集重新上传一次完整版的the Cauldron数据集。
+
+## 微调方法与代码实现
+
+### 冻结模型参数微调
+
+整体微调方法采用了CLM模型通常的Teacher Forcing的学习方法，损失就是标准的交叉熵损失。考虑到此次本教程的目标是先确保模型具备中文多模态能力（优化模型性能等之后撰写其他博客），因此为了实验效率，在对齐微调阶段**采用冻结视觉模型与文本模型，仅微调特征映射器和语言模型头**的方法。
+
+冻结模型参数的核心代码如下：
+
+```python
+def freeze_model(qwen_smvl):
+    for _, param in qwen_smvl.model.text_model.named_parameters():
+        param.requires_grad = False
+    for _, param in qwen_smvl.model.vision_model.named_parameters():
+        param.requires_grad = False
+    return qwen_smvl
+```
+
+冻结后训练参数、模型总参数、与占比如下：
+
+```txt
+trainable params: 12.00M || all params: 662.87M || trainable%: 1.81
+```
+
+### 文本长度，损失掩码和截断策略
+
+**文本长度**
+
+由于视觉特征需要占据大量的文本长度，笔者简单测试了下the_cauldron图像占0.8K到1.3K左右的token。而数据集中大多数文本token数在200-500左右，极少情况会有3-4K的情况。因此笔者统一采用2K的文本长度，超出部分截断处理。
+
+这里有一个不同于文本微调的细节要注意，文本截断长度不能小于图像token，否则会导致模型在进行特征拼接时报错（当然图像特征如果被截断了，这条训练数据也就没意义了）。因此对于显存不足64G的同学如果需要适当缩短文本长度（不建议低于1.5K），最好连同图像分辨率也缩小些。在后面的博客我们会专门增加对减少图片token占用的研究。
+
+同样由于文本长度受限，且图像特征没法截断，我们也没使用“packing dataset”的方法提升模型的训练效率。
+
+考虑到部分数据集存在多张图片的情况，考虑到本次训练仅采用2k的文本长度（与之对比HF在训练SmolVLM-256M版本采用的是8K的文本长度，2.2B版使用了16K的文本长度）。针对单条数据中存在多张图片的情况仅仅选用第一张。
+
+**损失掩码**
+
+在采用Teacher Forcing的学习方法时，文本微调中损失掩码有两种策略：
+
+* 对包含“用户问题”和“模型回复”的完整文本进行微调优化
+* 仅对“模型回复”部分进行微调优化
+
+这两种策略的对比如下图：
+
+<div align="center">
+  <figure>
+  <img src="./images/mask.png" alt="mask" width="800" />
+  <figcaption>两种微调掩码策略的差异，通常建议选择“仅微调模型回答部分”以增强泛化性</figcaption>
+  </figure>
+</div>
+
+通常来说使用“仅微调模型回复部分”的策略模型更容易泛化（这点与HF在SmolVLM2的论文提到的trick）。然而笔者为了提高训练效率选择了完整文本微调。可以在后续博客中增加消融实验做进一步对比。
+
+值得注意的是，在进行完整文本微调时，需要单独屏蔽Image Token以防止对图像占位token计算损失，影响模型表现。
+
+**关键代码如下：**
+
+```python
+def data_collate_fix2k(examples, processor, device, max_length=2048):
+    batch_text = []
+    batch_image = []
+    for example in examples:
+        images = example["images"][:1]  # 只允许一张图，不然显存压力太大
+        batch_image.append(images)
+        image_num = len(images)
+        chat_texts = example["texts"][0]
+        messages = [
+            {
+                "role": "user",
+                "content": [{"type": "image"}] * image_num
+                + [{"type": "text", "text": chat_texts["user"]}],
+            },
+            {
+                "role": "assistant",
+                "content": [{"type": "text", "text": chat_texts["assistant"]}],
+            },
+        ]
+        text = processor.apply_chat_template(
+            messages, enable_thinking=False, add_generation_prompt=False
+        )
+
+        batch_text.append(text)
+
+    batch = processor(
+        text=batch_text,
+        images=batch_image,
+        max_length=max_length,
+        return_tensors="pt",
+        padding="max_length",
+        truncation=True,
+    )
+    labels = batch["input_ids"].clone()
+    labels[labels == processor.tokenizer.pad_token_id] = -100
+    labels[labels == processor.image_token_id] = -100
+    batch["labels"] = labels
+    return batch.to(device, dtype=torch.bfloat16)
+```
+
+### 微调超参数设置
+
+**学习率**
+
+由于仅仅针对特征映射层（connector）进行训练，且conntector由于要对齐Qwen3的维度因此参数为随机初始化（理论上可以采用一些独特的初始化策略提升性能，但考虑到模型较小因此笔者没关注初始化策略）。因此学习率设置为lora中较为流行的1e-4学习率策略。
+
+为了保障有效收敛，学习率衰减基本是必备的trick，采用的是社区比较流行的cosine学习率衰减，衰减至0。warm up为整体步长的10%（在超过1000k step的情况下固定为50）。
+
+**batch size**
+
+Batch size通常来说越大越好，然而由于VLM模型的文本长度太大，因此采用每卡1 batch和4梯度累加（grad accelerate），在8卡训练中等效32 Batch size。
+
+**训练参数设置代码**
+
+```python
+training_args = TrainingArguments(
+    seed=42,
+    data_seed=42,
+    max_steps=200,
+    # num_train_epochs=1,  # 训练1个epoch 约1k steps
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=4,
+    dataloader_pin_memory=False,
+    warmup_ratio=0.1,
+    learning_rate=1e-4,
+    lr_scheduler_type="cosine",
+    weight_decay=0.01,
+    logging_steps=5,
+    eval_strategy="steps",
+    eval_steps=0.125,
+    save_strategy="steps",
+    save_steps=0.125,
+    save_total_limit=8,
+    optim="adamw_torch",
+    bf16=True,
+    output_dir=f"./model/freeze_except_connector_cocovqa",
+    overwrite_output_dir=False,
+    report_to="swanlab",
+    run_name="freeze_except_connector_cocovqa",
+    remove_unused_columns=False,
+    gradient_checkpointing=False,
+)
+```
+
+### 训练环境
+
+微调代码基于沐曦的C500国产通用计算GPU实现，显存为64G。沐曦的AI芯片基本完全兼容pytorch和huggingface transformers场景，并且在做多模态训练时相比较其他国产AI芯片罕见的没有兼容性问题。读者在尝试本项目代码时可以采用Nvidia显存40G以上的显卡运行本教程。
+
+**笔者个人感觉沐曦的GPU整体适配效果还是非常好的，没遇到适配性的问题。体验上和用NV的GPU做训练没什么区别**。笔者自己也用过好几款国产GPU，沐曦的体验肯定是名列前茅的，包括代码中有指定flash attention在沐曦GPU上都能成功迁移，这点非常值得给沐曦团队点个赞。希望国产GPU生态能越发展越好，造福广大炼丹师；）。
+
+<div align="center">
+  <figure>
+  <img src="./images/muxi-gpu.jpg" alt="muxi-gpu" width="400" />
+  <figcaption>沐曦国产GPU，笔者用的云端服务器没见过真机，因此找了张网图</figcaption>
+  </figure>
+</div>
+
+训练环境的话除了安装GPU对应的驱动和pytorch外，本教程需要额外安装Huggingface全家桶，如下：
+
+```txt
+torch   # 推荐版本>=6.0
+torchvision
+transformers>=4.53.0
+accelerate
+datasets
+num2words   # SmolVLM2需要
+```
+
+额外补充一句，如果采用沐曦GPU训练的话，需要在沐曦官方文档处寻找[沐曦版torch](https://developer.metax-tech.com/softnova/index)的安装方式进行下载。其他HF环境和NV基本一样。附赠一个沐曦查看GPU的命令：
+
+```bash
+mx-smi
+```
+
+效果如下：
+
+```bash
+=================== MetaX System Management Interface Log ===================
+Timestamp                                         : Sat Jul 12 14:58:51 2025
+
+Attached GPUs                                     : 8
++---------------------------------------------------------------------------------+
+| MX-SMI 2.1.12                       Kernel Mode Driver Version: 2.12.13         |
+| MACA Version: 2.29.0.19             BIOS Version: 1.22.3.0                      |
+|------------------------------------+---------------------+----------------------+
+| GPU         NAME                   | Bus-id              | GPU-Util             |
+| Temp        Pwr:Usage/Cap          | Memory-Usage        |                      |
+|====================================+=====================+======================|
+| 0           MetaX C500             | 0000:0e:00.0        | 0%                   |
+| 36C         69W / 350W             | 5680/65536 MiB      |                      |
++------------------------------------+---------------------+----------------------+
+| 1           MetaX C500             | 0000:0f:00.0        | 0%                   |
+| 38C         70W / 350W             | 4986/65536 MiB      |                      |
++------------------------------------+---------------------+----------------------+
+| 2           MetaX C500             | 0000:10:00.0        | 0%                   |
+| 37C         69W / 350W             | 4986/65536 MiB      |                      |
++------------------------------------+---------------------+----------------------+
+| 3           MetaX C500             | 0000:12:00.0        | 1%                   |
+| 37C         71W / 350W             | 4986/65536 MiB      |                      |
++------------------------------------+---------------------+----------------------+
+| 4           MetaX C500             | 0000:35:00.0        | 0%                   |
+| 37C         70W / 350W             | 4986/65536 MiB      |                      |
++------------------------------------+---------------------+----------------------+
+| 5           MetaX C500             | 0000:36:00.0        | 1%                   |
+| 36C         68W / 350W             | 4986/65536 MiB      |                      |
++------------------------------------+---------------------+----------------------+
+| 6           MetaX C500             | 0000:37:00.0        | 0%                   |
+| 39C         73W / 350W             | 4986/65536 MiB      |                      |
++------------------------------------+---------------------+----------------------+
+| 7           MetaX C500             | 0000:38:00.0        | 0%                   |
+| 38C         71W / 350W             | 4986/65536 MiB      |                      |
++------------------------------------+---------------------+----------------------+
+
++---------------------------------------------------------------------------------+
+| Process:                                                                        |
+|  GPU                    PID         Process Name                 GPU Memory     |
+|                                                                  Usage(MiB)     |
+|=================================================================================|
+|  0                  3496691         python3.10                   4066           |
+|  0                  3496692         python3.10                   102            |
+|  0                  3496693         python3.10                   102            |
+|  0                  3496694         python3.10                   102            |
+|  0                  3496695         python3.10                   102            |
+|  0                  3496696         python3.10                   102            |
+|  0                  3496697         python3.10                   102            |
+|  0                  3496698         python3.10                   170            |
+|  1                  3496692         python3.10                   4154           |
+|  2                  3496693         python3.10                   4154           |
+|  3                  3496694         python3.10                   4154           |
+|  4                  3496695         python3.10                   4154           |
+|  5                  3496696         python3.10                   4154           |
+|  6                  3496697         python3.10                   4154           |
+|  7                  3496698         python3.10                   4154           |
++---------------------------------------------------------------------------------+
+```
+
+### 训练代码实现
+
+在构建训练代码时，笔者使用HuggingFace Transfomers框架的Trainer类来完成训练代码。Trainer类实现的训练逻辑基本能完成大部分微调任务。这里唯一需要提到的是笔者使用了Qwen3-0.6B而非通常此类任务该使用的Qwen3-0.6B-Base模型，Qwen3-0.6B相比于Qwen3-0.6B-Base模型经过了指令遵从微调、对齐等，能实现聊天问答功能。
+
+通常来说对经过微调的模型进行持续训练会一定程度带来性能损失，然而此次微调时笔者冻结了LLM参数，因此需要选用经过微调的模型来实现多模态问答能力。
+
+笔者在训练过程中使用的是bfloat16精度，相比于float16来说bfloat16增加了尾数位数，训练过程中精度会更高些。
+
+在前期进行方案验证阶段笔者采用的是cocoqa数据集，并且进行200steps的微调训练。在确定方案可行后笔者计划使用完整数据集进行微调训练，然而考虑到训练数据量仅仅只有整个模型的12M，因此笔者按参数量与训练Token的比值为1:10采样数据集，即总共从数据集中采样出60K条数据用于实际训练（文本长度按照2k计算，实际上有padding部分因此实际参与token数小于120M）。笔者认为参与训练的数量是足以令模型收敛的，后续实验也证明了模型确实能达到我们所期望的效果。
+
+**训练关键代码实现**
+
+代码比较长是因为增加了断点续训的能力
+
+```python
+################
+# 开启训练
+################
+last_checkpoint = None  # load last checkpoint if available
+if (
+    os.path.isdir(training_args.output_dir)
+    and not training_args.overwrite_output_dir
+):
+    last_checkpoint = get_last_checkpoint(training_args.output_dir)
+    if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+        raise ValueError(
+            f"Output directory ({training_args.output_dir}) already exists"
+        )
+    print(
+        f"Checkpoint detected, resuming training at {last_checkpoint}."
+    )
+# Init Trainer
+trainer = Trainer(
+    model=qwen_smvl,
+    args=training_args,
+    train_dataset=raw_data["train"],
+    eval_dataset=raw_data["test"],
+    data_collator=collate_fn,
+)
+trainer.train(resume_from_checkpoint=last_checkpoint)
+qwen_smvl.save_pretrained(training_args.output_dir)
+```
+
+完整代码见[代码及数据集链接汇总](#代码及数据集链接汇总)
+
+或者直接由[完整项目GitHub地址]()
+
+## 微调训练&结果展示
+
+### 环境安装与微调代码执行
+
+**代码准备与环境安装**
+
+可以在[GitHub仓库地址](https://github.com/ShaohonChen/Qwen3-SmVL)处找到实验的完整代码。使用git clone后使用如下命令安装环境
+
+```bash
+pip install -r requirements.txt
+```
+
+**数据集和模型下载**
+
+笔者附上自动下载脚本，注意该脚本使用[魔塔社区](https://modelscope.cn/)完成模型与数据集的下载
+
+```bash
+bash download_resource.sh
+```
+
+### 小批量微调训练
+
+为了进行快速验证，笔者首先使用cocoqa数据集并且进行了200steps的训练，所有参数与前文所述一致。通过
+
+运行实验命令如下，推荐使用8卡进行训练，在8张沐曦GPU卡上预计需要使用20min
+
+```bash
+# 单GPU训练
+CUDA_VISIBLE_DEVICES=0 python train.py ./cocoqa_train.yaml
+# 8GPU训练
+accelerate launch --num_process 8 train.py ./cocoqa_train.yaml
+```
+
+注意，本项目使用SwanLab进行训练日志记录与分析，如果未登陆SwanLab需要使用`swanlab login`进行登陆。运行后看到如下结果即代表实验成功开启：
+
+<div align="center">
+  <figure>
+  <img src="./images/run.png" alt="run" width="800" />
+  <figcaption>成功训练后可以看到SwanLab链接</figcaption>
+  </figure>
+</div>
+
+下面是笔者完成小批量微调训练的训练损失、测试损失结果图
+
+<div align="center">
+  <figure>
+  <img src="./images/cocoqa_swanlab.png" alt="cocoqa_swanlab" width="800" />
+  <figcaption>SwanLab训练可视化分析结果，可以看到最后训练损失和测试损失都收敛在0.65左右</figcaption>
+  </figure>
+</div>
+
+模型在完成训练后会自动使用一张狗狗图片配合问题“图中有什么动物？”让模型根据图片进行推理，推理结果如下：
+
+<div align="center">
+  <figure>
+  <img src="./images/bad_case.png" alt="bad_case" width="800" />
+  <figcaption>SwanLab记录了模型训练好后的推理结果，可以看到模型能正常理解和回复中文</figcaption>
+  </figure>
+</div>
+
+当时看到模型对着三只狗的图片回答“兔子”时笔者一时认为炼丹失败了，当然如果实际炼丹失败后模型是不会输出动物类型的，而是输出一些乱码或者告诉用户并没有看到图片。识别错误的原因实际上是由于训练步数过少导致的。后续加大训练步数与数据量后模型能正常识别出狗狗并且能准确的说出有三只狗。
+
+<div align="center">
+  <figure>
+  <img src="./images/dog.png" alt="dog" width="250" />
+  <figcaption>附上三只眼神忧伤的狗子，难道长得很像兔子吗？</figcaption>
+  </figure>
+</div>
+
+PS: 作者公开了在[SwanLab上的训练结果](https://swanlab.cn/@ShaohonChen/Qwen3-SmVL/overview)，感兴趣的读者可以自己查看，SwanLab也支持Clone作者的训练日志，大家可以在自己训练时clone笔者的项目去做对照。
+
+### 完整微调训练结果展示
+
+运行实验命令如下，推荐使用8卡进行训练，在8片沐曦C500芯片上预计需要使用1.5h
+
+```bash
+# 单GPU训练
+CUDA_VISIBLE_DEVICES=0 python train.py ./full_train.yaml
+# 8GPU训练
+accelerate launch --num_processes 8 train.py ./full_train.yaml
+```
+
+下图展示了使用完整微调数据对比于小批量训练，可以看到全量数据微调时loss变得更为抖动，这是由于数据类型的丰富给模型的学习带来了一定的挑战。
+
+<div align="center">
+  <figure>
+  <img src="./images/fulldata_swanlab.png" alt="fulldata_swanlab" width="800" />
+  <figcaption>红色为完整训练loss，黄色为小批量训练结果</figcaption>
+  </figure>
+</div>
+
+进一步对比完整训练和小批量训练的训练和测试损失，可以看到完整训练的模型训练损失达到了0.61，远低于仅仅使用cocoqa模型的效果，评估损失也远低于前者，维持在0.58左右。
+
+<div align="center">
+  <figure>
+  <img src="./images/evalloss.png" alt="evalloss" width="800" />
+  <figcaption>红色为完整训练loss，黄色为小批量训练结果</figcaption>
+  </figure>
+</div>
+
+这里值得一提的是，由于我们选用的测试集比较小（仅有64条数据），因此训练损失和测试损失的差距并不能直接理解为过拟合的证据。实际上在大模型训练上，如果数据集足够大的情况下，通常可以认为训练损失等同于评估损失。
+
+此外，模型通过分析1k步之后的训练损失、平均梯度范数（Grad Norm）变化。此时训练任务已过半，且学习率开始快速衰减。如下图，可以看到学习率快速衰减的情况下模型损失并没有明显的进一步下降，这说明模型已经实现了充分训练。
+
+<div align="center">
+  <figure>
+  <img src="./images/1kstep.png" alt="1kstep" width="800" />
+  <figcaption>1k step之后模型的训练损失变化</figcaption>
+  </figure>
+</div>
+
+在训练效率方面，可以看到我们仍没有充分榨干沐曦GPU的性能，当然这也是由于多模态任务的网络本身架构上比较复杂，其中包含许多对图像、文本的拼接工作，这也导致了GPU性能没法完全利用。
+
+<div align="center">
+  <figure>
+  <img src="./images/mx-gpu-use.png" alt="mx-gpu-use" width="800" />
+  <figcaption>SwanLab对沐曦C500训效率自动记录</figcaption>
+  </figure>
+</div>
+
+同样在完成训练后使用狗狗图进行了测试，这次模型能理解图片、中文以及给出正确的回复。更为关键的是模型完全保留了Qwen3-0.6B原有的全部能力，包括函数调用、推理等。在此基础上，仅仅增加了0.09B参数量的情况下为模型带来了图像理解能力！
+
+<div align="center">
+  <figure>
+  <img src="./images/good_case.png" alt="good_case" width="800" />
+  <figcaption>同样的图片与问题，更大的数据量和更充足的数据使得模型能够正确给出回复</figcaption>
+  </figure>
+</div>
+
+### 模型推理与效果分析
+
+等笔者下完数据集后未来补一下测试环节 ; ）
+
+可以关注[swanlab教程集合](https://docs.swanlab.cn/examples/qwen3_smolvlm_muxi.html)获取最新更新教程！
+
+## 代码及数据集链接汇总
+
+微调用The Cauldron数据集下载链接：
+
+* HuggingFace Hub: [https://huggingface.co/datasets/HuggingFaceM4/the_cauldron](https://huggingface.co/datasets/HuggingFaceM4/the_cauldron)
+* ModelScope: [https://modelscope.cn/datasets/AI-ModelScope/the_cauldron](https://modelscope.cn/datasets/AI-ModelScope/the_cauldron)
+
+Qwen3-0.6B模型下载：
+
+* HuggingFace Hub: [https://huggingface.co/Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B)
+* ModelScope: [https://modelscope.cn/Qwen/Qwen3-0.6B](https://modelscope.cn/Qwen/Qwen3-0.6B)
+
+本实验完整代码GitHub链接：
+
+* 完整项目GitHub地址：[https://github.com/ShaohonChen/Qwen3-SmVL](https://github.com/ShaohonChen/Qwen3-SmVL)
+
+本实验SwanLab日志：
+
+* SwanLab训练过程查看：[https://swanlab.cn/@ShaohonChen/Qwen3-SmVL/overview](https://swanlab.cn/@ShaohonChen/Qwen3-SmVL/overview)
+
+## 参考资料
+
+* Huggingface SmolVLM2技术报告：[https://arxiv.org/pdf/2504.05299](https://arxiv.org/pdf/2504.05299)

Extra-Chapter/why-fine-tune-small-large-language-models/readme.md

@@ -0,0 +1,388 @@
+# 大模型都这么厉害了，微调0.6B的小模型有什么意义？
+
+大家在日常使用Deepseek-R1或者是阿里新发布的Qwen3模型，他们的模型都是能力很强，所提供的API服也都可以满足大家的日常或者是公司开发所需。但大家也可以想一个简单的问题几个简单的问题，如下：
+
+1. 公司的数据是够敏感，是否需要保密？
+1. 日常使用大模型的任务是否很困难，对推理链是否刚需？
+1. 任务调用的大模型API并发量是多少？每日资金消耗有多少？
+
+对于问题1，如果公司数据敏感，那我建议不要调用供应商提供的大模型API。就算供应商保证不会拿你们数据做训练，但你们的数据还是泄漏了（会有不必要的风险），建议本地部署大模型。
+
+对于问题2，如果使用大模型的场景问题很困难并且刚需推理链，那可以使用供应商的API，这样可以保证推理链的上下文不会爆显存。如果问题很简单，没有刚需推理链，那建议本地部署小模型即可。
+
+对于问题3，如果任务很简单，且调用的大模型API并发量很高，那我建议微调一个特定任务的小模型，本地部署。这样可以满足高并发，并且可以减少资金消耗。（本地部署，默认硬件环境单卡4090）
+
+看到这里，想必大家已经思考完了以上三个问题，心中有了答案。那我给出一个小小的案例。
+
+## 微调模型的需求性
+
+假如你的公司有一个从投诉的文本中抽取用户信息的任务。比如，你需要从以下文本中抽取用户姓名、住址、邮箱、投诉的问题等等。
+
+> 这只是一个小小的案例，数据也是我用大模型批量制造的。真正的投诉数据不会这么“干净、整洁”。
+
+INPUT：
+```text
+龙琳，宁夏回族自治区璐市城东林街g座 955491，邮箱 nafan@example.com。小区垃圾堆积成山，晚上噪音扰人清梦，停车难上加难，简直无法忍受！
+```
+
+OUTPUT：
+```json
+{
+    "name": "龙琳",
+    "address": "宁夏回族自治区璐市城东林街g座 955491",
+    "email": "nafan@example.com",
+    "question": "小区垃圾堆积成山，晚上噪音扰人清梦，停车难上加难，简直无法忍受！"
+}
+```
+
+
+那你当然可以调用 Deepseek最强大的模型R1，也可以调用阿里最新发布最强大的模型 Qwen3-235B-A22B等等，这些模型的信息抽取效果也很非常的棒。
+
+但有个问题，如果你有几百万条这样的数据要处理，全部调用最新的，最好的大模型可能需要消耗几万块钱。并且，如果这些投诉数据，比如电信投诉数据，电网投诉数据，这些数据是敏感的不可以直接放到外网的。
+
+所以，综合数据敏感，和资金消耗。最好的选择就是微调一个小模型（如Qwen3-0.6B），既可以保证高并发，可以保证数据不泄漏，保证模型抽取的效果，还可以省钱！！！
+
+那下面，用一个小案例带大家实操一下，微调Qwen3-0.6B小模型完成文本信息抽取任务。
+
+
+
+## 配置环境 下载数据
+
+> Colab 文件地址：https://colab.research.google.com/drive/18ByY11KVhIy6zWx1uKUjSzqeHTme-TtU?usp=drive_link
+
+```python
+!pip install datasets swanlab -q
+```
+
+```python
+!wget --no-check-certificate 'https://docs.google.com/uc?export=download&id=1a0sf5C209CLW5824TJkUM4olMy0zZWpg' -O fake_sft.json
+```
+
+## 处理数据
+
+```python
+from datasets import Dataset
+import pandas as pd
+from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer, GenerationConfig
+from peft import LoraConfig, TaskType, get_peft_model
+import torch
+```
+
+```python
+# 将JSON文件转换为CSV文件
+df = pd.read_json('fake_sft.json')
+ds = Dataset.from_pandas(df)
+ds[:3]
+```
+
+```python
+model_id = "Qwen/Qwen3-0.6B"
+```
+
+```python
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
+tokenizer
+```
+
+对大语言模型进行 `supervised-finetuning`（`sft`，有监督微调）的数据格式如下：
+
+```json
+{
+  "instruction": "回答以下用户问题，仅输出答案。",
+  "input": "1+1等于几?",
+  "output": "2"
+}
+```
+
+其中，`instruction` 是用户指令，告知模型其需要完成的任务；`input` 是用户输入，是完成用户指令所必须的输入内容；`output` 是模型应该给出的输出。
+
+有监督微调的目标是让模型具备理解并遵循用户指令的能力。因此，在构建数据集时，我们应针对我们的目标任务，针对性构建数据。比如，如果我们的目标是通过大量人物的对话数据微调得到一个能够 role-play 甄嬛对话风格的模型，因此在该场景下的数据示例如下：
+
+```json
+{
+  "instruction": "你父亲是谁？",
+  "input": "",
+  "output": "家父是大理寺少卿甄远道。"
+}
+```
+
+`Qwen3` 采用的 `Chat Template`格式如下：
+
+由于 `Qwen3` 是混合推理模型，因此可以手动选择开启思考模式
+
+不开启 `thinking mode`
+
+
+```python
+messages = [
+    {"role": "system", "content": "You are a helpful AI"},
+    {"role": "user", "content": "How are you?"},
+    {"role": "assistant", "content": "I'm fine, think you. and you?"},
+]
+
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
+    enable_thinking=False
+)
+print(text)
+```
+```
+<|im_start|>system
+You are a helpful AI<|im_end|>
+<|im_start|>user
+How are you?<|im_end|>
+<|im_start|>assistant
+<think>
+
+</think>
+
+I'm fine, think you. and you?<|im_end|>
+<|im_start|>assistant
+<think>
+
+</think>
+``` 
+
+
+`LoRA`（`Low-Rank Adaptation`）训练的数据是需要经过格式化、编码之后再输入给模型进行训练的，我们需要先将输入文本编码为 `input_ids`，将输出文本编码为 `labels`，编码之后的结果是向量。我们首先定义一个预处理函数，这个函数用于对每一个样本，同时编码其输入、输出文本并返回一个编码后的字典：
+
+
+```python
+def process_func(example):
+    MAX_LENGTH = 1024 # 设置最大序列长度为1024个token
+    input_ids, attention_mask, labels = [], [], [] # 初始化返回值
+    # 适配chat_template
+    instruction = tokenizer(
+        f"<s><|im_start|>system\n{example['system']}<|im_end|>\n"
+        f"<|im_start|>user\n{example['instruction'] + example['input']}<|im_end|>\n"
+        f"<|im_start|>assistant\n<think>\n\n</think>\n\n",
+        add_special_tokens=False
+    )
+    response = tokenizer(f"{example['output']}", add_special_tokens=False)
+    # 将instructio部分和response部分的input_ids拼接，并在末尾添加eos token作为标记结束的token
+    input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
+    # 注意力掩码，表示模型需要关注的位置
+    attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1]
+    # 对于instruction，使用-100表示这些位置不计算loss（即模型不需要预测这部分）
+    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]
+    if len(input_ids) > MAX_LENGTH:  # 超出最大序列长度截断
+        input_ids = input_ids[:MAX_LENGTH]
+        attention_mask = attention_mask[:MAX_LENGTH]
+        labels = labels[:MAX_LENGTH]
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "labels": labels
+    }
+
+```
+
+
+```python
+tokenized_id = ds.map(process_func, remove_columns=ds.column_names)
+tokenized_id
+```
+
+
+```python
+tokenizer.decode(tokenized_id[0]['input_ids'])
+```
+
+```python
+tokenizer.decode(list(filter(lambda x: x != -100, tokenized_id[1]["labels"])))
+```
+
+## 加载模型
+
+加载模型并配置LoraConfig
+
+```python
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto",torch_dtype=torch.bfloat16)
+model
+```
+
+
+```
+Qwen3ForCausalLM(
+  (model): Qwen3Model(
+    (embed_tokens): Embedding(151936, 1024)
+    (layers): ModuleList(
+      (0-27): 28 x Qwen3DecoderLayer(
+        (self_attn): Qwen3Attention(
+          (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+          (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+          (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+          (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+          (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+          (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+        )
+        (mlp): Qwen3MLP(
+          (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+          (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+          (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+          (act_fn): SiLU()
+        )
+        (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+      )
+    )
+    (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+    (rotary_emb): Qwen3RotaryEmbedding()
+  )
+  (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+)
+```
+
+```python
+model.enable_input_require_grads() # 开启梯度检查点时，要执行该方法
+```
+
+
+## Lora Config
+
+`LoraConfig`这个类中可以设置很多参数，比较重要的如下
+
+- `task_type`：模型类型，现在绝大部分 `decoder_only` 的模型都是因果语言模型 `CAUSAL_LM`
+- `target_modules`：需要训练的模型层的名字，主要就是 `attention`部分的层，不同的模型对应的层的名字不同
+- `r`：`LoRA` 的秩，决定了低秩矩阵的维度，较小的 `r` 意味着更少的参数
+- `lora_alpha`：缩放参数，与 `r` 一起决定了 `LoRA` 更新的强度。实际缩放比例为`lora_alpha/r`，在当前示例中是 `32 / 8 = 4` 倍
+- `lora_dropout`：应用于 `LoRA` 层的 `dropout rate`，用于防止过拟合
+
+
+```python
+from peft import LoraConfig, TaskType, get_peft_model
+
+config = LoraConfig(
+    task_type=TaskType.CAUSAL_LM,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+    inference_mode=False, # 训练模式
+    r=8, # Lora 秩
+    lora_alpha=32, # Lora alaph，具体作用参见 Lora 原理
+    lora_dropout=0.1# Dropout 比例
+)
+config
+```
+
+
+```python
+model = get_peft_model(model, config)
+config
+```
+
+
+```python
+model.print_trainable_parameters()  # 模型参数训练量只有0.8395%
+```
+
+> trainable params: 5,046,272 || all params: 601,096,192 || trainable%: 0.8395
+
+
+## Training Arguments
+
+- `output_dir`：模型的输出路径
+- `per_device_train_batch_size`：每张卡上的 `batch_size`
+- `gradient_accumulation_steps`: 梯度累计
+- `num_train_epochs`：顾名思义 `epoch`
+
+
+```python
+args = TrainingArguments(
+    output_dir="Qwen3_instruct_lora",
+    per_device_train_batch_size=4,
+    gradient_accumulation_steps=4,
+    logging_steps=1,
+    num_train_epochs=3,
+    save_steps=50,
+    learning_rate=1e-4,
+    save_on_each_node=True,
+    gradient_checkpointing=True,
+    report_to="none",
+)
+```
+
+## SwanLab 简介
+
+[SwanLab](https://github.com/swanhubx/swanlab) 是一个开源的模型训练记录工具，面向 AI 研究者，提供了训练可视化、自动日志记录、超参数记录、实验对比、多人协同等功能。在 `SwanLab` 上，研究者能基于直观的可视化图表发现训练问题，对比多个实验找到研究灵感，并通过在线链接的分享与基于组织的多人协同训练，打破团队沟通的壁垒。
+
+**为什么要记录训练**
+
+相较于软件开发，模型训练更像一个实验科学。一个品质优秀的模型背后，往往是成千上万次实验。研究者需要不断尝试、记录、对比，积累经验，才能找到最佳的模型结构、超参数与数据配比。在这之中，如何高效进行记录与对比，对于研究效率的提升至关重要。
+
+`(2) Use an existing SwanLab account` 并使用 private API Key 登录
+
+```python
+import swanlab
+from swanlab.integration.transformers import SwanLabCallback
+
+# 实例化SwanLabCallback
+swanlab_callback = SwanLabCallback(
+    project="Qwen3-Lora",  # 注意修改
+    experiment_name="Qwen3-8B-LoRA-experiment"  # 注意修改
+)
+```
+
+
+```python
+import swanlab
+from swanlab.integration.transformers import SwanLabCallback
+
+# 实例化SwanLabCallback
+swanlab_callback = SwanLabCallback(
+    project="Qwen3-Lora",
+    experiment_name="Qwen3-0.6B-extarct-lora-2"
+)
+```
+
+
+```python
+trainer = Trainer(
+    model=model,
+    args=args,
+    train_dataset=tokenized_id,
+    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
+    callbacks=[swanlab_callback]
+)
+```
+
+
+```python
+trainer.train()
+```
+
+
+## 测试文本
+
+
+```python
+prompt = "龙琳   ，宁夏回族自治区璐市城东林街g座 955491，nafan@example.com。小区垃圾堆积成山，晚上噪音扰人清梦，停车难上加难，简直无法忍受！太插件了阿萨德看见啊啥的健康仨都会撒娇看到撒谎的、"
+
+messages = [
+    {"role": "system", "content": "将文本中的name、address、email、question提取出来，以json格式输出，字段为name、address、email、question，值为文本中提取出来的内容。"},
+    {"role": "user", "content": prompt}
+]
+
+inputs = tokenizer.apply_chat_template(messages,
+                                       add_generation_prompt=True,
+                                       tokenize=True,
+                                       return_tensors="pt",
+                                       return_dict=True,
+                                       enable_thinking=False).to('cuda')
+
+gen_kwargs = {"max_length": 2500, "do_sample": True, "top_k": 1}
+with torch.no_grad():
+    outputs = model.generate(**inputs, **gen_kwargs)
+    outputs = outputs[:, inputs['input_ids'].shape[1]:]
+    print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+
+```json
+{
+    "name": "龙琳",
+    "address": "宁夏回族自治区璐市城东林街g座 955491",
+    "email": "nafan@example.com",
+    "question": "小区垃圾堆积成山，晚上噪音扰人清梦，停车难上加难，简直无法忍受！太插件了阿萨德看见啊啥的健康仨都会撒娇看到撒谎的、"
+}
+```

README.md

@@ -4,9 +4,15 @@
 </div>
 
 <div align="center">
-  <img src="https://img.shields.io/github/stars/datawhalechina/happy-llm?style=for-the-badge&logo=github" alt="GitHub stars"/>
-  <img src="https://img.shields.io/github/forks/datawhalechina/happy-llm?style=for-the-badge&logo=github" alt="GitHub forks"/>
-  <img src="https://img.shields.io/badge/language-Chinese-brightgreen?style=for-the-badge" alt="Language"/>
+  <img src="https://img.shields.io/github/stars/datawhalechina/happy-llm?style=flat&logo=github" alt="GitHub stars"/>
+  <img src="https://img.shields.io/github/forks/datawhalechina/happy-llm?style=flat&logo=github" alt="GitHub forks"/>
+  <img src="https://img.shields.io/badge/language-Chinese-brightgreen?style=flat" alt="Language"/>
+  <a href="https://github.com/datawhalechina/happy-llm"><img src="https://img.shields.io/badge/GitHub-Project-blue?style=flat&logo=github" alt="GitHub Project"></a>
+  <a href="https://swanlab.cn/@kmno4/Happy-LLM/overview"><img src="https://raw.githubusercontent.com/SwanHubX/assets/main/badge1.svg" alt="SwanLab"></a>
+</div>
+
+<div align="center">
+  <a href="https://trendshift.io/repositories/14175" target="_blank"><img src="https://trendshift.io/api/badge/repositories/14175" alt="datawhalechina%2Fhappy-llm | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
 </div>
 
 <div align="center">
@@ -16,6 +22,7 @@
 </div>
 
 <div align="center">
+  <p><a href="https://datawhalechina.github.io/happy-llm/">📚 在线阅读地址</a></p>
   <h3>📚 从零开始的大语言模型原理与实践教程</h3>
   <p><em>深入理解 LLM 核心原理，动手实现你的第一个大模型</em></p>
 </div>
@@ -42,14 +49,55 @@
 
 | 章节 | 关键内容 | 状态 |
 | --- | --- | --- |
-| [前言](./docs/README.md) | 本项目的缘起、背景及读者建议 | ✅ |
+| [前言](./docs/前言.md) | 本项目的缘起、背景及读者建议 | ✅ |
 | [第一章 NLP 基础概念](./docs/chapter1/第一章%20NLP基础概念.md) | 什么是 NLP、发展历程、任务分类、文本表示演进 | ✅ |
 | [第二章 Transformer 架构](./docs/chapter2/第二章%20Transformer架构.md) | 注意力机制、Encoder-Decoder、手把手搭建 Transformer | ✅ |
 | [第三章 预训练语言模型](./docs/chapter3/第三章%20预训练语言模型.md) | Encoder-only、Encoder-Decoder、Decoder-Only 模型对比 | ✅ |
 | [第四章 大语言模型](./docs/chapter4/第四章%20大语言模型.md) | LLM 定义、训练策略、涌现能力分析 | ✅ |
 | [第五章 动手搭建大模型](./docs/chapter5/第五章%20动手搭建大模型.md) | 实现 LLaMA2、训练 Tokenizer、预训练小型 LLM | ✅ |
-| [第六章 大模型训练实践](./docs/chapter6/第六章%20大模型训练流程实践.md) | 预训练、有监督微调、LoRA/QLoRA 高效微调 | ✅ |
+| [第六章 大模型训练实践](./docs/chapter6/第六章%20大模型训练流程实践.md) | 预训练、有监督微调、LoRA/QLoRA 高效微调 | 🚧 |
 | [第七章 大模型应用](./docs/chapter7/第七章%20大模型应用.md) | 模型评测、RAG 检索增强、Agent 智能体 | ✅ |
+| [Extra Chapter LLM Blog](./Extra-Chapter/) | 优秀的大模型 学习笔记/Blog ，欢迎大家来 PR ！| 🚧 |
+
+### Extra Chapter LLM Blog
+
+- [大模型都这么厉害了，微调0.6B的小模型有什么意义？](./Extra-Chapter/why-fine-tune-small-large-language-models/readme.md) @[不要葱姜蒜](https://github.com/KMnO4-zx) 2025-7-11
+
+- [Transformer 整体模块设计解读](./Extra-Chapter/transformer-architecture/) @[ditingdapeng](https://github.com/ditingdapeng) 2025-7-14
+
+- [文本数据处理详解](./Extra-Chapter/text-data-processing/readme.md) @[蔡鋆捷](https://github.com/xinala-781) 2025-7-14
+
+- [Qwen3-"VL"——超小中文多模态模型的“拼接微调”之路](./Extra-Chapter/vlm-concatenation-finetune/README.md) @[ShaohonChen](https://github.com/ShaohonChen) 2025-7-30
+
+- [S1: Thinking Budget with vLLM](./Extra-Chapter/s1-vllm-thinking-budget/readme.md) @[不要葱姜蒜](https://github.com/kmno4-zx) 2025-8-03
+
+
+- [CDDRS: 使用细粒度语义信息指导增强的RAG检索方法](./Extra-Chapter/CDDRS/readme.md) @[Hongru0306](https://github.com/Hongru0306) 2025-8-21
+
+- [大模型生成 Token 的方式有哪些？](./Extra-Chapter/generation-method/readme.md) @[不要葱姜蒜](https://github.com/kmno4-zx) 2025-10-17
+
+
+> &emsp;&emsp;*如果大家在学习 Happy-LLM 项目或 LLM 相关知识中有自己独到的见解、认知、实践，欢迎大家 PR 在 [Extra Chapter LLM Blog](./Extra-Chapter/) 中。请遵守 Extra Chapter LLM Blog 的 [PR 规范](./Extra-Chapter/Readme.md)，我们会视 PR 内容的质量和价值来决定是否合并或补充到 Happy-LLM 正文中来。*
+
+### 模型下载
+
+| 模型名称 | 下载地址 |
+| --- | --- |
+| Happy-LLM-Chapter5-Base-215M | [🤖 ModelScope](https://www.modelscope.cn/models/kmno4zx/happy-llm-215M-base) |
+| Happy-LLM-Chapter5-SFT-215M | [🤖 ModelScope](https://www.modelscope.cn/models/kmno4zx/happy-llm-215M-sft) |
+
+> *ModelScope 创空间体验地址：[🤖 创空间](https://www.modelscope.cn/studios/kmno4zx/happy_llm_215M_sft)*
+
+
+### PDF 版本下载
+
+&emsp;&emsp;***本 Happy-LLM PDF 教程完全开源免费。为防止各类营销号加水印后贩卖给大模型初学者，我们特地在 PDF 文件中预先添加了不影响阅读的 Datawhale 开源标志水印，敬请谅解～***
+
+> *Happy-LLM PDF : https://github.com/datawhalechina/happy-llm/releases/tag/v1.0.2*  
+
+### PPT 资源下载
+
+&emsp;&emsp;***本项目配套教学讲义PPT课件资源获取链接：https://github.com/HZAI-ZJNU/happy-llm-ppt 或可在本项目的 [Releases](https://github.com/datawhalechina/happy-llm/releases) 页面下载。***
 
 ## 💡 如何学习
 
@@ -61,6 +109,8 @@
 
 &emsp;&emsp;最后，欢迎每一位读者在学习完本项目后加入到 LLM 开发者的行列。作为国内 AI 开源社区，我们希望充分聚集共创者，一起丰富这个开源 LLM 的世界，打造更多、更全面特色 LLM 的教程。星火点点，汇聚成海。我们希望成为 LLM 与普罗大众的阶梯，以自由、平等的开源精神，拥抱更恢弘而辽阔的 LLM 世界。
 
+> - 中国计算机学会(CCF) × Datawhale × GitLink开源平台联合推出AI普惠课程，免费算力报名参加 [【报名地址】](https://mp.weixin.qq.com/s/P03f3e2vUUh7OxDP40Ra6w)[【GitLink 地址】](https://gitlink.org.cn/datawhalechina/happy-llm)
+
 ## 🤝 如何贡献
 
 我们欢迎任何形式的贡献！
@@ -73,10 +123,17 @@
 ## 🙏 致谢
 
 ### 核心贡献者
-- [宋志学-项目负责人](https://github.com/KMnO4-zx) (Datawhale成员-中国矿业大学(北京))
+- [宋志学-项目负责人](https://github.com/KMnO4-zx) (Datawhale成员)
 - [邹雨衡-项目负责人](https://github.com/logan-zou) (Datawhale成员-对外经济贸易大学)
 - [朱信忠-指导专家](https://xinzhongzhu.github.io/)（Datawhale首席科学家-浙江师范大学杭州人工智能研究院教授）
 
+### Extra-Chapter 贡献者
+
+- [ditingdapeng](https://github.com/ditingdapeng)（内容贡献者-云原生基础架构工程师）
+- [蔡鋆捷](https://github.com/xinala-781)（内容贡献者-福州大学）
+- [ShaohonChen](https://github.com/ShaohonChen) （情感机器实验室研究员-西安电子科技大学在读硕士）
+- [肖鸿儒, 庄健琨](https://github.com/Hongru0306) (内容贡献者-同济大学)
+
 ### 特别感谢
 - 感谢 [@Sm1les](https://github.com/Sm1les) 对本项目的帮助与支持
 - 感谢所有为本项目做出贡献的开发者们 ❤️
@@ -90,7 +147,7 @@
 ## Star History
 
 <div align='center'>
-    <img src="./images/star-history-202566.png" alt="Datawhale" width="90%">
+    <img src="./images/star-history-20251017.png" alt="Datawhale" width="90%">
 </div>
 
 <div align="center">

README_en.md

@@ -4,11 +4,12 @@
 </div>
 
 <div align="center">
-  <img src="https://img.shields.io/github/stars/datawhalechina/happy-llm?style=for-the-badge&logo=github" alt="GitHub stars"/>
-  <img src="https://img.shields.io/github/forks/datawhalechina/happy-llm?style=for-the-badge&logo=github" alt="GitHub forks"/>
-  <img src="https://img.shields.io/badge/language-English-brightgreen?style=for-the-badge" alt="Language"/>
+  <img src="https://img.shields.io/github/stars/datawhalechina/happy-llm?style=flat&logo=github" alt="GitHub stars"/>
+  <img src="https://img.shields.io/github/forks/datawhalechina/happy-llm?style=flat&logo=github" alt="GitHub forks"/>
+  <img src="https://img.shields.io/badge/language-Chinese-brightgreen?style=flat" alt="Language"/>
+  <a href="https://github.com/datawhalechina/happy-llm"><img src="https://img.shields.io/badge/GitHub-Project-blue?style=flat&logo=github" alt="GitHub Project"></a>
+  <a href="https://swanlab.cn/@kmno4/Happy-LLM/overview"><img src="https://raw.githubusercontent.com/SwanHubX/assets/main/badge1.svg" alt="SwanLab"></a>
 </div>
-
 <div align="center">
 
 [中文](./README.md) | [English](./README_en.md)
@@ -16,6 +17,7 @@
 </div>
 
 <div align="center">
+  <p><a href="https://datawhalechina.github.io/happy-llm/">📚 Online Reading</a></p>
   <h3>📚 A Comprehensive Tutorial on Large Language Model Principles and Practice from Scratch</h3>
   <p><em>Deep understanding of LLM core principles, hands-on implementation of your first large model</em></p>
 </div>
@@ -48,8 +50,44 @@
 | [Chapter 3: Pre-trained Language Models](./docs/chapter3/第三章%20预训练语言模型.md) | Comparison of Encoder-only, Encoder-Decoder, Decoder-Only models | ✅ |
 | [Chapter 4: Large Language Models](./docs/chapter4/第四章%20大语言模型.md) | LLM definition, training strategies, emergent ability analysis | ✅ |
 | [Chapter 5: Building Large Models from Scratch](./docs/chapter5/第五章%20动手搭建大模型.md) | Implementing LLaMA2, training Tokenizer, pre-training small LLM | ✅ |
-| [Chapter 6: Large Model Training Practice](./docs/chapter6/第六章%20大模型训练流程实践.md) | Pre-training, supervised fine-tuning, LoRA/QLoRA efficient fine-tuning | ✅ |
+| [Chapter 6: Large Model Training Practice](./docs/chapter6/第六章%20大模型训练流程实践.md) | Pre-training, supervised fine-tuning, LoRA/QLoRA efficient fine-tuning | 🚧 |
 | [Chapter 7: Large Model Applications](./docs/chapter7/第七章%20大模型应用.md) | Model evaluation, RAG retrieval enhancement, Agent intelligent agents | ✅ |
+| [Extra Chapter LLM Blog](./Extra-Chapter/) | Excellent Learning Notes/Blog on LLMs ，Welcome PR ！| 🚧 |
+
+### Extra Chapter LLM Blog
+
+- [With large models becoming so powerful, what’s the significance of fine-tuning a 0.6B small model?](./Extra-Chapter/why-fine-tune-small-large-language-models/readme.md) @[不要葱姜蒜](https://github.com/KMnO4-zx) 2025-7-11
+
+- [Details of the Transformer modules](./Extra-Chapter/transformer-architecture/) @[ditingdapeng](https://github.com/ditingdapeng) 2025-7-14
+
+- [Detailed Explanation of Text Data Processing](./Extra-Chapter/text-data-processing/readme.md) @[蔡鋆捷](https://github.com/xinala-781) 2025-7-14
+
+- [Qwen3-"VL"——Path to 'Concatenation Fine-tuning' for Ultra-small Chinese Multimodal Models](./Extra-Chapter/vlm-concatenation-finetune/README.md) @[ShaohonChen](https://github.com/ShaohonChen) 2025-7-30
+
+- [S1: Thinking Budget with vLLM](./Extra-Chapter/s1-vllm-thinking-budget/readme.md) @[kmno4-zx](https://github.com/kmno4-zx) 2025-8-03
+
+- [CDDRS: Key elements guided Enhancement for RAG-based Retrieval Methods](./Extra-Chapter/CDDRS/readme.md) @[Hongru0306](https://github.com/Hongru0306) 2025-8-21
+
+
+> &emsp;&emsp;* If anyone has unique insights, knowledge, or practices related to the Happy-LLM project or LLMs in general, you are welcome to submit a PR to the [Extra Chapter LLM Blog](./Extra-Chapter/). Please adhere to the [PR Guidances](./Extra-Chapter/Readme.md). We will decide whether to merge or supplement the content into the main Happy-LLM text based on the quality and value of the PR.*
+
+
+### Model Downloads
+
+| Model Name | Download Link |
+| --- | --- |
+| Happy-LLM-Chapter5-Base-215M | [🤖 ModelScope](https://www.modelscope.cn/models/kmno4zx/happy-llm-215M-base) |
+| Happy-LLM-Chapter5-SFT-215M | [🤖 ModelScope](https://www.modelscope.cn/models/kmno4zx/happy-llm-215M-sft) |
+
+> *ModelScope Studio Experience: [🤖 Studio](https://www.modelscope.cn/studios/kmno4zx/happy_llm_215M_sft)*
+
+
+### PDF Version Download
+
+&emsp;&emsp;***This Happy-LLM PDF tutorial is completely open source and free. To prevent various marketing accounts from adding watermarks and selling to LLM beginners, we have pre-added Datawhale open source logo watermarks that do not affect reading in the PDF files. Please understand~***
+
+> *Happy-LLM PDF : https://github.com/datawhalechina/happy-llm/releases/tag/PDF*  
+> *Happy-LLM PDF Domestic Download: https://www.datawhale.cn/learn/summary/179*
 
 ## 💡 How to Learn
 
@@ -73,7 +111,7 @@ We welcome any form of contribution!
 ## 🙏 Acknowledgments
 
 ### Core Contributors
-- [Song Zhixue - Project Leader](https://github.com/KMnO4-zx) (Datawhale Member - China University of Mining and Technology, Beijing)
+- [Song Zhixue - Project Leader](https://github.com/KMnO4-zx) (Datawhale Member)
 - [Zou Yuheng - Project Leader](https://github.com/logan-zou) (Datawhale Member - University of International Business and Economics)
 - [Zhu Xinzhong - Expert Advisor](https://xinzhongzhu.github.io/) (Datawhale Chief Scientist - Professor at Hangzhou Institute for Advanced Study, Zhejiang Normal University)
 
@@ -90,7 +128,7 @@ We welcome any form of contribution!
 ## Star History
 
 <div align='center'>
-    <img src="./images/star-history-202566.png" alt="Datawhale" width="90%">
+    <img src="./images/star-history-20251017.png" alt="Datawhale" width="90%">
 </div>
 
 <div align="center">
@@ -101,11 +139,17 @@ We welcome any form of contribution!
 
 <div align='center'>
     <img src="./images/datawhale.png" alt="Datawhale" width="30%">
-    <p>Scan the QR code to follow Datawhale WeChat Official Account for more quality open-source content</p>
+    <p>Scan the QR code to follow Datawhale WeChat Official Account for more quality open source content</p>
 </div>
 
 ---
 
 ## 📜 Open Source License
 
+This work is licensed under a [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License](http://creativecommons.org/licenses/by-nc-sa/4.0/).
+
+---
+
+## 📜 Open Source License
+
 This work is licensed under a [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License](http://creativecommons.org/licenses/by-nc-sa/4.0/).

docs/README.md

@@ -4,15 +4,15 @@
 </div>
 
 <div align="center">
-  <img src="https://img.shields.io/github/stars/datawhalechina/happy-llm?style=for-the-badge&logo=github" alt="GitHub stars"/>
-  <img src="https://img.shields.io/github/forks/datawhalechina/happy-llm?style=for-the-badge&logo=github" alt="GitHub forks"/>
-  <img src="https://img.shields.io/badge/language-Chinese-brightgreen?style=for-the-badge" alt="Language"/>
+  <img src="https://img.shields.io/github/stars/datawhalechina/happy-llm?style=flat&logo=github" alt="GitHub stars"/>
+  <img src="https://img.shields.io/github/forks/datawhalechina/happy-llm?style=flat&logo=github" alt="GitHub forks"/>
+  <img src="https://img.shields.io/badge/language-Chinese-brightgreen?style=flat" alt="Language"/>
+  <a href="https://github.com/datawhalechina/happy-llm"><img src="https://img.shields.io/badge/GitHub-Project-blue?style=flat&logo=github" alt="GitHub Project"></a>
+  <a href="https://swanlab.cn/@kmno4/Happy-LLM/overview"><img src="https://raw.githubusercontent.com/SwanHubX/assets/main/badge1.svg" alt="SwanLab"></a>
 </div>
 
 <div align="center">
-
-[中文](./README.md) | [English](./README_en.md)
-
+  <a href="https://trendshift.io/repositories/14175" target="_blank"><img src="https://trendshift.io/api/badge/repositories/14175" alt="datawhalechina%2Fhappy-llm | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
 </div>
 
 <div align="center">
@@ -42,14 +42,30 @@
 
 | 章节 | 关键内容 | 状态 |
 | --- | --- | --- |
-| [前言](./docs/README.md) | 本项目的缘起、背景及读者建议 | ✅ |
-| [第一章 NLP 基础概念](./docs/chapter1/第一章%20NLP基础概念.md) | 什么是 NLP、发展历程、任务分类、文本表示演进 | ✅ |
-| [第二章 Transformer 架构](./docs/chapter2/第二章%20Transformer架构.md) | 注意力机制、Encoder-Decoder、手把手搭建 Transformer | ✅ |
-| [第三章 预训练语言模型](./docs/chapter3/第三章%20预训练语言模型.md) | Encoder-only、Encoder-Decoder、Decoder-Only 模型对比 | ✅ |
-| [第四章 大语言模型](./docs/chapter4/第四章%20大语言模型.md) | LLM 定义、训练策略、涌现能力分析 | ✅ |
-| [第五章 动手搭建大模型](./docs/chapter5/第五章%20动手搭建大模型.md) | 实现 LLaMA2、训练 Tokenizer、预训练小型 LLM | ✅ |
-| [第六章 大模型训练实践](./docs/chapter6/第六章%20大模型训练流程实践.md) | 预训练、有监督微调、LoRA/QLoRA 高效微调 | ✅ |
-| [第七章 大模型应用](./docs/chapter7/第七章%20大模型应用.md) | 模型评测、RAG 检索增强、Agent 智能体 | ✅ |
+| [前言](./前言.md) | 本项目的缘起、背景及读者建议 | ✅ |
+| [第一章 NLP 基础概念](./chapter1/第一章%20NLP基础概念.md) | 什么是 NLP、发展历程、任务分类、文本表示演进 | ✅ |
+| [第二章 Transformer 架构](./chapter2/第二章%20Transformer架构.md) | 注意力机制、Encoder-Decoder、手把手搭建 Transformer | ✅ |
+| [第三章 预训练语言模型](./chapter3/第三章%20预训练语言模型.md) | Encoder-only、Encoder-Decoder、Decoder-Only 模型对比 | ✅ |
+| [第四章 大语言模型](./chapter4/第四章%20大语言模型.md) | LLM 定义、训练策略、涌现能力分析 | ✅ |
+| [第五章 动手搭建大模型](./chapter5/第五章%20动手搭建大模型.md) | 实现 LLaMA2、训练 Tokenizer、预训练小型 LLM | ✅ |
+| [第六章 大模型训练实践](./chapter6/第六章%20大模型训练流程实践.md) | 预训练、有监督微调、LoRA/QLoRA 高效微调 | 🚧 |
+| [第七章 大模型应用](./chapter7/第七章%20大模型应用.md) | 模型评测、RAG 检索增强、Agent 智能体 | ✅ |
+
+### 模型下载
+
+| 模型名称 | 下载地址 |
+| --- | --- |
+| Happy-LLM-Chapter5-Base-215M | [🤖 ModelScope](https://www.modelscope.cn/models/kmno4zx/happy-llm-215M-base) |
+| Happy-LLM-Chapter5-SFT-215M | [🤖 ModelScope](https://www.modelscope.cn/models/kmno4zx/happy-llm-215M-sft) |
+
+> *ModelScope 创空间体验地址：[🤖 创空间](https://www.modelscope.cn/studios/kmno4zx/happy_llm_215M_sft)*
+
+### PDF 版本下载
+
+&emsp;&emsp;***本 Happy-LLM PDF 教程完全开源免费。为防止各类营销号加水印后贩卖给大模型初学者，我们特地在 PDF 文件中预先添加了不影响阅读的 Datawhale 开源标志水印，敬请谅解～***
+
+> *Happy-LLM PDF : https://github.com/datawhalechina/happy-llm/releases/tag/PDF*  
+> *Happy-LLM PDF 国内下载地址 : https://www.datawhale.cn/learn/summary/179* 
 
 ## 💡 如何学习
 
@@ -90,7 +106,7 @@
 ## Star History
 
 <div align='center'>
-    <img src="./images/star-history-202566.png" alt="Datawhale" width="90%">
+    <img src="./images/star-history-20251017.png" alt="Datawhale" width="90%">
 </div>
 
 <div align="center">

docs/chapter1/第一章 NLP基础概念.md

@@ -67,9 +67,9 @@ NLP 的早期探索始于二战后，当时人们认识到将一种语言自动
 子词切分的方法有很多种，常见的有Byte Pair Encoding (BPE)、WordPiece、Unigram、SentencePiece等。这些方法的基本思想是将单词分解成更小的、频繁出现的片段，这些片段可以是单个字符、字符组合或者词根和词缀。
 
 ```
-输出：unhappiness
+输入：unhappiness
 
-不使用子词切分：整个单词作为一个单位：“unhappiness”
+不使用子词切分：整个单词作为一个单位，输出：“unhappiness”
 使用子词切分（假设BPE算法）：单词被分割为：“un”、“happi”、“ness”
 ```
 
@@ -219,6 +219,7 @@ vector = [0, 0, ..., 1, 0, ..., 1, 0, ..., 1, 0, ..., 1, 0, ..., 1, 0, ...]
 # 实际有效维度：仅5维（非零维度）
 # 稀疏率：(16384-5)/16384 ≈ 99.97%
 ```
+> 词汇表是一个包含所有可能出现的词语的集合。在向量空间模型中，每个词对应词汇表中的一个位置，通过这种方式可以将词语转换为向量表示。例如，如果词汇表大小为 16384 ，那么每个词都会被表示为一个 16384 维的向量，其中只有该词对应的位置为 1，其他位置都为 0。
 
 为了解决这些问题，研究者们对向量空间模型的研究主要集中在两个方面：一是改进特征表示方法，如借助图方法、主题方法等进行关键词抽取；二是改进和优化特征项权重的计算方法，可以在现有方法的基础上进行融合计算或提出新的计算方法.
 

docs/chapter2/code/transformer.py

@@ -25,22 +25,19 @@ class MultiHeadAttention(nn.Module):
         # args: 配置对象
         super().__init__()
         # 隐藏层维度必须是头数的整数倍，因为后面我们会将输入拆成头数个矩阵
-        assert args.n_embd % args.n_heads == 0
-        # 模型并行处理大小，默认为1。
-        model_parallel_size = 1
-        # 本地计算头数，等于总头数除以模型并行处理大小。
-        self.n_local_heads = args.n_heads // model_parallel_size
+        assert args.dim % args.n_heads == 0
         # 每个头的维度，等于模型维度除以头的总数。
         self.head_dim = args.dim // args.n_heads
+        self.n_heads = args.n_heads
 
-        # Wq, Wk, Wv 参数矩阵，每个参数矩阵为 n_embd x n_embd
+        # Wq, Wk, Wv 参数矩阵，每个参数矩阵为 n_embd x dim
         # 这里通过三个组合矩阵来代替了n个参数矩阵的组合，其逻辑在于矩阵内积再拼接其实等同于拼接矩阵再内积，
         # 不理解的读者可以自行模拟一下，每一个线性层其实相当于n个参数矩阵的拼接
-        self.wq = nn.Linear(args.n_embd, args.n_heads * self.head_dim, bias=False)
-        self.wk = nn.Linear(args.n_embd, args.n_heads * self.head_dim, bias=False)
-        self.wv = nn.Linear(args.n_embd, args.n_heads * self.head_dim, bias=False)
-        # 输出权重矩阵，维度为 n_embd x n_embd（head_dim = n_embeds / n_heads）
-        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+        self.wq = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
+        # 输出权重矩阵，维度为 dim x dim（head_dim = dim / n_heads）
+        self.wo = nn.Linear(self.n_heads * self.head_dim, args.dim, bias=False)
         # 注意力的 dropout
         self.attn_dropout = nn.Dropout(args.dropout)
         # 残差连接的 dropout
@@ -60,16 +57,16 @@ class MultiHeadAttention(nn.Module):
         # 获取批次大小和序列长度，[batch_size, seq_len, dim]
         bsz, seqlen, _ = q.shape
 
-        # 计算查询（Q）、键（K）、值（V）,输入通过参数矩阵层，维度为 (B, T, n_embed) x (n_embed, n_embed) -> (B, T, n_embed)
+        # 计算查询（Q）、键（K）、值（V）,输入通过参数矩阵层，维度为 (B, T, n_embed) x (n_embed, dim) -> (B, T, dim)
         xq, xk, xv = self.wq(q), self.wk(k), self.wv(v)
 
-        # 将 Q、K、V 拆分成多头，维度为 (B, T, n_head, C // n_head)，然后交换维度，变成 (B, n_head, T, C // n_head)
+        # 将 Q、K、V 拆分成多头，维度为 (B, T, n_head, dim // n_head)，然后交换维度，变成 (B, n_head, T, dim // n_head)
         # 因为在注意力计算中我们是取了后两个维度参与计算
         # 为什么要先按B*T*n_head*C//n_head展开再互换1、2维度而不是直接按注意力输入展开，是因为view的展开方式是直接把输入全部排开，
         # 然后按要求构造，可以发现只有上述操作能够实现我们将每个头对应部分取出来的目标
-        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
-        xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
-        xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+        xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_heads, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_heads, self.head_dim)
         xq = xq.transpose(1, 2)
         xk = xk.transpose(1, 2)
         xv = xv.transpose(1, 2)
@@ -90,7 +87,7 @@ class MultiHeadAttention(nn.Module):
         output = torch.matmul(scores, xv)
 
         # 恢复时间维度并合并头。
-        # 将多头的结果拼接起来, 先交换维度为 (B, T, n_head, C // n_head)，再拼接成 (B, T, n_head * C // n_head)
+        # 将多头的结果拼接起来, 先交换维度为 (B, T, n_head, dim // n_head)，再拼接成 (B, T, n_head * dim // n_head)
         # contiguous 函数用于重新开辟一块新内存存储，因为Pytorch设置先transpose再view会报错，
         # 因为view直接基于底层存储得到，然而transpose并不会改变底层存储，因此需要额外存储
         output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
@@ -103,7 +100,7 @@ class MultiHeadAttention(nn.Module):
 class LayerNorm(nn.Module):
     ''' Layer Norm 层'''
     def __init__(self, features, eps=1e-6):
-        super(LayerNorm, self).__init__()
+        super().__init__()
         # 线性矩阵做映射
         self.a_2 = nn.Parameter(torch.ones(features))
         self.b_2 = nn.Parameter(torch.zeros(features))
@@ -130,7 +127,6 @@ class MLP(nn.Module):
     def forward(self, x):
         # 前向传播函数
         # 首先，输入x通过第一层线性变换和RELU激活函数
-        # 然后，结果乘以输入x通过第三层线性变换的结果
         # 最后，通过第二层线性变换和dropout层
         return self.dropout(self.w2(F.relu(self.w1(x))))
     
@@ -215,7 +211,7 @@ class PositionalEncoding(nn.Module):
     def __init__(self, args):
         super(PositionalEncoding, self).__init__()
         # Dropout 层
-        self.dropout = nn.Dropout(p=args.dropout)
+        # self.dropout = nn.Dropout(p=args.dropout)
 
         # block size 是序列的最大长度
         pe = torch.zeros(args.block_size, args.n_embd)
@@ -233,7 +229,7 @@ class PositionalEncoding(nn.Module):
     def forward(self, x):
         # 将位置编码加到 Embedding 结果上
         x = x + self.pe[:, : x.size(1)].requires_grad_(False)
-        return self.dropout(x)
+        return x
 
 
 class Transformer(nn.Module):
@@ -268,7 +264,7 @@ class Transformer(nn.Module):
         n_params = sum(p.numel() for p in self.parameters())
         # 如果不统计 embedding 的参数，就减去
         if non_embedding:
-            n_params -= self.transformer.wpe.weight.numel()
+            n_params -= self.transformer.wte.weight.numel()
         return n_params
 
     '''初始化权重'''

docs/chapter2/第二章 Transformer架构.md

@@ -6,24 +6,24 @@
 
 随着 NLP 从统计机器学习向深度学习迈进，作为 NLP 核心问题的文本表示方法也逐渐从统计学习向深度学习迈进。正如我们在第一章所介绍的，文本表示从最初的通过统计学习模型进行计算的向量空间模型、语言模型，通过 Word2Vec 的单层神经网络进入到通过神经网络学习文本表示的时代。但是，从 计算机视觉（Computer Vision，CV）为起源发展起来的神经网络，其核心架构有三种：
 
-- 前馈神经网络（Feedforward Neural Network，FNN），即每一层的神经元都和上下两层的每一个神经元完全连接，如图2.1所示:
+- 前馈神经网络（Feedforward Neural Network，FNN），即数据从输入层单向流动到输出层，无循环结构，各层之间通过全连接或特定方式传递信息，其中多层感知机（Multi-Layer Perceptron，MLP）是最常见的形式，每一层的神经元都和上下两层的每一个神经元完全连接，如图2.1所示:
 
 <div align="center">
-  <img src="../images/2-figures/1-0.png" alt="图片描述" width="90%"/>
-  <p>图2.1 前馈神经网络</p>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/2-figures/1-0.png" alt="图片描述" width="90%"/>
+  <p>图2.1 全连接神经网络</p>
 </div>
 
-- 卷积神经网络（Convolutional Neural Network，CNN），即训练参数量远小于前馈神经网络的卷积层来进行特征提取和学习，如图2.2所示:
+- 卷积神经网络（Convolutional Neural Network，CNN），即训练参数量远小于全连接神经网络的卷积层来进行特征提取和学习，如图2.2所示:
 
 <div align="center">
-  <img src="../images/2-figures/1-1.png" alt="图片描述" width="90%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/2-figures/1-1.png" alt="图片描述" width="90%"/>
   <p>图2.2 卷积神经网络</p>
 </div>
 
 - 循环神经网络（Recurrent Neural Network，RNN），能够使用历史信息作为输入、包含环和自重复的网络，如图2.3所示:
 
 <div align="center">
-  <img src="../images/2-figures/1-2.png" alt="图片描述" width="90%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/2-figures/1-2.png" alt="图片描述" width="90%"/>
   <p>图2.3 循环神经网络</p>
 </div>
 
@@ -87,7 +87,7 @@ $$
 
 根据词向量的定义，语义相似的两个词对应的词向量的点积应该大于0，而语义不相似的词向量点积应该小于0。
 
-那么，我们就可以用点积来计算词之间的相似度。假设我们的 Query 为“fruit”，对应的词向量为 $q$；我们的 Key 对应的词向量为 $k = [v_{apple} v_{banana} v_{chair}]$,则我们可以计算 Query 和每一个键的相似程度：
+那么，我们就可以用点积来计算词之间的相似度。假设我们的 Query 为“fruit”，对应的词向量为 $q$ ；我们的 Key 对应的词向量为 $k = [v_{apple} v_{banana} v_{chair}]$ ,则我们可以计算 Query 和每一个键的相似程度：
 
 $$
 x = qK^T
@@ -155,7 +155,7 @@ def attention(query, key, value, dropout=None):
 
 但是，在我们的实际应用中，我们往往只需要计算 Query 和 Key 之间的注意力结果，很少存在额外的真值 Value。也就是说，我们其实只需要拟合两个文本序列。​在经典的 注意力机制中，Q 往往来自于一个序列，K 与 V 来自于另一个序列，都通过参数矩阵计算得到，从而可以拟合这两个序列之间的关系。例如在 Transformer 的 Decoder 结构中，Q 来自于 Decoder 的输入，K 与 V 来自于 Encoder 的输出，从而拟合了编码信息与历史信息之间的关系，便于综合这两种信息实现未来的预测。
 
-​但在 Transformer 的 Encoder 结构中，使用的是 注意力机制的变种 —— 自注意力（self-attention，自注意力）机制。所谓自注意力，即是计算本身序列中每个元素都其他元素的注意力分布，即在计算过程中，Q、K、V 都由同一个输入通过不同的参数矩阵计算得到。在 Encoder 中，Q、K、V 分别是输入对参数矩阵 $W_q、W_k、W_v$ 做积得到，从而拟合输入语句中每一个 token 对其他所有 token 的关系。
+​但在 Transformer 的 Encoder 结构中，使用的是 注意力机制的变种 —— 自注意力（self-attention，自注意力）机制。所谓自注意力，即是计算本身序列中每个元素对其他元素的注意力分布，即在计算过程中，Q、K、V 都由同一个输入通过不同的参数矩阵计算得到。在 Encoder 中，Q、K、V 分别是输入对参数矩阵 $W_q、W_k、W_v$ 做积得到，从而拟合输入语句中每一个 token 对其他所有 token 的关系。
 
 通过自注意力机制，我们可以找到一段文本中每一个 token 与其他所有 token 的相关关系大小，从而建模文本之间的依赖关系。​在代码中的实现，self-attention 机制其实是通过给 Q、K、V 的输入传入同一个参数实现的：
 
@@ -187,7 +187,7 @@ attention(x, x, x)
     <BOS>    I   【MASK】 【MASK】【MASK】
     <BOS>    I     like  【MASK】【MASK】
     <BOS>    I     like    you  【MASK】
-    <BoS>    I     like    you   </EOS>
+    <BOS>    I     like    you   </EOS>
 
 在每一行输入中，模型仍然是只看到前面的 token，预测下一个 token。但是注意，上述输入不再是串行的过程，而可以一起并行地输入到模型中，模型只需要每一个样本根据未被遮蔽的 token 来预测下一个 token 即可，从而实现了并行的语言模型。
 
@@ -222,7 +222,7 @@ scores = F.softmax(scores.float(), dim=-1).type_as(xq)
 在原论文中，作者也通过实验证实，多头注意力计算中，每个不同的注意力头能够拟合语句中的不同信息，如图2.4所示：
 
 <div align="center">
-  <img src="../images/2-figures/1-3.jpeg" alt="图片描述" width="90%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/2-figures/1-3.jpeg" alt="图片描述" width="90%"/>
   <p>图2.4 多头注意力机制</p>
 </div>
 
@@ -252,55 +252,52 @@ class MultiHeadAttention(nn.Module):
         # args: 配置对象
         super().__init__()
         # 隐藏层维度必须是头数的整数倍，因为后面我们会将输入拆成头数个矩阵
-        assert args.n_embd % args.n_heads == 0
-        # 模型并行处理大小，默认为1。
-        model_parallel_size = 1
-        # 本地计算头数，等于总头数除以模型并行处理大小。
-        self.n_local_heads = args.n_heads // model_parallel_size
+        assert args.dim % args.n_heads == 0
         # 每个头的维度，等于模型维度除以头的总数。
         self.head_dim = args.dim // args.n_heads
+        self.n_heads = args.n_heads
 
-        # Wq, Wk, Wv 参数矩阵，每个参数矩阵为 n_embd x n_embd
+        # Wq, Wk, Wv 参数矩阵，每个参数矩阵为 n_embd x dim
         # 这里通过三个组合矩阵来代替了n个参数矩阵的组合，其逻辑在于矩阵内积再拼接其实等同于拼接矩阵再内积，
         # 不理解的读者可以自行模拟一下，每一个线性层其实相当于n个参数矩阵的拼接
-        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
-        self.wk = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
-        self.wv = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
-        # 输出权重矩阵，维度为 n_embd x n_embd（head_dim = n_embeds / n_heads）
-        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+        self.wq = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(args.n_embd, self.n_heads * self.head_dim, bias=False)
+        # 输出权重矩阵，维度为 dim x dim（head_dim = dim / n_heads）
+        self.wo = nn.Linear(self.n_heads * self.head_dim, args.dim, bias=False)
         # 注意力的 dropout
         self.attn_dropout = nn.Dropout(args.dropout)
         # 残差连接的 dropout
         self.resid_dropout = nn.Dropout(args.dropout)
-         
+        self.is_causal = is_causal
+
         # 创建一个上三角矩阵，用于遮蔽未来信息
         # 注意，因为是多头注意力，Mask 矩阵比之前我们定义的多一个维度
         if is_causal:
-           mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
-           mask = torch.triu(mask, diagonal=1)
-           # 注册为模型的缓冲区
-           self.register_buffer("mask", mask)
+            mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
+            mask = torch.triu(mask, diagonal=1)
+            # 注册为模型的缓冲区
+            self.register_buffer("mask", mask)
 
     def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
 
         # 获取批次大小和序列长度，[batch_size, seq_len, dim]
         bsz, seqlen, _ = q.shape
 
-        # 计算查询（Q）、键（K）、值（V）,输入通过参数矩阵层，维度为 (B, T, n_embed) x (n_embed, n_embed) -> (B, T, n_embed)
+        # 计算查询（Q）、键（K）、值（V）,输入通过参数矩阵层，维度为 (B, T, n_embed) x (n_embed, dim) -> (B, T, dim)
         xq, xk, xv = self.wq(q), self.wk(k), self.wv(v)
 
-        # 将 Q、K、V 拆分成多头，维度为 (B, T, n_head, C // n_head)，然后交换维度，变成 (B, n_head, T, C // n_head)
+        # 将 Q、K、V 拆分成多头，维度为 (B, T, n_head, dim // n_head)，然后交换维度，变成 (B, n_head, T, dim // n_head)
         # 因为在注意力计算中我们是取了后两个维度参与计算
         # 为什么要先按B*T*n_head*C//n_head展开再互换1、2维度而不是直接按注意力输入展开，是因为view的展开方式是直接把输入全部排开，
         # 然后按要求构造，可以发现只有上述操作能够实现我们将每个头对应部分取出来的目标
-        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
-        xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
-        xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+        xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_heads, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_heads, self.head_dim)
         xq = xq.transpose(1, 2)
         xk = xk.transpose(1, 2)
         xv = xv.transpose(1, 2)
 
-
         # 注意力计算
         # 计算 QK^T / sqrt(d_k)，维度为 (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
         scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
@@ -317,7 +314,7 @@ class MultiHeadAttention(nn.Module):
         output = torch.matmul(scores, xv)
 
         # 恢复时间维度并合并头。
-        # 将多头的结果拼接起来, 先交换维度为 (B, T, n_head, C // n_head)，再拼接成 (B, T, n_head * C // n_head)
+        # 将多头的结果拼接起来, 先交换维度为 (B, T, n_head, dim // n_head)，再拼接成 (B, T, n_head * dim // n_head)
         # contiguous 函数用于重新开辟一块新内存存储，因为Pytorch设置先transpose再view会报错，
         # 因为view直接基于底层存储得到，然而transpose并不会改变底层存储，因此需要额外存储
         output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
@@ -326,7 +323,6 @@ class MultiHeadAttention(nn.Module):
         output = self.wo(output)
         output = self.resid_dropout(output)
         return output
-
 ```
 
 ## 2.2 Encoder-Decoder
@@ -337,7 +333,7 @@ class MultiHeadAttention(nn.Module):
 
 ### 2.2.1 Seq2Seq 模型
 
-Seq2Seq，即序列到序列，是一种经典 NLP 任务。具体而言，是指模型输入的是一个自然语言序列 $input = (x_1, x_2, x_3...x_n)$，输出的是一个可能不等长的自然语言序列 $output = (y_1, y_2, y_3...y_m)$。事实上，Seq2Seq 是 NLP 最经典的任务，几乎所有的 NLP 任务都可以视为 Seq2Seq 任务。例如文本分类任务，可以视为输出长度为 1 的目标序列（如在上式中 $m$ = 1）；词性标注任务，可以视为输出与输入序列等长的目标序列（如在上式中 $m$ = $n$）。
+Seq2Seq，即序列到序列，是一种经典 NLP 任务。具体而言，是指模型输入的是一个自然语言序列 $input = (x_1, x_2, x_3...x_n)$ ，输出的是一个可能不等长的自然语言序列 $output = (y_1, y_2, y_3...y_m)$ 。事实上，Seq2Seq 是 NLP 最经典的任务，几乎所有的 NLP 任务都可以视为 Seq2Seq 任务。例如文本分类任务，可以视为输出长度为 1 的目标序列（如在上式中 $m$ = 1）；词性标注任务，可以视为输出与输入序列等长的目标序列（如在上式中 $m$ = $n$ ）。
 
 机器翻译任务即是一个经典的 Seq2Seq 任务，例如，我们的输入可能是“今天天气真好”，输出是“Today is a good day.”。Transformer 是一个经典的 Seq2Seq 模型，即模型的输入为文本序列，输出为另一个文本序列。事实上，Transformer 一开始正是应用在机器翻译任务上的。
 
@@ -346,7 +342,7 @@ Seq2Seq，即序列到序列，是一种经典 NLP 任务。具体而言，是
 Transformer 中的 Encoder，就是用于上述的编码过程；Decoder 则用于上述的解码过程。Transformer 结构，如图2.5所示：
 
 <div align="center">
-  <img src="../images/2-figures/2-0.jpg" alt="图片描述" width="90%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/2-figures/2-0.jpg" alt="图片描述" width="90%"/>
   <p>图2.5 编码器-解码器结构</p>
 </div>
 
@@ -356,7 +352,7 @@ Transformer 由 Encoder 和 Decoder 组成，每一个 Encoder（Decoder）又
 
 ### 2.2.2 前馈神经网络
 
-前馈神经网络（Feed Forward Neural Network，下简称 FFN），也就是我们在上一节提过的每一层的神经元都和上下两层的每一个神经元完全连接的网络结构。每一个 Encoder Layer 都包含一个上文讲的注意力机制和一个前馈神经网络。前馈神经网络的实现是较为简单的：
+前馈神经网络（Feed Forward Neural Network，下简称 FNN），也就是我们在上一节提过的每一层的神经元都和上下两层的每一个神经元完全连接的网络结构。每一个 Encoder Layer 都包含一个上文讲的注意力机制和一个前馈神经网络。前馈神经网络的实现是较为简单的：
 
 ```python
 class MLP(nn.Module):
@@ -373,13 +369,12 @@ class MLP(nn.Module):
     def forward(self, x):
         # 前向传播函数
         # 首先，输入x通过第一层线性变换和RELU激活函数
-        # 然后，结果乘以输入x通过第三层线性变换的结果
         # 最后，通过第二层线性变换和dropout层
         return self.dropout(self.w2(F.relu(self.w1(x))))
     
 ```
 
-注意，Transformer 的前馈神经网络是由两个线性层中间加一个 RELU 激活函数组成的，以及前馈神经网络还加入了一个 Dropout 层来防止过拟合。
+注意，Transformer 的前馈神经网络是由两个线性层中间加一个 RELU 激活函数组成的，以及前馈神经网络还加入了一个 Dropout 层来防止过拟合。Dropout 层只在训练时开启，推理/测试阶段关闭，所以许多Transformer结构示意图中不会画出该层。
 
 ### 2.2.3 层归一化
 
@@ -393,7 +388,7 @@ $$
 \mu_j = \frac{1}{m}\sum^{m}_{i=1}Z_j^{i}
 $$
 
-其中，$Z_j^{i}$ 是样本 i 在第 j 个维度上的值，m 就是 mini-batch 的大小。
+其中， $Z_j^{i}$ 是样本 i 在第 j 个维度上的值，m 就是 mini-batch 的大小。
 
 再计算样本的方差：
 
@@ -424,18 +419,18 @@ $$
 class LayerNorm(nn.Module):
     ''' Layer Norm 层'''
     def __init__(self, features, eps=1e-6):
-	super(LayerNorm, self).__init__()
-    # 线性矩阵做映射
-	self.a_2 = nn.Parameter(torch.ones(features))
-	self.b_2 = nn.Parameter(torch.zeros(features))
-	self.eps = eps
+        super().__init__()
+        # 线性矩阵做映射
+        self.a_2 = nn.Parameter(torch.ones(features))
+        self.b_2 = nn.Parameter(torch.zeros(features))
+        self.eps = eps
 	
     def forward(self, x):
-	# 在统计每个样本所有维度的值，求均值和方差
-	mean = x.mean(-1, keepdim=True) # mean: [bsz, max_len, 1]
-	std = x.std(-1, keepdim=True) # std: [bsz, max_len, 1]
-    # 注意这里也在最后一个维度发生了广播
-	return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
+        # 在统计每个样本所有维度的值，求均值和方差
+        mean = x.mean(-1, keepdim=True) # mean: [bsz, max_len, 1]
+        std = x.std(-1, keepdim=True) # std: [bsz, max_len, 1]
+        # 注意这里也在最后一个维度发生了广播
+        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
 ```
 注意，在我们上文实现的 Layer Norm 层中，有两个线性矩阵进行映射。
 
@@ -479,7 +474,7 @@ class EncoderLayer(nn.Module):
         # Encoder 不需要掩码，传入 is_causal=False
         self.attention = MultiHeadAttention(args, is_causal=False)
         self.fnn_norm = LayerNorm(args.n_embd)
-        self.feed_forward = MLP(args)
+        self.feed_forward = MLP(args.dim, args.dim, args.dropout)
 
     def forward(self, x):
         # Layer Norm
@@ -529,7 +524,7 @@ class DecoderLayer(nn.Module):
         self.attention = MultiHeadAttention(args, is_causal=False)
         self.ffn_norm = LayerNorm(args.n_embd)
         # 第三个部分是 MLP
-        self.feed_forward = MLP(args)
+        self.feed_forward = MLP(args.dim, args.dim, args.dropout)
 
     def forward(self, x, enc_out):
         # Layer Norm
@@ -540,7 +535,7 @@ class DecoderLayer(nn.Module):
         norm_x = self.attention_norm_2(x)
         h = x + self.attention.forward(norm_x, enc_out, enc_out)
         # 经过前馈神经网络
-        out = h + self.feed_forward.forward(self.fnn_norm(h))
+        out = h + self.feed_forward.forward(self.ffn_norm(h))
         return out
 ```
 
@@ -568,7 +563,7 @@ class Decoder(nn.Module):
 
 在前两章，我们分别深入剖析了 Attention 机制和 Transformer 的核心——Encoder、Decoder 结构，接下来，我们就可以基于上一章实现的组件，搭建起一个完整的 Transformer 模型。
 
-### 2.3.1 Embeddng 层
+### 2.3.1 Embedding 层
 
 正如我们在第一章所讲过的，在 NLP 任务中，我们往往需要将自然语言的输入转化为机器可以处理的向量。在深度学习中，承担这个任务的组件就是 Embedding 层。
 
@@ -590,7 +585,7 @@ output: 2
 因此，Embedding 层的输入往往是一个形状为 （batch_size，seq_len，1）的矩阵，第一个维度是一次批处理的数量，第二个维度是自然语言序列的长度，第三个维度则是 token 经过 tokenizer 转化成的 index 值。例如，对上述输入，Embedding 层的输入会是：
 
 ```
-[[0,1,2]]
+[[[0],[1],[2]]]
 ```
 
 其 batch_size 为1，seq_len 为3，转化出来的 index 如上。
@@ -616,18 +611,38 @@ PE(pos, 2i) = sin(pos/10000^{2i/d_{model}})\\
 PE(pos, 2i+1) = cos(pos/10000^{2i/d_{model}})
 $$
 
-​上式中，pos 为 token 在句子中的位置，2i 和 2i+1 则是指示了 token 是奇数位置还是偶数位置，从上式中我们可以看出对于奇数位置的 token 和偶数位置的 token，Transformer 采用了不同的函数进行编码。
+上式中，pos 为 token 在句子中的位置，2i 和 2i+1 则指示了位置编码向量的维度索引是奇数还是偶数，从上式中我们可以看出对于奇数维度和偶数维度，Transformer 采用了不同的函数进行编码。
 
-我们以一个简单的例子来说明位置编码的计算过程：假如我们输入的是一个长度为 4 的句子"I like to code"，我们可以得到下面的词向量矩阵$\rm x$，其中每一行代表的就是一个词向量，$\rm x_0=[0.1,0.2,0.3,0.4]$对应的就是“I”的词向量，它的pos就是为0，以此类推，第二行代表的是“like”的词向量，它的pos就是1：
+我们以一个简单的例子来说明位置编码的计算过程：假如我们输入的是一个长度为 4 的句子"I like to code"，我们可以得到下面的词向量矩阵 $\rm x$ ，其中每一行代表的就是一个词向量， $\rm x_0=[0.1,0.2,0.3,0.4]$ 对应的就是“I”的词向量，它的pos就是为0，以此类推，第二行代表的是“like”的词向量，它的pos就是1：
 
 $$
-\rm x = \begin{bmatrix} 0.1 & 0.2 & 0.3 & 0.4 \\ 0.2 & 0.3 & 0.4 & 0.5 \\ 0.3 & 0.4 & 0.5 & 0.6 \\ 0.4 & 0.5 & 0.6 & 0.7 \end{bmatrix}
+\rm x = \begin{bmatrix} 
+0.1 & 0.2 & 0.3 & 0.4 \\   
+0.2 & 0.3 & 0.4 & 0.5 \\    
+0.3 & 0.4 & 0.5 & 0.6 \\    
+0.4 & 0.5 & 0.6 & 0.7
+\end{bmatrix}
 $$
 
 ​则经过位置编码后的词向量为：
 
 $$
-\rm x_{PE} = \begin{bmatrix} 0.1 & 0.2 & 0.3 & 0.4 \\ 0.2 & 0.3 & 0.4 & 0.5 \\ 0.3 & 0.4 & 0.5 & 0.6 \\ 0.4 & 0.5 & 0.6 & 0.7 \end{bmatrix} + \begin{bmatrix} \sin(\frac{0}{10000^0}) & \cos(\frac{0}{10000^0}) & \sin(\frac{0}{10000^{2/4}}) & \cos(\frac{0}{10000^{2/4}}) \\ \sin(\frac{1}{10000^0}) & \cos(\frac{1}{10000^0}) & \sin(\frac{1}{10000^{2/4}}) & \cos(\frac{1}{10000^{2/4}}) \\ \sin(\frac{2}{10000^0}) & \cos(\frac{2}{10000^0}) & \sin(\frac{2}{10000^{2/4}}) & \cos(\frac{2}{10000^{2/4}}) \\ \sin(\frac{3}{10000^0}) & \cos(\frac{3}{10000^0}) & \sin(\frac{3}{10000^{2/4}}) & \cos(\frac{3}{10000^{2/4}}) \end{bmatrix} = \begin{bmatrix} 0.1 & 1.2 & 0.3 & 1.4 \\ 1.041 & 0.84 & 0.41 & 1.49 \\ 1.209 & -0.016 & 0.52 & 1.59 \\ 0.541 & -0.489 & 0.895 & 1.655 \end{bmatrix}
+\rm x_{PE} = \begin{bmatrix} 
+0.1 & 0.2 & 0.3 & 0.4 \\     
+0.2 & 0.3 & 0.4 & 0.5 \\    
+0.3 & 0.4 & 0.5 & 0.6 \\     
+0.4 & 0.5 & 0.6 & 0.7 
+\end{bmatrix} + \begin{bmatrix}
+\sin(\frac{0}{10000^0}) & \cos(\frac{0}{10000^0}) & \sin(\frac{0}{10000^{2/4}}) & \cos(\frac{0}{10000^{2/4}}) \\ 
+\sin(\frac{1}{10000^0}) & \cos(\frac{1}{10000^0}) & \sin(\frac{1}{10000^{2/4}}) & \cos(\frac{1}{10000^{2/4}}) \\ 
+\sin(\frac{2}{10000^0}) & \cos(\frac{2}{10000^0}) & \sin(\frac{2}{10000^{2/4}}) & \cos(\frac{2}{10000^{2/4}}) \\ 
+\sin(\frac{3}{10000^0}) & \cos(\frac{3}{10000^0}) & \sin(\frac{3}{10000^{2/4}}) & \cos(\frac{3}{10000^{2/4}}) 
+\end{bmatrix} = \begin{bmatrix} 
+0.1 & 1.2 & 0.3 & 1.4 \\ 
+1.041 & 0.84 & 0.41 & 1.49 \\ 
+1.209 & -0.016 & 0.52 & 1.59 \\ 
+0.541 & -0.489 & 0.895 & 1.655 
+\end{bmatrix}
 $$
 
 我们可以使用如下的代码来获取上述例子的位置编码：
@@ -671,13 +686,13 @@ $$
 \begin{equation}\tilde{f}(\cdots,\boldsymbol{x}_m,\cdots,\boldsymbol{x}_n,\cdots)=f(\cdots,\boldsymbol{x}_m + \boldsymbol{p}_m,\cdots,\boldsymbol{x}_n + \boldsymbol{p}_n,\cdots)\end{equation}
 $$
 
-这里加上的 $p_m$，$p_n$ 就是位置编码。接下来我们将 $f(...,x_m+p_m,...,x_n+p_n)$ 在 m,n 两个位置上做泰勒展开：
+这里加上的 $p_m$， $p_n$ 就是位置编码。接下来我们将 $f(...,x_m+p_m,...,x_n+p_n)$ 在 m,n 两个位置上做泰勒展开：
 
 $$
 \begin{equation}\tilde{f}\approx f + \boldsymbol{p}_m^{\top} \frac{\partial f}{\partial \boldsymbol{x}_m} + \boldsymbol{p}_n^{\top} \frac{\partial f}{\partial \boldsymbol{x}_n} + \frac{1}{2}\boldsymbol{p}_m^{\top} \frac{\partial^2 f}{\partial \boldsymbol{x}_m^2}\boldsymbol{p}_m + \frac{1}{2}\boldsymbol{p}_n^{\top} \frac{\partial^2 f}{\partial \boldsymbol{x}_n^2}\boldsymbol{p}_n + \underbrace{\boldsymbol{p}_m^{\top} \frac{\partial^2 f}{\partial \boldsymbol{x}_m \partial \boldsymbol{x}_n}\boldsymbol{p}_n}_{\boldsymbol{p}_m^{\top} \boldsymbol{\mathcal{H}} \boldsymbol{p}_n}\end{equation}
 $$
 
-可以看到第1项与位置无关，2～5项仅依赖单一位置，第6项（f 分别对 m、n 求偏导）与两个位置有关，所以我们希望第六项（$p_m^THp_n$）表达相对位置信息，即求一个函数 g 使得:
+可以看到第1项与位置无关，2～5项仅依赖单一位置，第6项（f 分别对 m、n 求偏导）与两个位置有关，所以我们希望第六项（ $p_m^THp_n$ ）表达相对位置信息，即求一个函数 g 使得:
 
 $$
 p_m^THp_n = g(m-n)
@@ -732,7 +747,7 @@ $$
 上述​编码结果，如图2.6所示：
 
 <div align="center">
-  <img src="../images/2-figures/3-0.png" alt="图片描述" width="90%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/2-figures/3-0.png" alt="图片描述" width="90%"/>
   <p>图2.6 编码结果</p>
 </div>
 
@@ -747,7 +762,7 @@ class PositionalEncoding(nn.Module):
     def __init__(self, args):
         super(PositionalEncoding, self).__init__()
         # Dropout 层
-        self.dropout = nn.Dropout(p=args.dropout)
+        # self.dropout = nn.Dropout(p=args.dropout)
 
         # block size 是序列的最大长度
         pe = torch.zeros(args.block_size, args.n_embd)
@@ -765,7 +780,7 @@ class PositionalEncoding(nn.Module):
     def forward(self, x):
         # 将位置编码加到 Embedding 结果上
         x = x + self.pe[:, : x.size(1)].requires_grad_(False)
-        return self.dropout(x)
+        return x
 ```
 
 ### 2.3.3 一个完整的 Transformer
@@ -773,10 +788,12 @@ class PositionalEncoding(nn.Module):
 上述所有组件，再按照下图的 Tranfromer 结构拼接起来就是一个完整的 Transformer 模型了，如图2.7所示：
 
 <div align="center">
-  <img src="../images/2-figures/3-1.png" alt="图片描述" width="80%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/2-figures/3-1.png" alt="图片描述" width="80%"/>
   <p>图2.7 Transformer 模型结构</p>
 </div>
 
+但需要注意的是，上图是原论文《Attention is all you need》配图，LayerNorm 层放在了 Attention 层后面，也就是“Post-Norm”结构，但在其发布的源代码中，LayerNorm 层是放在 Attention 层前面的，也就是“Pre Norm”结构。考虑到目前 LLM 一般采用“Pre-Norm”结构（可以使 loss 更稳定），本文在实现时采用“Pre-Norm”结构。
+
 如图，经过 tokenizer 映射后的输出先经过 Embedding 层和 Positional Embedding 层编码，然后进入上一节讲过的 N 个 Encoder 和 N 个 Decoder（在 Transformer 原模型中，N 取为6），最后经过一个线性层和一个 Softmax 层就得到了最终输出。
 
 基于之前所实现过的组件，我们实现完整的 Transformer 模型：
@@ -812,7 +829,7 @@ class Transformer(nn.Module):
         n_params = sum(p.numel() for p in self.parameters())
         # 如果不统计 embedding 的参数，就减去
         if non_embedding:
-            n_params -= self.transformer.wpe.weight.numel()
+            n_params -= self.transformer.wte.weight.numel()
         return n_params
 
     '''初始化权重'''

docs/chapter3/第三章 预训练语言模型.md

@@ -26,7 +26,7 @@ BERT 是一个统一了多种思想的预训练模型。其所沿承的核心思
 BERT 的模型架构是取了 Transformer 的 Encoder 部分堆叠而成，其主要结构如图3.1所示：
 
 <div align="center">
-  <img src="../images/3-figures/1-0.png" alt="图片描述" width="100%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/1-0.png" alt="图片描述" width="100%"/>
   <p>图3.1 BERT 模型结构</p>
 </div>
 
@@ -35,30 +35,33 @@ BERT 是针对于 NLU 任务打造的预训练模型，其输入一般是文本
 模型整体既是由 Embedding、Encoder 加上 prediction_heads 组成：
 
 <div align="center">
-  <img src="../images/3-figures/1-1.png" alt="图片描述" width="70%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/1-1.png" alt="图片描述" width="70%"/>
   <p>图3.2 BERT 模型简略结构</p>
 </div>
 
-输入的文本序列会首先通过 tokenizer（分词器） 转化成 input_ids（基本每一个模型在 tokenizer 的操作都类似，可以参考 Transformer 的 tokenizer 机制，后文不再赘述），然后进入 Embedding 层转化为特定维度的 hidden_states，再经过 Encoder 块。Encoder 块中是对叠起来的 N 层 Encoder Layer，BERT 有两种规模的模型，分别是 base 版本（12层 Encoder Layer，768 的隐藏层维度，总参数量 110M），large 版本（24层 Encoder Layer，1024 的隐藏层维度，总参数量 340M）。通过Encoder 编码之后的最顶层 hidden_states 最后经过 prediction_heads 就得到了最后的类别概率，经过 Softmax 计算就可以计算出模型预测的类别。
+输入的文本序列会首先通过 tokenizer（分词器） 转化成 input_ids（基本每一个模型在 tokenizer 的操作都类似，可以参考 Transformer 的 tokenizer 机制，后文不再赘述），然后进入 Embedding 层转化为特定维度的 hidden_states，再经过 Encoder 块。Encoder 块中是堆叠起来的 N 层 Encoder Layer，BERT 有两种规模的模型，分别是 base 版本（12层 Encoder Layer，768 的隐藏层维度，总参数量 110M），large 版本（24层 Encoder Layer，1024 的隐藏层维度，总参数量 340M）。通过Encoder 编码之后的最顶层 hidden_states 最后经过 prediction_heads 就得到了最后的类别概率，经过 Softmax 计算就可以计算出模型预测的类别。
+
+
+> BERT 采用 WordPiece 作为分词方法。WordPiece 是一种基于统计的子词切分算法，其核心在于将单词拆解为子词（例如，"playing" -> ["play", "##ing"]）。其合并操作的依据是最大化语言模型的似然度。对于中文等非空格分隔的语言，通常将单个汉字作为原子分词单位（token）处理。
 
 prediction_heads 其实就是线性层加上激活函数，一般而言，最后一个线性层的输出维度和任务的类别数相等，如图3.3所示：
 
 <div align="center">
-  <img src="../images/3-figures/1-5.png" alt="图片描述" width="20%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/1-5.png" alt="图片描述" width="20%"/>
   <p>图3.3 prediction_heads 结构</p>
 </div>
 
 而每一层 Encoder Layer 都是和 Transformer 中的 Encoder Layer 结构类似的层，如图3.4所示：
 
 <div align="center">
-  <img src="../images/3-figures/1-2.png" alt="图片描述" width="40%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/1-2.png" alt="图片描述" width="40%"/>
   <p>图3.4 Encoder Layer 结构</p>
 </div>
 
 如图3.5所示，已经通过 Embedding 层映射的 hidden_states 进入核心的 attention 机制，然后通过残差连接的机制和原输入相加，再经过一层 Intermediate 层得到最终输出。Intermediate 层是 BERT 的特殊称呼，其实就是一个线性层加上激活函数：
 
 <div align="center">
-  <img src="../images/3-figures/1-3.png" alt="图片描述" width="40%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/1-3.png" alt="图片描述" width="40%"/>
   <p>图3.5 Intermediate 结构</p>
 </div>
 
@@ -71,12 +74,14 @@ GELU 的核心思路为将随机正则的思想引入激活函数，通过输入
 BERT 的 注意力机制和 Transformer 中 Encoder 的 自注意力机制几乎完全一致，但是 BERT 将相对位置编码融合在了注意力机制中，将相对位置编码同样视为可训练的权重参数，如图3.6所示：
 
 <div align="center">
-  <img src="../images/3-figures/1-4.png" alt="图片描述" width="40%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/1-4.png" alt="图片描述" width="40%"/>
   <p>图3.6 BERT 注意力机制结构</p>
 </div>
 
 如图，BERT 的注意力计算过程和 Transformer 的唯一差异在于，在完成注意力分数的计算之后，先通过 Position Embedding 层来融入相对位置信息。这里的 Position Embedding 层，其实就是一层线性矩阵。通过可训练的参数来拟合相对位置，相对而言比 Transformer 使用的绝对位置编码 Sinusoidal 能够拟合更丰富的相对位置信息，但是，这样也增加了不少模型参数，同时完全无法处理超过模型训练长度的输入（例如，对 BERT 而言能处理的最大上下文长度是 512 个 token）。
 
+注：原始 BERT（即论文提出）使用和 Transformer 一致的绝对位置编码，后续改进（包括 BERT 的各种变体）使用了上述相对位置编码，为帮助读者了解更全面的模型结构设计，此处选择了改进版 BERT。
+
 可以看出，BERT 的模型架构既是建立在 Transformer 的 Encoder 之上的，这也是为什么说 BERT 沿承了 Transformer 的思想。
 
 #### （3）预训练任务——MLM + NSP
@@ -161,7 +166,7 @@ RoBERTa 使用了更大量的无监督语料进行预训练，除去 BERT 所使
 
 #### （3）优化三：更大的 bpe 词表
 
-RoBERTa、BERT 和 Transformer 一样，都使用了 BPE 作为 Tokenizer 的编码策略。BPE，即 Byte Pair Encoding，字节对编码，是指以子词对作为分词的单位。例如，对“Hello World”这句话，可能会切分为“Hel，lo，Wor，ld”四个子词对。而对于以字为基本单位的中文，一般会按照 字节编码进行切分。例如，在 UTF-8 编码中，“我”会被编码为“E68891”，那么在 BPE 中可能就会切分成“E68”，“891”两个字词对。
+与 BERT 使用的 WordPiece 算法不同，RoBERTa 使用了 BPE 作为 Tokenizer 的编码策略。BPE，即 Byte Pair Encoding，字节对编码，是指以子词对作为分词的单位。例如，对“Hello World”这句话，可能会切分为“Hel，lo，Wor，ld”四个子词对。而对于以字为基本单位的中文，一般会按照字节编码进行切分。例如，在 UTF-8 编码中，“我”会被编码为“E68891”，那么在 BPE 中可能就会切分成“E68”，“891”两个字词对。
 
 一般来说，BPE 编码的词典越大，编码效果越好。当然，由于 Embedding 层就是把 token 从词典空间映射到隐藏空间（也就是说 Embedding 的形状为 (vocab_size, hidden_size)，越大的词表也会带来模型参数的增加。
 
@@ -175,11 +180,11 @@ BERT 原始的 BPE 词表大小为 30K，RoBERTa 选择了 50K 大小的词表
 
 #### （1）优化一：将 Embedding 参数进行分解
 
-BERT 等预训练模型具有远超传统神经网络的参数量，如前所述，BERT-large 具有 24层 Encoder Layer，1024 的隐藏层维度，总共参数量达 340M。而这其中，Embedding 层的参数矩阵维度为 $V*H$，此处的 V 为词表大小 30K，H 即为隐藏层大小 768，也就是 Embedding 层参数达到了 23M。而这样的设置还会带来一个更大的问题，即 Google 探索尝试搭建更宽（也就是隐藏层维度更大）的模型时发现，隐藏层维度的增加会带来 Embedding 层参数的巨大上升，如果把隐藏层维度增加到 2048，Embedding 层参数就会膨胀到 61M，这无疑是极大增加了模型的计算开销。
+BERT 等预训练模型具有远超传统神经网络的参数量，如前所述，BERT-large 具有 24层 Encoder Layer，1024 的隐藏层维度，总共参数量达 340M。而这其中，Embedding 层的参数矩阵维度为 $V*H$，此处的 V 为词表大小 30K，H 即为隐藏层大小 1024，也就是 Embedding 层参数达到了 30M。而这样的设置还会带来一个更大的问题，即 Google 探索尝试搭建更宽（也就是隐藏层维度更大）的模型时发现，隐藏层维度的增加会带来 Embedding 层参数的巨大上升，如果把隐藏层维度增加到 2048，Embedding 层参数就会膨胀到 61M，这无疑是极大增加了模型的计算开销。
 
 而从另一个角度看，Embedding 层输出的向量是我们对文本 token 的稠密向量表示，从 Word2Vec 的成功经验来看，这种词向量并不需要很大的维度，Word2Vec 仅使用了 100维大小就取得了很好的效果。因此，Embedding 层的输出也许不需要和隐藏层大小一致。
 
-因此，ALBERT 对 Embedding 层的参数矩阵进行了分解，让 Embedding 层的输出维度和隐藏层维度解绑，也就是在 Embedding 层的后面加入一个线性矩阵进行维度变换。ALBERT 设置了 Embedding 层的输出为 128，因此在 Embedding 层后面加入了一个 $128*768$ 的线性矩阵来将 Embedding 层的输出再升维到隐藏层大小。也就是说，Embedding 层的参数从 $V*H$ 降低到了 $V*E + E*H$，当 E 的大小远小于 H 时，该方法对 Embedding 层参数的优化就会很明显。
+因此，ALBERT 对 Embedding 层的参数矩阵进行了分解，让 Embedding 层的输出维度和隐藏层维度解绑，也就是在 Embedding 层的后面加入一个线性矩阵进行维度变换。ALBERT 设置了 Embedding 层的输出为 128，因此在 Embedding 层后面加入了一个 $128*1024$ 的线性矩阵来将 Embedding 层的输出再升维到隐藏层大小。也就是说，Embedding 层的参数从 $V*H$ 降低到了 $V*E + E*H$，当 E 的大小远小于 H 时，该方法对 Embedding 层参数的优化就会很明显。
 
 #### （2）优化二：跨层进行参数共享
 
@@ -230,14 +235,14 @@ T5 的大一统思想将不同的 NLP 任务如文本分类、问答、翻译等
 BERT 采用了 Encoder-Only 结构，只包含编码器部分；而 GPT 采用了 Decoder-Only 结构，只包含解码器部分。T5 则采用了 Encoder-Decoder 结构，其中编码器和解码器都是基于 Transformer 架构设计。编码器用于处理输入文本，解码器用于生成输出文本。编码器和解码器之间通过注意力机制进行信息交互，从而实现输入文本到输出文本的转换。其主要结构如图3.7所示：
 
 <div align="center">
-  <img src="../images/3-figures/2-1.png" alt="图片描述" width="100%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/2-1.png" alt="图片描述" width="100%"/>
   <p>图3.7 T5 模型详细结构</p>
 </div>
 
 如图3.8所示，从整体来看 T5 的模型结构包括 Tokenizer 部分和 Transformer 部分。Tokenizer 部分主要负责将输入文本转换为模型可接受的输入格式，包括分词、编码等操作。Transformer 部分又分为 EncoderLayers 和 DecoderLayers 两部分，他们分别由一个个小的 Block组成，每个 Block 包含了多头注意力机制、前馈神经网络和 Norm 层。Block 的设计可以使模型更加灵活，像乐高一样可以根据任务的复杂程度和数据集的大小来调整 Block 的数量和层数。
 
 <div align="center">
-  <img src="../images/3-figures/2-2.png" alt="图片描述" width="70%"/>
+  <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/2-2.png" alt="图片描述" width="70%"/>
   <p>图3.8 T5 模型整体结构</p>
 </div>
 
@@ -246,28 +251,28 @@ T5 模型的 Encoder 和 Decoder 部分都是基于 Transformer 架构设计的
 和 Encoder 不一样的是，在 Decoder 中还包含了 Encoder-Decoder Attention 结构，用于捕捉输入和输出序列之间的依赖关系。这两种 Attention 结构几乎完全一致，只有在位置编码和 Mask 机制上有所不同。如图3.9所示，Encoder 和 Decoder 的结构如下：
 
 <div align='center'>
-    <img src="../images/3-figures/2-3.png" alt="alt text" width="50%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/2-3.png" alt="alt text" width="50%">
     <p>图3.9 Encoder 和 Decoder</p>
 </div>
 
 T5 的 Self-Attention 机制和 BERT 的 Attention 机制是一样的，都是基于 Self-Attention 机制设计的。Self-Attention 机制是一种全局依赖关系建模方法，通过计算 Query、Key 和 Value 之间的相似度来捕捉输入序列中的全局依赖关系。Encoder-Decoder Attention 仅仅在位置编码和 Mask 机制上有所不同，主要是为了区分输入和输出序列。如图3.10所示，Self-Attention 结构如下：
 
 <div align='center'>
-    <img src="../images/3-figures/2-4.png" alt="alt text" width="50%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/2-4.png" alt="alt text" width="50%">
     </p>图3.10 Self-Attention 结构</p>
 </div>
 
-与原始 Transformer 模型不同，T5 模型的LayerNorm 采用了 RMSNorm，通过计算每个神经元的均方根（Root Mean Square）来归一化每个隐藏层的激活值。RMSNorm 的参数设置与Layer Normalization 相比更简单，只有一个可调参数，可以更好地适应不同的任务和数据集。RMSNorm函数可以用以下数学公式表示：
+与原始 Transformer 模型不同，T5 模型的LayerNorm 采用了 RMSNorm，通过计算每个神经元的均方根（Root Mean Square）来归一化每个隐藏层的激活值。RMSNorm 的参数设置与Layer Normalization 相比更简单，只有一个可学参数，可以更好地适应不同的任务和数据集。RMSNorm函数可以用以下数学公式表示：
 
 $$
-\text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{n}\sum_{i=1}^{n}w_i^2 + \epsilon}}
+\text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{n}\sum_{i=1}^{n}x_i^2 + \epsilon}} \cdot \gamma
 $$
 
 其中：
-- \( $x$ \) 是层的输入。
-- \( $w_i$ \) 代表层的权重。
-- \( $n$ \) 是权重的数量。
-- \( $\epsilon$ \) 是一个小常数，用于数值稳定性（以避免除以零的情况）。
+- $x_i$ 是输入向量的第 $i$ 个元素
+- $\gamma$ 是可学习的缩放参数
+- $n$ 是输入向量的维度数量
+- $\epsilon$ 是一个小常数，用于数值稳定性（以避免除以零的情况）
 
 这种归一化有助于通过确保权重的规模不会变得过大或过小来稳定学习过程，这在具有许多层的深度学习模型中特别有用。
 
@@ -298,7 +303,7 @@ T5通过大规模的文本数据进行预训练，然后在具体任务上进行
 我们可以通过图3.11，更加直观地理解 T5 的大一统思想：
 
 <div align='center'>
-    <img src="../images/3-figures/2-0.png" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/2-0.png" alt="alt text" width="90%">
     <p>图3.11 T5 的大一统思想</p>
 </div>
 
@@ -323,7 +328,7 @@ GPT，即 Generative Pre-Training Language Model，是由 OpenAI 团队于 2018
 #### （1） 模型架构——Decoder Only
 
 <div align='center'>
-    <img src="../images/3-figures/3-0.png" alt="alt text" width="100%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/3-0.png" alt="alt text" width="100%">
     <p>图3.12 GPT 模型结构</p>
 </div>
 
@@ -339,7 +344,7 @@ GPT，即 Generative Pre-Training Language Model，是由 OpenAI 团队于 2018
 
 #### （2）预训练任务——CLM
 
-Decoder-Only 的模型结构往往更适合于文本生成任务，因此，Decoder-Only 模型往往选择了最传统也最直接的预训练任务——因果语言模型，Casual Language Model，下简称 CLM。
+Decoder-Only 的模型结构往往更适合于文本生成任务，因此，Decoder-Only 模型往往选择了最传统也最直接的预训练任务——因果语言模型，Causal Language Model，下简称 CLM。
 
 CLM 可以看作 N-gram 语言模型的一个直接扩展。N-gram 语言模型是基于前 N 个 token 来预测下一个 token，CLM 则是基于一个自然语言序列的前面所有 token 来预测下一个 token，通过不断重复该过程来实现目标文本序列的生成。也就是说，CLM 是一个经典的补全形式。例如，CLM 的输入和输出可以是：
 
@@ -394,7 +399,7 @@ LLaMA模型是由Meta（前Facebook）开发的一系列大型预训练语言模
 与GPT系列模型一样，LLaMA模型也是基于Decoder-Only架构的预训练语言模型。LLaMA模型的整体结构与GPT系列模型类似，只是在模型规模和预训练数据集上有所不同。如图3.13是LLaMA模型的架构示意图：
 
 <div align='center'>
-    <img src="../images/3-figures/3-1.png" alt="alt text" width="100%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/3-1.png" alt="alt text" width="100%">
     <p>图3.13 LLaMA-3 模型结构</p>
 </div>
 
@@ -428,7 +433,7 @@ LLaMA模型是由Meta（前Facebook）开发的一系列大型预训练语言模
 - LLaMA-3支持8K长文本，并采用了编码效率更高的tokenizer，词表大小为128K。
 - 使用了超过15T token的预训练语料，是LLaMA-2的7倍多。
 
-LLaMA模型以其技术创新、多参数版本、大规模预训练和高效架构设计而著称。模型支持从7亿到数百亿不等的参数量，适应不同规模的应用需求。LLaMA-1以其开源性和优异性能迅速受到社区欢迎，而LLaMA-2和LLaMA-3进一步通过引入分组查询注意力机制和支持更长文本输入，显著提升了模型性能和应用范围。特别是LLaMA-3，通过采用128K词表大小的高效tokenizer和15T token的庞大训练数据，实现了在多语言和多任务处理上的重大进步。Meta对模型安全性和社区支持的持续关注，预示着LLaMA将继续作为AI技术发展的重要推动力，促进全球范围内的技术应用和创新。
+LLaMA模型以其技术创新、多参数版本、大规模预训练和高效架构设计而著称。模型支持从70亿到数百亿不等的参数量，适应不同规模的应用需求。LLaMA-1以其开源性和优异性能迅速受到社区欢迎，而LLaMA-2和LLaMA-3进一步通过引入分组查询注意力机制和支持更长文本输入，显著提升了模型性能和应用范围。特别是LLaMA-3，通过采用128K词表大小的高效tokenizer和15T token的庞大训练数据，实现了在多语言和多任务处理上的重大进步。Meta对模型安全性和社区支持的持续关注，预示着LLaMA将继续作为AI技术发展的重要推动力，促进全球范围内的技术应用和创新。
 
 ### 3.3.3 GLM
 
@@ -446,7 +451,7 @@ GLM 最初是由清华计算机系推出的一种通用语言模型基座，其
 
 2. 使用单个线性层实现最终 token 的预测，而不是使用 MLP；这样的结构更加简单也更加鲁棒，即减少了最终输出的参数量，将更大的参数量放在了模型本身；
 
-3. 激活函数从 ReLU 换成了 GeLUS。ReLU 是传统的激活函数，其核心计算逻辑为去除小于 0的传播，保留大于 0的传播；GeLUS 核心是对接近于 0的正向传播，做了一个非线性映射，保证了激活函数后的非线性输出，具有一定的连续性。
+3. 激活函数从 ReLU 换成了 GeLUs。ReLU 是传统的激活函数，其核心计算逻辑为去除小于 0的传播，保留大于 0的传播；GeLUs 核心是对接近于 0的正向传播，做了一个非线性映射，保证了激活函数后的非线性输出，具有一定的连续性。
 
 #### （2）预训练任务-GLM
 
@@ -460,7 +465,7 @@ GLM 通过优化一个自回归空白填充任务来实现 MLM 与 CLM 思想的
 通过将 MLM 与 CLM 思想相结合，既适配逐个 token 生成的生成类任务，也迫使模型从前后两个方向学习输入文本的隐含关系从而适配了理解类任务。使用 GLM 预训练任务产出的 GLM 模型，在一定程度上展现了其超出同体量 BERT 系模型的优越性能：
 
 <div align='center'>
-    <img src="../images/3-figures/3-2.png" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/3-2.png" alt="alt text" width="90%">
     <p>图3.14 alt text</p>
 </div>
 
@@ -479,7 +484,7 @@ ChatGLM3-6B 发布于 23年 10月，相对于二代在语义、数学、推理
 图3.15展示了 GLM 系列模型在基准集上的表现演进：
 
 <div align='center'>
-    <img src="../images/3-figures/3-3.png" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/3-figures/3-3.png" alt="alt text" width="90%">
     <p>图3.15 alt text</p>
 </div>
 

docs/chapter4/第四章 大语言模型.md

@@ -96,13 +96,13 @@ LLM 的强大能力也为其带来了跨模态的强大表现。随着 LLM 的
 在上一节，我们分析了 LLM 的定义及其特有的强大能力，通过更大规模的参数和海量的训练语料获得远超传统预训练模型的涌现能力，展现出强大的上下文学习、指令遵循及逐步推理能力，带来 NLP 领域的全新变革。那么，通过什么样的步骤，我们才可以训练出一个具有涌现能力的 LLM 呢？训练一个 LLM，与训练传统的预训练模型，又有什么区别？
 
 <div align='center'>
-    <img src="../images/4-figures/2-0.jpg" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/4-figures/2-0.jpg" alt="alt text" width="90%">
     <p>图4.1 训练 LLM 的三个阶段</p>
 </div>
 
 一般而言，训练一个完整的 LLM 需要经过图1中的三个阶段——Pretrain、SFT 和 RLHF。在这一节，我们将详细论述训练 LLM 的三个阶段，并分析每一个阶段的过程及其核心难点、注意事项，帮助读者们从理论上了解要训练一个 LLM，需要经过哪些步骤。
 
-### 4.2.2 Pretrain
+### 4.2.1 Pretrain
 
 Pretrain，即预训练，是训练 LLM 最核心也是工程量最大的第一步。LLM 的预训练和传统预训练模型非常类似，同样是使用海量无监督文本对随机初始化的模型参数进行训练。正如我们在第三章中所见，目前主流的 LLM 几乎都采用了 Decoder-Only 的类 GPT 架构（LLaMA 架构），它们的预训练任务也都沿承了 GPT 模型的经典预训练任务——因果语言模型（Causal Language Model，CLM）。
 
@@ -128,7 +128,7 @@ GPT-3|96|12288|96|175B|300B
 也正因如此，分布式训练框架也成为 LLM 训练必不可少的组成部分。分布式训练框架的核心思路是数据并行和模型并行。所谓数据并行，是指训练模型的尺寸可以被单个 GPU 内存容纳，但是由于增大训练的 batch_size 会增大显存开销，无法使用较大的 batch_size 进行训练；同时，训练数据量非常大，使用单张 GPU 训练时长难以接受。
 
 <div align='center'>
-    <img src="../images/4-figures/2-1.jpg" alt="alt text" width="60%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/4-figures/2-1.jpg" alt="alt text" width="60%">
     <p>图4.2 模型、数据并行</p>
 </div>
 
@@ -137,7 +137,7 @@ GPT-3|96|12288|96|175B|300B
 但是，当 LLM 扩大到上百亿参数，单张 GPU 内存往往就无法存放完整的模型参数。如图4.3所示，在这种情况下，可以将模型拆分到多个 GPU 上，每个 GPU 上存放不同的层或不同的部分，从而实现模型并行。
 
 <div align='center'>
-    <img src="../images/4-figures/2-2.jpg" alt="alt text" width="30%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/4-figures/2-2.jpg" alt="alt text" width="30%">
     <p>图4.3 模型并行</p>
 </div>
 
@@ -179,11 +179,11 @@ StackExchange|2.0%|78 GB
 
 目前，已有很多经过处理的高质量预训练语料和专用于预训练数据处理的框架。例如，有基于 LLaMA 思路收集、清洗的预训练数据集[RedPajama-1T](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T)，以及在 RedPajama 基础上进行筛选去重的[SlimPajama-627B](https://huggingface.co/datasets/cerebras/SlimPajama-627B/tree/main/train)数据集，实验证明高质量的 627B Slimpajama 数据集能够获得比 1T 的 RedPajama 数据集更好的效果。
 
-### 4.2.3 SFT
+### 4.2.2 SFT
 
 预训练是 LLM 强大能力的根本来源，事实上，LLM 所覆盖的海量知识基本都是源于预训练语料。LLM 的性能本身，核心也在于预训练的工作。但是，预训练赋予了 LLM 能力，却还需要第二步将其激发出来。经过预训练的 LLM 好像一个博览群书但又不求甚解的书生，对什么样的偏怪问题，都可以流畅地接出下文，但他偏偏又不知道问题本身的含义，只会“死板背书”。这一现象的本质是因为，LLM 的预训练任务就是经典的 CLM，也就是训练其预测下一个 token 的能力，在没有进一步微调之前，其无法与其他下游任务或是用户指令适配。
 
-因此，我们还需要第二步来教这个博览群书的学生如何去使用它的知识，也就是 SFT——Supervisor Finetune，有监督微调。所谓有监督微调，其实就是我们在第三章中讲过的预训练-微调中的微调，稍有区别的是，对于能力有限的传统预训练模型，我们需要针对每一个下游任务单独对其进行微调以训练模型在该任务上的表现。例如要解决文本分类问题，需要对 BERT 进行文本分类的微调；要解决实体识别的问题，就需要进行实体识别任务的微调。
+因此，我们还需要第二步来教这个博览群书的学生如何去使用它的知识，也就是 SFT（Supervised Fine-Tuning，有监督微调）。所谓有监督微调，其实就是我们在第三章中讲过的预训练-微调中的微调，稍有区别的是，对于能力有限的传统预训练模型，我们需要针对每一个下游任务单独对其进行微调以训练模型在该任务上的表现。例如要解决文本分类问题，需要对 BERT 进行文本分类的微调；要解决实体识别的问题，就需要进行实体识别任务的微调。
 
 而面对能力强大的 LLM，我们往往不再是在指定下游任务上构造有监督数据进行微调，而是选择训练模型的“通用指令遵循能力”，也就是一般通过`指令微调`的方式来进行 SFT。
 
@@ -196,7 +196,7 @@ StackExchange|2.0%|78 GB
 
 首先是指令数据量及覆盖范围。为了使 LLM 能够获得泛化的指令遵循能力，即能够在未训练的指令上表现良好，需要收集大量类别各异的用户指令和对应回复对 LLM 进行训练。一般来说，在单个任务上 500~1000 的训练样本就可以获得不错的微调效果。但是，为了让 LLM 获得泛化的指令遵循能力，在多种任务指令上表现良好，需要在训练数据集中覆盖多种类型的任务指令，同时也需要相对较大的训练数据量，表现良好的开源 LLM SFT 数据量一般在数 B token 左右。
 
-为提高 LLM 的泛化能力，指令数据集的覆盖范围自然是越大越好。但是，多种不同类型的指令数据之间的配比也是 LLM 训练的一大挑战。OpenAI 训练的 InstructGPT（即 ChatGPT 前身）使用了源自于用户使用其 API 的十种指令：
+为提高 LLM 的泛化能力，指令数据集的覆盖范围自然是越大越好。**但是**，多种不同类型的指令数据之间的配比也是 LLM 训练的一大挑战。OpenAI 训练的 InstructGPT（即 ChatGPT 前身）使用了源自于用户使用其 API 的十种指令：
 
 指令类型|占比
 -------|-----
@@ -290,14 +290,14 @@ StackExchange|2.0%|78 GB
 
 显然可知，第一种方式会丢失大量中间信息，第二种方式造成了大量重复计算，只有第三种方式是最合理的多轮对话构造。我们之所以可以以第三种方式来构造多轮对话样本，是因为 LLM 本质还是进行的 CLM 任务，进行单向注意力计算，因此在预测时会从左到右依次进行拟合，前轮的输出预测不会影响后轮的预测。目前，绝大部分 LLM 均使用了多轮对话的形式来进行 SFT。
 
-## 4.2.4 RLHF
+### 4.2.3 RLHF
 
 RLHF，全称是 Reinforcement Learning from Human Feedback，即人类反馈强化学习，是利用强化学习来训练 LLM 的关键步骤。相较于在 GPT-3 就已经初见雏形的 SFT，RLHF 往往被认为是 ChatGPT 相较于 GPT-3 的最核心突破。事实上，从功能上出发，我们可以将 LLM 的训练过程分成预训练与对齐（alignment）两个阶段。预训练的核心作用是赋予模型海量的知识，而所谓对齐，其实就是让模型与人类价值观一致，从而输出人类希望其输出的内容。在这个过程中，SFT 是让 LLM 和人类的指令对齐，从而具有指令遵循能力；而 RLHF 则是从更深层次令 LLM 和人类价值观对齐，令其达到安全、有用、无害的核心标准。
 
 如图4.4所示，ChatGPT 在技术报告中将对齐分成三个阶段，后面两个阶段训练 RM 和 PPO 训练，就是 RLHF 的步骤：
 
 <div align='center'>
-    <img src="../images/4-figures/2-3.png" alt="alt text" width="100%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/4-figures/2-3.png" alt="alt text" width="100%">
     <p>图4.4 ChatGPT 训练三个的阶段</p>
 </div>
 
@@ -331,7 +331,7 @@ RM，Reward Model，即奖励模型。RM 是用于拟合人类偏好，来给 LL
 在具体 PPO 训练过程中，会存在四个模型。如图4.5所示，两个 LLM 和两个 RM。两个 LLM 分别是进行微调、参数更新的 actor model 和不进行参数更新的 ref model，均是从 SFT 之后的 LLM 初始化的。两个 RM 分别是进行参数更新的 critic model 和不进行参数更新的 reward model，均是从上一步训练的 RM 初始化的。
 
 <div align='center'>
-    <img src="../images/4-figures/2-4.jpg" alt="alt text" width="100%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/4-figures/2-4.jpg" alt="alt text" width="100%">
     <p>图4.5 PPO 训练流程</p>
 </div>
 
@@ -339,9 +339,9 @@ RM，Reward Model，即奖励模型。RM 是用于拟合人类偏好，来给 LL
 
 1. 从 SFT 之后的 LLM 初始化两个模型分别作为 Actor Model 和 Ref Model；从训练的 RM 初始化两个模型分别作为 Reward Model 和 Critic Model；
 2. 输入一个 Prompt，Actor Model 和 Ref Model 分别就 Prompt 生成回复；
-3. Actor Response 和 Ref Response 计算 KL 散度：$r_{KL} = -\theta_{KL}D_{KL}(\pi_{PPO}(y|x)||\pi_{base}(y|x))$其中，$\pi_{PPO}(y|x)$即为 Actor Model 的输出，而 $\pi_{base}(y|x)$即为 Ref Model 的输出，$theta_{KL}D_{KL}$即是计算 KL 散度的方法；
+3. Actor Response 和 Ref Response 计算 KL 散度： $r_{KL} = -\theta_{KL}D_{KL}(\pi_{PPO}(y|x)||\pi_{base}(y|x))$ 其中， $\pi_{PPO}(y|x)$ 即为 Actor Model 的输出，而 $\pi_{base}(y|x)$ 即为 Ref Model 的输出， $\theta_{KL}D_{KL}$ 即是计算 KL 散度的方法；
 4. Actor Response 分别输入到 Reward Model 和 Critic Model 进行打分，其中，Reward Model 输出的是回复对应的标量奖励，Critic Model 还会输出累加奖励（即从i位置到最后的累积奖励）；
-5. 计算的 KL 散度、两个模型的打分均输入到奖励函数中，计算奖励：$loss = -(kl_{ctl}*r_{KL} + \gamma * V_{t+1} - V_{t})logP(A_t|V_t)$，这里的 $kl_{ctl}是控制 KL 散度对结果影响的权重参数，$\gamma$ 是控制下一个时间（也就是样本）打分对结果影响的权重参数，$V_t$ 是 Critic Model 的打分输出，$A_t$ 则是 Reward Model 的打分输出;
+5. 计算的 KL 散度、两个模型的打分均输入到奖励函数中，计算奖励： $loss = -(kl_{ctl} \cdot r_{KL} + \gamma \cdot V_{t+1} - V_{t}) \log P(A_t|V_t)$ ，这里的 $kl_{ctl}$ 是控制 KL 散度对结果影响的权重参数， $\gamma$ 是控制下一个时间（也就是样本）打分对结果影响的权重参数， $V_t$ 是 Critic Model 的打分输出， $A_t$ 则是 Reward Model 的打分输出;
 6. 根据奖励函数分别计算出的 actor loss 和 critic loss，更新 Actor Model 的参数和 Critic Model 的参数；注意，Actor Model 和 Critic Model 的参数更新方法是不同的，此处就不再一一赘述了，感兴趣的读者可以深入研究强化学习的相关理论。
 
 在上述过程中，因为要使用到四个模型，显存占用会数倍于 SFT。例如，如果我们 RM 和 LLM 都是用 7B 的体量，PPO 过程中大概需要 240G（4张 80G A100，每张卡占用 60G）显存来进行模型加载。那么，为什么我们需要足足四个模型呢？Actor Model 和 Critic Model 较为容易理解，而之所以我们还需要保持原参数不更新的 Ref Model 和 Reward Model，是为了限制模型的更新不要过于偏离原模型以至于丢失了 Pretrain 和 SFT 赋予的能力。

docs/chapter5/code/dataset.py

@@ -1,12 +1,10 @@
 import json
 import random
 import re
-
 import pandas as pd
 import numpy as np
 from torch.utils.data import Dataset, DataLoader
 import torch
-from sklearn.model_selection import train_test_split
 import os
     
 class PretrainDataset(Dataset):
@@ -15,15 +13,23 @@ class PretrainDataset(Dataset):
         self.data_path = data_path
         self.tokenizer = tokenizer
         self.max_length = max_length
-        self.padding = 0
-        with open(data_path, 'r', encoding='utf-8') as f:
-            self.data = f.readlines()
+        self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+        # 预计算每行的起始字节偏移量
+        self._offsets = []
+        with open(data_path, 'rb') as f:
+            self._offsets.append(0)
+            while f.readline():
+                self._offsets.append(f.tell())
+        self._total_lines = len(self._offsets) - 1  # 最后一个 tell() 是 EOF
 
     def __len__(self):
-        return len(self.data)
+        return self._total_lines
 
     def __getitem__(self, index: int):
-        sample = json.loads(self.data[index])
+        with open(self.data_path, 'rb') as f:
+            f.seek(self._offsets[index])
+            line = f.readline().decode('utf-8')
+        sample = json.loads(line)
         text = f"{self.tokenizer.bos_token}{sample['text']}"
         input_id = self.tokenizer(text).data['input_ids'][:self.max_length]
         text_len = len(input_id)
@@ -38,25 +44,28 @@ class PretrainDataset(Dataset):
         Y = np.array(input_id[1:]).astype(np.int64)
         loss_mask = np.array(loss_mask[1:]).astype(np.int64)
         return torch.from_numpy(X), torch.from_numpy(Y), torch.from_numpy(loss_mask)
-
-
+    
 class SFTDataset(Dataset):
     def __init__(self, data_path, tokenizer, max_length=512):
         super().__init__()
         self.data_path = data_path
         self.tokenizer = tokenizer
         self.max_length = max_length
-        self.padding = 0
-        with open(data_path, 'r', encoding='utf-8') as f:
-            self.data = f.readlines()
+        self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+        self._offsets = []
+        with open(data_path, 'rb') as f:
+            self._offsets.append(0)
+            while f.readline():
+                self._offsets.append(f.tell())
+        self._total_lines = len(self._offsets) - 1 
 
     def __len__(self):
-        return len(self.data)
+        return self._total_lines
 
     def generate_loss_mask(self, input_ids):
         # 生成 loss mask, 0 表示不计算损失, 1 表示计算损失
         mask = [0] * len(input_ids)
-        a_sequence = [3, 1074, 537, 500, 203]  # <|im_start|>assistant\n
+        a_sequence = self.tokenizer("<|im_start|>assistant\n")['input_ids']  # <|im_start|>assistant\n
         a_length = len(a_sequence)
         n = len(input_ids)
         i = 0
@@ -69,10 +78,10 @@ class SFTDataset(Dataset):
                     match = False
                     break
             if match:
-                # 从子序列结束的位置开始查找第一个4
+                # 从子序列结束的位置开始查找第一个 4 (eos_token_id)
                 j = None
                 for idx in range(i + a_length, n):
-                    if input_ids[idx] == 4:
+                    if input_ids[idx] == self.tokenizer.eos_token_id:
                         j = idx
                         break
                 if j is not None:
@@ -90,7 +99,10 @@ class SFTDataset(Dataset):
         return mask
 
     def __getitem__(self, index: int):
-        sample = json.loads(self.data[index])
+        with open(self.data_path, 'rb') as f:
+            f.seek(self._offsets[index])
+            line = f.readline().decode('utf-8')
+        sample = json.loads(line)
         text = self.tokenizer.apply_chat_template(sample, tokenize=False, add_generation_prompt=False)
         input_id = self.tokenizer(text).data['input_ids'][:self.max_length]
         text_len = len(input_id)
@@ -104,4 +116,4 @@ class SFTDataset(Dataset):
         X = np.array(input_id[:-1]).astype(np.int64)
         Y = np.array(input_id[1:]).astype(np.int64)
         loss_mask = np.array(loss_mask[1:]).astype(np.int64)
-        return torch.from_numpy(X), torch.from_numpy(Y), torch.from_numpy(loss_mask)
+        return torch.from_numpy(X), torch.from_numpy(Y), torch.from_numpy(loss_mask)

docs/chapter5/code/ddp_pretrain.py

@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 import os
 import platform
 import argparse
@@ -17,176 +18,308 @@ from dataset import PretrainDataset
 
 import swanlab
 
+# 忽略警告信息
 warnings.filterwarnings('ignore')
 
 
 def Logger(content):
+    """
+    简单的日志记录函数
+    
+    Args:
+        content (str): 要打印的内容
+    """
     print(content)
 
 def get_lr(it, all):
-    warmup_iters = args.warmup_iters
-    lr_decay_iters = all
-    min_lr = args.learning_rate / 10
+    """
+    计算当前迭代的学习率，使用余弦退火调度策略
+    
+    学习率调度策略：
+    1. Warmup阶段：学习率从0线性增长到目标学习率
+    2. 余弦退火阶段：学习率按余弦函数衰减到最小学习率
+    3. 超出训练步数后：保持最小学习率
+    
+    Args:
+        it (int): 当前迭代步数
+        all (int): 总迭代步数
+        
+    Returns:
+        float: 当前步数对应的学习率
+    """
+    warmup_iters = args.warmup_iters  # 预热迭代次数
+    lr_decay_iters = all  # 学习率衰减的总迭代次数
+    min_lr = args.learning_rate / 10  # 最小学习率，为初始学习率的1/10
 
+    # Warmup阶段：线性增长
     if it < warmup_iters:
         return args.learning_rate * it / warmup_iters
     
+    # 超出训练步数：保持最小学习率
     if it > lr_decay_iters:
         return min_lr
     
+    # 余弦退火阶段
     decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
     assert 0 <= decay_ratio <= 1
-    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # 余弦系数
     return min_lr + coeff * (args.learning_rate - min_lr)
 
 def train_epoch(epoch):
-    start_time = time.time()
+    """
+    训练一个epoch的函数
+    
+    实现了完整的训练循环，包括：
+    1. 数据加载和设备转移
+    2. 动态学习率调整
+    3. 前向传播和损失计算
+    4. 梯度累积和反向传播
+    5. 梯度裁剪和优化器更新
+    6. 日志记录和模型保存
+    
+    Args:
+        epoch (int): 当前epoch编号
+    """
+    start_time = time.time()  # 记录开始时间
+    
+    # 遍历数据加载器中的每个batch
     for step, (X, Y, loss_mask) in enumerate(train_loader):
-        X = X.to(args.device)
-        Y = Y.to(args.device)
-        loss_mask = loss_mask.to(args.device)
+        # 将数据转移到指定设备（GPU/CPU）
+        X = X.to(args.device)  # 输入序列
+        Y = Y.to(args.device)  # 目标序列
+        loss_mask = loss_mask.to(args.device)  # 损失掩码，用于忽略padding token
 
+        # 计算当前步骤的学习率
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch)
+        # 更新优化器中所有参数组的学习率
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
+        # 使用混合精度训练上下文
         with ctx:
+            # 前向传播
             out = model(X, Y)
+            # 计算损失并除以累积步数（用于梯度累积）
             loss = out.last_loss / args.accumulation_steps
+            # 将loss_mask展平为一维
             loss_mask = loss_mask.view(-1)
+            # 应用掩码计算有效损失（忽略padding位置）
             loss = torch.sum(loss * loss_mask) / loss_mask.sum()
 
+        # 使用scaler进行混合精度的反向传播
         scaler.scale(loss).backward()
 
+        # 每accumulation_steps步执行一次优化器更新
         if (step + 1) % args.accumulation_steps == 0:
+            # 取消梯度缩放，准备梯度裁剪
             scaler.unscale_(optimizer)
+            # 梯度裁剪，防止梯度爆炸
             torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
 
+            # 执行优化器步骤
             scaler.step(optimizer)
+            # 更新scaler的缩放因子
             scaler.update()
 
+            # 清零梯度，set_to_none=True可以节省内存
             optimizer.zero_grad(set_to_none=True)
 
+        # 每log_interval步记录一次日志
         if step % args.log_interval == 0:
             spend_time = time.time() - start_time
+            # 打印训练进度信息
             Logger(
-                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
+                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min;'.format(
                     epoch + 1,
                     args.epochs,
                     step,
                     iter_per_epoch,
-                    loss.item() * args.accumulation_steps,
+                    loss.item() * args.accumulation_steps,  # 恢复真实的loss值
                     optimizer.param_groups[-1]['lr'],
                     spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
+            
+            # 如果启用SwanLab，记录训练指标
             if args.use_swanlab:
                 swanlab.log({
                     "loss": loss.item() * args.accumulation_steps,
                     "lr": optimizer.param_groups[-1]['lr']
                 })
 
+        # 每save_interval步保存一次模型
         if (step + 1) % args.save_interval == 0:
-            model.eval()
+            model.eval()  # 切换到评估模式
+            # 构建检查点文件名
             ckp = f'{args.save_dir}/pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}.pth'
 
-            # 处理多卡保存
+            # 处理多卡保存：如果是DataParallel模型，需要访问.module属性
             state_dict = model.module.state_dict() if isinstance(model, torch.nn.DataParallel) else model.state_dict()
             torch.save(state_dict, ckp)
-            model.train()
+            model.train()  # 切换回训练模式
         
+        # 每20000步保存一个带步数标记的检查点
         if (step + 1) % 20000 == 0:
             model.eval()
+            # 构建带步数的检查点文件名
             ckp = f'{args.save_dir}/pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}_step{step+1}.pth'
 
+            # 保存模型状态字典
             state_dict = model.module.state_dict() if isinstance(model, torch.nn.DataParallel) else model.state_dict()
             torch.save(state_dict, ckp)
             model.train()
 
 
 def init_model():
+    """
+    初始化模型和分词器
+    
+    功能包括：
+    1. 加载预训练的分词器
+    2. 创建Transformer模型
+    3. 设置多GPU并行训练（如果可用）
+    4. 将模型移动到指定设备
+    5. 统计并打印模型参数量
+    
+    Returns:
+        tuple: (model, tokenizer) 初始化后的模型和分词器
+    """
     def count_parameters(model):
+        """
+        统计模型中可训练参数的数量
+        
+        Args:
+            model: PyTorch模型
+            
+        Returns:
+            int: 可训练参数总数
+        """
         return sum(p.numel() for p in model.parameters() if p.requires_grad)
 
+    # 从本地路径加载预训练的分词器
     tokenizer = AutoTokenizer.from_pretrained('./tokenizer_k/')
+    if tokenizer.pad_token_id is not None:
+        lm_config.pad_token_id = tokenizer.pad_token_id
 
+    # 根据配置创建Transformer模型
     model = Transformer(lm_config)
     
-    # 多卡初始化
+    # 多卡初始化：检查可用GPU数量并设置DataParallel
     num_gpus = torch.cuda.device_count()
     if num_gpus > 1:
         Logger(f"Using {num_gpus} GPUs with DataParallel!")
+        # 使用DataParallel包装模型以支持多GPU训练
         model = torch.nn.DataParallel(model)
     
+    # 将模型移动到指定设备（GPU或CPU）
     model = model.to(args.device)
+    
+    # 计算并打印模型参数量（以百万为单位）
     Logger(f'LLM总参数量：{count_parameters(model) / 1e6:.3f} 百万')
     return model, tokenizer
 
 
 if __name__ == "__main__":
+    # ==================== 命令行参数解析 ====================
     parser = argparse.ArgumentParser(description="Tiny-LLM Pretraining")
-    parser.add_argument("--out_dir", type=str, default="base_monkey_215M", help="Output directory")
-    parser.add_argument("--epochs", type=int, default=1, help="Number of epochs")
-    parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
-    parser.add_argument("--learning_rate", type=float, default=2e-4, help="Learning rate")
-    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="Device to use")
-    parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type")
-    parser.add_argument("--use_swanlab", type=bool, default=True, help="Use Weights & Biases")
-    parser.add_argument("--num_workers", type=int, default=8, help="Number of workers for data loading")
-    parser.add_argument("--data_path", type=str, default="", help="Path to training data")
-    parser.add_argument("--accumulation_steps", type=int, default=8, help="Gradient accumulation steps")
-    parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold")
-    parser.add_argument("--warmup_iters", type=int, default=0, help="Number of warmup iterations")
-    parser.add_argument("--log_interval", type=int, default=100, help="Logging interval")
-    parser.add_argument("--save_interval", type=int, default=1000, help="Model saving interval")
-    # 添加多卡参数
-    parser.add_argument("--gpus", type=str, default='0,1', help="Comma-separated GPU IDs (e.g. '0,1,2')")
+    
+    # 基础训练参数
+    parser.add_argument("--out_dir", type=str, default="base_model_215M", help="模型输出目录")
+    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=64, help="批次大小")
+    parser.add_argument("--learning_rate", type=float, default=2e-4, help="学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型")
+    
+    # 实验跟踪和数据加载参数
+    parser.add_argument("--use_swanlab", action="store_true", help="是否使用SwanLab进行实验跟踪")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载的工作进程数")
+    parser.add_argument("--data_path", type=str, default="./seq_monkey_datawhale.jsonl", help="训练数据路径")
+    
+    # 训练优化参数
+    parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--warmup_iters", type=int, default=0, help="学习率预热迭代次数")
+    
+    # 日志和保存参数
+    parser.add_argument("--log_interval", type=int, default=100, help="日志记录间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
+    
+    # 多GPU训练参数
+    parser.add_argument("--gpus", type=str, default='0,1,2,3,4,5,6,7', help="使用的GPU ID，用逗号分隔 (例如: '0,1,2')")
 
     args = parser.parse_args()
 
-    # 设置可见GPU
+    # ==================== GPU环境设置 ====================
+    # 设置可见的GPU设备
     if args.gpus is not None:
         os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
-        # 自动设置主设备为第一个GPU
+        # 自动设置主设备为第一个可用GPU
         if torch.cuda.is_available():
             args.device = "cuda:0"
         else:
             args.device = "cpu"
 
+    # ==================== 实验跟踪初始化 ====================
     if args.use_swanlab:
-        swanlab.login(api_key='your key')
+        # 注意：使用前需要先登录 swanlab.login(api_key='your key')
         run = swanlab.init(
-            project="Tiny-LLM",
-            experiment_name="Pretrain-215M",
-            config=args,
+            project="Happy-LLM",  # 项目名称
+            experiment_name="Pretrain-215M",  # 实验名称
+            config=args,  # 保存所有超参数
         )
 
+    # ==================== 模型配置 ====================
+    # 定义语言模型的配置参数
     lm_config = ModelConfig(
-        dim=1024,
-        n_layers=18,
+        dim=1024,      # 模型维度
+        n_layers=18,   # Transformer层数
     )
-    max_seq_len = lm_config.max_seq_len
-    args.save_dir = os.path.join(args.out_dir)
-    os.makedirs(args.save_dir, exist_ok=True)
+
+    # ==================== 训练环境设置 ====================
+    max_seq_len = lm_config.max_seq_len  # 最大序列长度
+    args.save_dir = os.path.join(args.out_dir)  # 模型保存目录
+    
+    # 创建必要的目录
     os.makedirs(args.out_dir, exist_ok=True)
+    
+    # 设置随机种子以确保结果可复现
     torch.manual_seed(42)
+    
+    # 确定设备类型（用于选择合适的上下文管理器）
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
+    # 设置混合精度训练的上下文管理器
+    # CPU训练时使用nullcontext，GPU训练时使用autocast
     ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
 
+    # ==================== 模型和数据初始化 ====================
+    # 初始化模型和分词器
     model, tokenizer = init_model()
     
+    # 创建训练数据集
     train_ds = PretrainDataset(args.data_path, tokenizer, max_length=max_seq_len)
+    
+    # 创建数据加载器
     train_loader = DataLoader(
         train_ds,
-        batch_size=args.batch_size,
-        pin_memory=True,
-        drop_last=False,
-        shuffle=True,
-        num_workers=args.num_workers
+        batch_size=args.batch_size,  # 批次大小
+        pin_memory=True,             # 将数据加载到固定内存中，加速GPU传输
+        drop_last=False,             # 不丢弃最后一个不完整的批次
+        shuffle=True,                # 随机打乱数据
+        num_workers=args.num_workers # 数据加载的并行工作进程数
     )
 
+    # ==================== 优化器和训练组件初始化 ====================
+    # 初始化混合精度训练的梯度缩放器
+    # 只有在使用float16或bfloat16时才启用
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    
+    # 初始化Adam优化器
     optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
 
+    # ==================== 开始训练 ====================
+    # 计算每个epoch的迭代次数
     iter_per_epoch = len(train_loader)
+    
+    # 开始训练循环
     for epoch in range(args.epochs):
-        train_epoch(epoch)
+        train_epoch(epoch)

docs/chapter5/code/ddp_sft_full.py

@@ -17,13 +17,18 @@ from dataset import SFTDataset
 
 import swanlab
 
+# 忽略警告
 warnings.filterwarnings('ignore')
 
 
 def Logger(content):
+    """日志记录器"""
     print(content)
 
 def get_lr(it, all):
+    """获取学习率"""
+    # 1) linear warmup for warmup_iters steps
+    # 1) 预热迭代的线性预热
     warmup_iters = args.warmup_iters
     lr_decay_iters = all
     min_lr = args.learning_rate / 10
@@ -31,33 +36,42 @@ def get_lr(it, all):
     if it < warmup_iters:
         return args.learning_rate * it / warmup_iters
     
+    # 2) if it > lr_decay_iters, return min learning rate
+    # 2) 如果迭代次数超过学习率衰减迭代次数，则返回最小学习率
     if it > lr_decay_iters:
         return min_lr
     
+    # 3) in between, use cosine decay down to min learning rate
+    # 3) 在两者之间，使用余弦衰减至最小学习率
     decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
     assert 0 <= decay_ratio <= 1
     coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
     return min_lr + coeff * (args.learning_rate - min_lr)
 
 def train_epoch(epoch):
+    """训练一个epoch"""
     start_time = time.time()
     for step, (X, Y, loss_mask) in enumerate(train_loader):
         X = X.to(args.device)
         Y = Y.to(args.device)
         loss_mask = loss_mask.to(args.device)
 
+        # 获取学习率并更新优化器
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
+        # 前向传播
         with ctx:
             out = model(X, Y)
             loss = out.last_loss / args.accumulation_steps
             loss_mask = loss_mask.view(-1)
             loss = torch.sum(loss * loss_mask) / loss_mask.sum()
 
+        # 反向传播
         scaler.scale(loss).backward()
 
+        # 更新权重
         if (step + 1) % args.accumulation_steps == 0:
             scaler.unscale_(optimizer)
             torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
@@ -67,6 +81,7 @@ def train_epoch(epoch):
 
             optimizer.zero_grad(set_to_none=True)
 
+        # 打印日志
         if step % args.log_interval == 0:
             spend_time = time.time() - start_time
             Logger(
@@ -84,6 +99,7 @@ def train_epoch(epoch):
                     "lr": optimizer.param_groups[-1]['lr']
                 })
 
+        # 保存模型
         if (step + 1) % args.save_interval == 0:
             model.eval()
             ckp = f'{args.save_dir}/sft_dim{lm_config.dim}_layers{lm_config.n_layers}_vocab_size{lm_config.vocab_size}.pth'
@@ -93,6 +109,7 @@ def train_epoch(epoch):
             torch.save(state_dict, ckp)
             model.train()
         
+        # 定期保存模型
         if (step + 1) % 20000 == 0:
             model.eval()
             ckp = f'{args.save_dir}/sft_dim{lm_config.dim}_layers{lm_config.n_layers}_vocab_size{lm_config.vocab_size}_step{step+1}.pth'
@@ -103,14 +120,21 @@ def train_epoch(epoch):
 
 
 def init_model():
+    """初始化模型"""
     def count_parameters(model):
+        """计算模型参数量"""
         return sum(p.numel() for p in model.parameters() if p.requires_grad)
 
+    # 加载分词器
     tokenizer = AutoTokenizer.from_pretrained('./tokenizer_k/')
+    if tokenizer.pad_token_id is not None:
+        lm_config.pad_token_id = tokenizer.pad_token_id
 
+    # 初始化模型
     model = Transformer(lm_config)
 
-    ckp = './base_monkey_215M/pretrain_1024_18_6144.pth'
+    # 加载预训练权重
+    ckp = './base_model_215M/pretrain_1024_18_6144.pth'
     state_dict = torch.load(ckp, map_location=args.device)
     unwanted_prefix = '_orig_mod.'
     for k, v in list(state_dict.items()):
@@ -131,22 +155,22 @@ def init_model():
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Tiny-LLM Pretraining")
-    parser.add_argument("--out_dir", type=str, default="BeelGroup_sft_model_215M", help="Output directory")
-    parser.add_argument("--epochs", type=int, default=1, help="Number of epochs")
-    parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
-    parser.add_argument("--learning_rate", type=float, default=2e-4, help="Learning rate")
-    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="Device to use")
-    parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type")
-    parser.add_argument("--use_swanlab", type=bool, default=True, help="Use Weights & Biases")
-    parser.add_argument("--num_workers", type=int, default=4, help="Number of workers for data loading")
-    parser.add_argument("--data_path", type=str, default="", help="Path to training data")
-    parser.add_argument("--accumulation_steps", type=int, default=4, help="Gradient accumulation steps")
-    parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold")
-    parser.add_argument("--warmup_iters", type=int, default=0, help="Number of warmup iterations")
-    parser.add_argument("--log_interval", type=int, default=100, help="Logging interval")
-    parser.add_argument("--save_interval", type=int, default=1000, help="Model saving interval")
+    parser.add_argument("--out_dir", type=str, default="sft_model_215M", help="输出目录")
+    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=64, help="批处理大小")
+    parser.add_argument("--learning_rate", type=float, default=2e-4, help="学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="使用的设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型")
+    parser.add_argument("--use_swanlab", action="store_true", help="是否使用SwanLab进行实验跟踪")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载的工作进程数")
+    parser.add_argument("--data_path", type=str, default="./BelleGroup_sft.jsonl", help="训练数据路径")
+    parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--warmup_iters", type=int, default=0, help="预热迭代次数")
+    parser.add_argument("--log_interval", type=int, default=100, help="日志记录间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
     # 添加多卡参数
-    parser.add_argument("--gpus", type=str, default='0,1', help="Comma-separated GPU IDs (e.g. '0,1,2')")
+    parser.add_argument("--gpus", type=str, default='0,1,2,3,4,5,6,7', help="逗号分隔的GPU ID (例如 '0,1,2')")
 
     args = parser.parse_args()
 
@@ -159,29 +183,32 @@ if __name__ == "__main__":
         else:
             args.device = "cpu"
 
+    # 初始化swanlab
     if args.use_swanlab:
-        swanlab.login(api_key='your key')
         run = swanlab.init(
-            project="Tiny-LLM",
-            experiment_name="BelleGropu-sft-215M",
+            project="Happy-LLM",
+            experiment_name="SFT-215M",
             config=args,
         )
 
+    # 模型配置
     lm_config = ModelConfig(
         dim=1024,
         n_layers=18,
     )
     max_seq_len = lm_config.max_seq_len
     args.save_dir = os.path.join(args.out_dir)
-    os.makedirs(args.save_dir, exist_ok=True)
     os.makedirs(args.out_dir, exist_ok=True)
     torch.manual_seed(42)
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
+    # 上下文管理器
     ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
 
+    # 初始化模型和分词器
     model, tokenizer = init_model()
     
+    # 创建数据集和数据加载器
     train_ds = SFTDataset(args.data_path, tokenizer, max_length=max_seq_len)
     train_loader = DataLoader(
         train_ds,
@@ -192,9 +219,11 @@ if __name__ == "__main__":
         num_workers=args.num_workers
     )
 
+    # 缩放器和优化器
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
-    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
 
+    # 开始训练
     iter_per_epoch = len(train_loader)
     for epoch in range(args.epochs):
-        train_epoch(epoch)
+        train_epoch(epoch)

docs/chapter5/code/deal_dataset.py

@@ -1,32 +1,24 @@
-import os 
-from tqdm import tqdm
+import os
 import json
+from tqdm import tqdm
 
-# 设置环境变量
-os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
-
-
-# 下载预训练数据集
-os.system("modelscope download --dataset ddzhu123/seq-monkey mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2 --local_dir your_local_dir")
-# 解压预训练数据集
-os.system("tar -xvf your_local_dir/mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2")
-
-# 下载SFT数据集
-os.system(f'huggingface-cli download --repo-type dataset --resume-download BelleGroup/train_3.5M_CN --local-dir BelleGroup')
-
+# pretrain_data 为运行download_dataset.sh时，下载的pretrain_data本地路径
+pretrain_data = 'your local pretrain_data'
+output_pretrain_data = 'seq_monkey_datawhale.jsonl'
 
+# sft_data 为运行download_dataset.sh时，下载的sft_data本地路径
+sft_data = 'your local sft_data'
+output_sft_data = 'BelleGroup_sft.jsonl'
 
 # 1 处理预训练数据
 def split_text(text, chunk_size=512):
     """将文本按指定长度切分成块"""
     return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
 
-input_file = 'mobvoi_seq_monkey_general_open_corpus.jsonl'
-
-with open('seq_monkey_datawhale.jsonl', 'a', encoding='utf-8') as pretrain:
-    with open(input_file, 'r', encoding='utf-8') as f:
+with open(output_pretrain_data, 'a', encoding='utf-8') as pretrain:
+    with open(pretrain_data, 'r', encoding='utf-8') as f:
         data = f.readlines()
-        for line in tqdm(data, desc=f"Processing lines in {input_file}", leave=False):  # 添加行级别的进度条
+        for line in tqdm(data, desc=f"Processing lines in {pretrain_data}", leave=False):  # 添加行级别的进度条
             line = json.loads(line)
             text = line['text']
             chunks = split_text(text)
@@ -34,7 +26,6 @@ with open('seq_monkey_datawhale.jsonl', 'a', encoding='utf-8') as pretrain:
                 pretrain.write(json.dumps({'text': chunk}, ensure_ascii=False) + '\n')
 
 # 2 处理SFT数据
-
 def convert_message(data):
     """
     将原始数据转换为标准格式
@@ -49,10 +40,10 @@ def convert_message(data):
             message.append({'role': 'assistant', 'content': item['value']})
     return message
 
-with open('BelleGroup_sft.jsonl', 'a', encoding='utf-8') as sft:
-    with open('BelleGroup/train_3.5M_CN.json', 'r') as f:
+with open(output_sft_data, 'a', encoding='utf-8') as sft:
+    with open(sft_data, 'r', encoding='utf-8') as f:
         data = f.readlines()
         for item in tqdm(data, desc="Processing", unit="lines"):
             item = json.loads(item)
             message = convert_message(item['conversations'])
-            sft.write(json.dumps(message, ensure_ascii=False) + '\n')
+            sft.write(json.dumps(message, ensure_ascii=False) + '\n')

docs/chapter5/code/download_dataset.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# 设置环境变量
+export HF_ENDPOINT=https://hf-mirror.com
+
+# dataset dir 下载到本地目录
+dataset_dir="your local dataset dir"
+
+# 下载预训练数据集， 需要预先安装modelscope，使用pip3 install modelscope安装
+modelscope download --dataset ddzhu123/seq-monkey mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2 --local_dir ${dataset_dir}
+
+# 解压预训练数据集
+tar -xvf "${dataset_dir}/mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2" -C "${dataset_dir}"
+
+# 下载SFT数据集
+huggingface-cli download \
+  --repo-type dataset \
+  --resume-download \
+  BelleGroup/train_3.5M_CN \
+  --local-dir "${dataset_dir}/BelleGroup"

docs/chapter5/code/export_model.py

@@ -15,6 +15,15 @@ def export_model(tokenizer_path, model_config, model_ckpt_path, save_directory):
     ModelConfig.register_for_auto_class()
     Transformer.register_for_auto_class("AutoModelForCausalLM")
 
+    # 加载tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer_path,
+        trust_remote_code=True,
+        use_fast=False
+    )
+    if tokenizer.pad_token_id is not None:
+        model_config.pad_token_id = tokenizer.pad_token_id
+
     # 初始化模型
     model = Transformer(model_config)
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -31,13 +40,6 @@ def export_model(tokenizer_path, model_config, model_ckpt_path, save_directory):
     model.load_state_dict(state_dict, strict=False)
     print(f'模型参数: {count_parameters(model)/1e6:.2f}M = {count_parameters(model)/1e9:.2f}B')
 
-    # 加载tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(
-        tokenizer_path,
-        trust_remote_code=True,
-        use_fast=False
-    )
-
     # 保存完整模型和tokenizer
     model.save_pretrained(save_directory, safe_serialization=False)
     tokenizer.save_pretrained(save_directory)
@@ -56,4 +58,4 @@ if __name__ == '__main__':
         model_config=config,
         model_ckpt_path='./BeelGroup_sft_model_215M/sft_dim1024_layers18_vocab_size6144.pth',
         save_directory="k-model-215M"
-    )
+    )

docs/chapter5/code/k_model.py

@@ -26,6 +26,7 @@ class ModelConfig(PretrainedConfig):
             max_seq_len: int = 512,
             dropout: float = 0.0,
             flash_attn: bool = True,
+            pad_token_id: int = 0,
             **kwargs,
     ):
         self.dim = dim
@@ -39,6 +40,7 @@ class ModelConfig(PretrainedConfig):
         self.max_seq_len = max_seq_len
         self.dropout = dropout
         self.flash_attn = flash_attn
+        self.pad_token_id = pad_token_id
         super().__init__(**kwargs)
 
 class RMSNorm(nn.Module):
@@ -177,7 +179,7 @@ class Attention(nn.Module):
             # 注册为模型的缓冲区
             self.register_buffer("mask", mask)
 
-    def forward(self, x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor):
+    def forward(self, x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
         # 获取批次大小和序列长度，[batch_size, seq_len, dim]
         bsz, seqlen, _ = x.shape
 
@@ -199,16 +201,40 @@ class Attention(nn.Module):
         xq = xq.transpose(1, 2)
         xk = xk.transpose(1, 2)
         xv = xv.transpose(1, 2)
+        key_padding_mask = None
+        if attention_mask is not None:
+            key_padding_mask = attention_mask[:, None, None, :].to(dtype=torch.bool)
 
         # 根据是否支持Flash Attention，选择实现方式。
         if self.flash:
             # 使用Flash Attention。
-            output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
+            if key_padding_mask is not None:
+                causal_mask = torch.ones((seqlen, seqlen), dtype=torch.bool, device=x.device).tril()
+                full_attn_mask = causal_mask[None, None, :, :] & key_padding_mask
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    xq,
+                    xk,
+                    xv,
+                    attn_mask=full_attn_mask,
+                    dropout_p=self.dropout if self.training else 0.0,
+                    is_causal=False,
+                )
+            else:
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    xq,
+                    xk,
+                    xv,
+                    attn_mask=None,
+                    dropout_p=self.dropout if self.training else 0.0,
+                    is_causal=True,
+                )
         else:
             # 使用手动实现的注意力机制。
             scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
             assert hasattr(self, 'mask')
             scores = scores + self.mask[:, :, :seqlen, :seqlen]
+            if key_padding_mask is not None:
+                scores = scores.masked_fill(~key_padding_mask, float("-inf"))
             scores = F.softmax(scores.float(), dim=-1).type_as(xq)
             scores = self.attn_dropout(scores)
             output = torch.matmul(scores, xv)
@@ -272,11 +298,11 @@ class DecoderLayer(nn.Module):
         # 定义前馈神经网络计算的归一化层
         self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
 
-    def forward(self, x, freqs_cos, freqs_sin):
+    def forward(self, x, freqs_cos, freqs_sin, attention_mask: Optional[torch.Tensor] = None):
         # 前向传播函数
         # 首先，输入x经过注意力归一化层，然后进行注意力计算，结果与输入x相加得到h
         # 然后，h经过前馈神经网络归一化层，然后进行前馈神经网络计算，结果与h相加得到输出
-        h = x + self.attention.forward(self.attention_norm(x), freqs_cos, freqs_sin)
+        h = x + self.attention.forward(self.attention_norm(x), freqs_cos, freqs_sin, attention_mask=attention_mask)
         out = h + self.feed_forward.forward(self.ffn_norm(h))
         return out
 
@@ -334,21 +360,61 @@ class Transformer(PreTrainedModel):
                 torch.nn.init.zeros_(module.bias)
         elif isinstance(module, nn.Embedding):
             torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def _prepare_attention_mask(self, attention_mask: Optional[torch.Tensor], tokens: torch.Tensor) -> Optional[torch.Tensor]:
+        if attention_mask is None:
+            return None
+        if attention_mask.dim() == 4:
+            attention_mask = attention_mask[:, 0, 0, :]
+        elif attention_mask.dim() == 3:
+            attention_mask = attention_mask[:, 0, :]
+        attention_mask = attention_mask.to(tokens.device)
+        if attention_mask.dtype != torch.bool:
+            attention_mask = attention_mask > 0
+        if attention_mask.shape != tokens.shape:
+            raise ValueError(f"attention_mask shape {attention_mask.shape} must match input_ids shape {tokens.shape}")
+        return attention_mask
+
+    def _left_pad_by_attention_mask(
+            self,
+            idx: torch.Tensor,
+            attention_mask: Optional[torch.Tensor],
+            pad_token_id: int
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if attention_mask is None or attention_mask.all():
+            return idx, attention_mask
+
+        bsz = idx.size(0)
+        lengths = attention_mask.long().sum(dim=1)
+        max_len = max(int(lengths.max().item()), 1)
+        packed_idx = idx.new_full((bsz, max_len), pad_token_id)
+        packed_mask = attention_mask.new_zeros((bsz, max_len), dtype=torch.bool)
+
+        for row in range(bsz):
+            valid_len = int(lengths[row].item())
+            if valid_len <= 0:
+                continue
+            valid_tokens = idx[row][attention_mask[row]]
+            packed_idx[row, max_len - valid_len:] = valid_tokens
+            packed_mask[row, max_len - valid_len:] = True
+
+        return packed_idx, packed_mask
     
-    def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, **keyargs) -> torch.Tensor:
+    def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
         """
         - tokens: Optional[torch.Tensor], 输入 token 张量。
         - targets: Optional[torch.Tensor], 目标 token 张量。
         - kv_cache: bool, 是否使用键值缓存。
-        - keyargs: 其他关键字参数。
+        - kwargs: 其他关键字参数。
 
         - self.OUT: CausalLMOutputWithPast, 包含 logits 和损失。
         """
 
-        if 'input_ids' in keyargs:
-            tokens = keyargs['input_ids']
-        if 'attention_mask' in keyargs:
-            targets = keyargs['attention_mask']
+        if 'input_ids' in kwargs:
+            tokens = kwargs['input_ids']
+        if 'labels' in kwargs:
+            targets = kwargs['labels']
+        attention_mask = self._prepare_attention_mask(kwargs.get('attention_mask'), tokens)
 
         # 前向传播函数
         _bsz, seqlen = tokens.shape
@@ -361,17 +427,30 @@ class Transformer(PreTrainedModel):
 
         # 通过Decoder层
         for layer in self.layers:
-            h = layer(h, freqs_cos, freqs_sin)
+            h = layer(h, freqs_cos, freqs_sin, attention_mask=attention_mask)
         # 通过归一化层
         h = self.norm(h)
 
         if targets is not None:
             # 如果给定了目标，计算损失
             logits = self.output(h)
-            self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=0, reduction='none')
+            ignore_index = self.args.pad_token_id if self.args.pad_token_id is not None else 0
+            if torch.any(targets == -100):
+                ignore_index = -100
+            self.last_loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=ignore_index,
+                reduction='none'
+            )
         else:
             # 推理时的小优化：只对最后一个位置的输出进行前向传播
-            logits = self.output(h[:, [-1], :]) 
+            if attention_mask is None:
+                logits = self.output(h[:, [-1], :])
+            else:
+                full_logits = self.output(h)
+                last_token_pos = attention_mask.long().sum(dim=1).clamp(min=1) - 1
+                logits = full_logits[torch.arange(_bsz, device=tokens.device), last_token_pos].unsqueeze(1)
             self.last_loss = None
 
         # 设置输出
@@ -381,18 +460,36 @@ class Transformer(PreTrainedModel):
 
     
     @torch.inference_mode()
-    def generate(self, idx, stop_id=None, max_new_tokens=256, temperature=1.0, top_k=None):
+    def generate(
+            self,
+            idx,
+            stop_id=None,
+            max_new_tokens=256,
+            temperature=1.0,
+            top_k=None,
+            attention_mask: Optional[torch.Tensor] = None,
+            pad_token_id: Optional[int] = None
+    ):
         """
         给定输入序列 idx（形状为 (bz,seq_len) 的长整型张量），通过多次生成新 token 来完成序列。
         在 model.eval() 模式下运行。效率较低的采样版本，没有使用键k/v cache。
         """
+        if pad_token_id is None:
+            pad_token_id = self.args.pad_token_id if self.args.pad_token_id is not None else 0
+        attention_mask = self._prepare_attention_mask(attention_mask, idx)
+        idx, attention_mask = self._left_pad_by_attention_mask(idx, attention_mask, pad_token_id)
+
+        finished = torch.zeros(idx.size(0), dtype=torch.bool, device=idx.device)
         index = idx.shape[1]
         for _ in range(max_new_tokens):
             # 如果序列上下文过长，截断它到最大长度
             idx_cond = idx if idx.size(1) <= self.args.max_seq_len else idx[:, -self.args.max_seq_len:]
+            mask_cond = None
+            if attention_mask is not None:
+                mask_cond = attention_mask if attention_mask.size(1) <= self.args.max_seq_len else attention_mask[:, -self.args.max_seq_len:]
             
             # 前向传播获取序列中最后一个位置的 logits
-            logits = self(idx_cond).logits
+            logits = self(idx_cond, attention_mask=mask_cond).logits
             logits = logits[:, -1, :] # 只保留最后一个时间步的输出
             
             if temperature == 0.0:
@@ -406,13 +503,273 @@ class Transformer(PreTrainedModel):
                     logits[logits < v[:, [-1]]] = -float('Inf')
                 probs = F.softmax(logits, dim=-1)
                 idx_next = torch.multinomial(probs, num_samples=1)
-            
 
-            if idx_next == stop_id:
-                break
+            prev_finished = finished.clone()
+            if stop_id is not None:
+                if prev_finished.any():
+                    fill_token = pad_token_id if pad_token_id is not None else stop_id
+                    idx_next = torch.where(prev_finished[:, None], torch.full_like(idx_next, fill_token), idx_next)
+                finished = prev_finished | idx_next[:, 0].eq(stop_id)
 
             # 将采样的索引添加到序列中并继续
             idx = torch.cat((idx, idx_next), dim=1)
+            if attention_mask is not None:
+                next_mask = torch.ones((attention_mask.size(0), 1), dtype=attention_mask.dtype, device=attention_mask.device)
+                if prev_finished.any():
+                    next_mask[prev_finished] = False
+                attention_mask = torch.cat((attention_mask, next_mask), dim=1)
+
+            if stop_id is not None and finished.all():
+                break
+
+        return idx[:, index:] # 只返回生成的token
+    
+    def _greedy_decode(self, logits: torch.Tensor) -> torch.Tensor:
+        """
+        贪婪解码：选择概率最大的token
+
+        Args:
+            logits: 模型输出的logits，形状为 (batch_size, vocab_size)
+
+        Returns:
+            选择的token索引，形状为 (batch_size, 1)
+        """
+        _, idx_next = torch.topk(logits, k=1, dim=-1)
+        return idx_next
+
+    def _random_sample(self, logits: torch.Tensor, temperature: float = 1.0, top_k: int = None) -> torch.Tensor:
+        """
+        随机采样：基于概率分布随机选择token
+
+        Args:
+            logits: 模型输出的logits，形状为 (batch_size, vocab_size)
+            temperature: 温度参数，控制随机性
+            top_k: 只考虑概率最高的k个token
+
+        Returns:
+            选择的token索引，形状为 (batch_size, 1)
+        """
+        # 缩放 logits
+        logits = logits / temperature
+
+        # 应用top-k过滤
+        if top_k is not None:
+            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+            # 将不在 top-k 内的 logits 设为负无穷
+            logits[logits < v[:, [-1]]] = -float('Inf')
+
+        # 计算概率并采样
+        probs = F.softmax(logits, dim=-1)
+        idx_next = torch.multinomial(probs, num_samples=1)
+        return idx_next
+
+    def _beam_search(self, idx: torch.Tensor, max_new_tokens: int, num_beams: int,
+                     temperature: float = 1.0, top_k: int = None, stop_id: int = None) -> torch.Tensor:
+        """
+        束搜索：维护多个候选序列，选择最优路径
+
+        束搜索的核心思想：在每一步生成时，不是只选择一个最佳token，
+        而是保留多个候选路径，最终选择累积概率最高的完整序列。
+
+        Args:
+            idx: 输入序列，形状为 (batch_size, seq_len)
+            max_new_tokens: 最大生成token数量
+            num_beams: 束宽度，表示保留的候选路径数量
+            temperature: 温度参数，控制分布的平滑程度
+            top_k: top-k过滤参数，限制候选token范围
+            stop_id: 停止生成的token ID，遇到则停止
+
+        Returns:
+            生成的token序列，形状为 (batch_size, generated_length)
+            只返回新生成的部分，不包含原始输入序列
+        """
+        # 获取输入序列的基本信息
+        batch_size = idx.shape[0]  # 批次大小，通常为1
+        seq_len = idx.shape[1]     # 输入序列长度
+
+        # 初始化束：创建 num_beams 个候选序列
+        beams = [idx.clone() for _ in range(num_beams)]
+        # 初始化每个候选序列的累积对数概率分数
+        beam_scores = torch.zeros(num_beams, device=idx.device)
+        # 第一个候选是原始输入序列，分数为0
+        beam_scores[0] = 0.0
+        # 其他候选初始分数设为负无穷，表示尚未生成
+        beam_scores[1:] = float('-inf')
+
+        # 主循环：逐步生成新的token，最多生成 max_new_tokens 个
+        for step in range(max_new_tokens):
+            # 每轮迭代收集新的候选序列和分数
+            new_beams = []   # 新的候选序列列表
+            new_scores = []  # 对应的分数列表
+
+            # 遍历当前的所有候选序列
+            for beam_idx, beam in enumerate(beams):
+                # 跳过无效候选（分数为负无穷的序列）
+                if beam_scores[beam_idx] == float('-inf'):
+                    continue
+
+                # 序列长度检查：如果超过最大长度，截取最后的部分
+                beam_cond = beam if beam.size(1) <= self.args.max_seq_len else beam[:, -self.args.max_seq_len:]
+
+                # 前向传播：获取模型对当前序列的预测
+                output = self(beam_cond)
+                # 提取最后一个位置的logits，用于预测下一个token
+                logits = output.logits[:, -1, :]  # 形状: (1, vocab_size)
+
+                # 温度缩放：调整logits的分布
+                if temperature != 1.0:
+                    logits = logits / temperature
+                    # 温度 > 1：分布更平滑，增加随机性
+                    # 温度 < 1：分布更尖锐，更确定
+
+                # Top-k过滤：限制候选token的范围，提高质量
+                if top_k is not None:
+                    # 找到logits中前top_k个最大的值
+                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                    # 将不在前top_k内的logits设为负无穷
+                    logits[logits < v[:, [-1]]] = -float('Inf')
+                    # 这样采样时只会考虑前top_k个token
+
+                # 计算对数概率：使用log_softmax避免数值不稳定
+                log_probs = F.log_softmax(logits, dim=-1)
+
+                # 获取前 num_beams 个最可能的候选token
+                # 注意：这里的top-k与上面的top-k不同
+                # 上面的top-k是全局过滤，这里是束搜索的分支选择
+                top_log_probs, top_indices = torch.topk(log_probs, k=num_beams, dim=-1)
+
+                # 为当前候选序列生成 num_beams 个扩展序列
+                for k in range(num_beams):
+                    # 选择第k个候选token
+                    token = top_indices[:, k:k+1]      # token ID
+                    log_prob = top_log_probs[:, k]     # 对应的对数概率
+
+                    # 扩展序列：将新token添加到当前序列末尾
+                    new_beam = torch.cat([beam, token], dim=1)
+                    # 更新累积分数：原序列分数 + 新token的对数概率
+                    new_score = beam_scores[beam_idx] + log_prob.item()
+
+                    # 保存新的候选序列和分数
+                    new_beams.append(new_beam)
+                    new_scores.append(new_score)
+
+            # 安全检查：如果没有生成任何有效候选，提前结束
+            if not new_beams:
+                break
+
+            # 筛选最佳候选：从所有新生成的候选中选择分数最高的 num_beams 个
+            # 按分数降序排序，获取索引
+            sorted_indices = sorted(range(len(new_scores)), key=lambda i: new_scores[i], reverse=True)
+            # 选择前 num_beams 个最佳候选
+            beams = [new_beams[i] for i in sorted_indices[:num_beams]]
+            beam_scores = [new_scores[i] for i in sorted_indices[:num_beams]]
+
+            # 停止条件检查：检查最佳序列是否以停止token结尾
+            if stop_id is not None and beams[0][0, -1] == stop_id:
+                break
+
+        # 返回得分最高的序列，只返回新生成的部分（去掉原始输入）
+        # beams[0] 是最终得分最高的完整序列
+        # [:, seq_len:] 切片只保留生成部分
+        return beams[0][:, seq_len:]
+
+    @torch.inference_mode()
+    def generate_super(self,
+                       idx,
+                       stop_id=None,
+                       max_new_tokens=256,
+                       temperature=1.0,
+                       top_k=None,
+                       do_sample=False,
+                       num_beams=1,
+                       attention_mask: Optional[torch.Tensor] = None,
+                       pad_token_id: Optional[int] = None
+                       ):
+        """
+        高级文本生成函数，支持三种解码策略：
+
+        1. 贪婪解码（Greedy Search）：
+           - 参数：do_sample=False, num_beams=1
+           - 特点：每步选择概率最大的token，速度快、结果确定
+
+        2. 随机采样（Random Sampling）：
+           - 参数：do_sample=True, num_beams=1
+           - 特点：基于概率分布随机采样，可配合temperature和top-k控制多样性
+
+        3. 束搜索（Beam Search）：
+           - 参数：do_sample=False, num_beams>1
+           - 特点：维护多条候选路径，选择总概率最高的序列，质量更高但速度较慢
+
+        Args:
+            idx: 输入序列张量，形状为 (batch_size, seq_len)
+            stop_id: 停止生成的token ID
+            max_new_tokens: 最大生成token数量
+            temperature: 温度参数，控制随机性，越高越随机
+            top_k: 只考虑概率最高的k个token，None表示不考虑
+            do_sample: 是否使用随机采样，False时使用确定性解码
+            num_beams: 束搜索的束宽度，1表示不使用束搜索
+
+        Returns:
+            生成的token序列，形状为 (batch_size, generated_length)
+        """
+        # 参数验证
+        if temperature <= 0:
+            temperature = 0.001  # 避免除零错误
+        if num_beams < 1:
+            num_beams = 1
+        if top_k is not None and top_k < 1:
+            top_k = None
+        if pad_token_id is None:
+            pad_token_id = self.args.pad_token_id if self.args.pad_token_id is not None else 0
+        attention_mask = self._prepare_attention_mask(attention_mask, idx)
+        idx, attention_mask = self._left_pad_by_attention_mask(idx, attention_mask, pad_token_id)
+
+        # 束搜索逻辑
+        if not do_sample and num_beams > 1:
+            return self._beam_search(idx, max_new_tokens, num_beams, temperature, top_k, stop_id)
+
+        # 贪婪解码和随机采样逻辑
+        finished = torch.zeros(idx.size(0), dtype=torch.bool, device=idx.device)
+        index = idx.shape[1]
+        for _ in range(max_new_tokens):
+            # 如果序列上下文过长，截断它到最大长度
+            idx_cond = idx if idx.size(1) <= self.args.max_seq_len else idx[:, -self.args.max_seq_len:]
+            mask_cond = None
+            if attention_mask is not None:
+                mask_cond = attention_mask if attention_mask.size(1) <= self.args.max_seq_len else attention_mask[:, -self.args.max_seq_len:]
+
+            # 前向传播获取序列中最后一个位置的 logits
+            logits = self(idx_cond, attention_mask=mask_cond).logits
+            logits = logits[:, -1, :] # 只保留最后一个时间步的输出
+
+            # 根据参数选择解码策略
+            if do_sample:
+                idx_next = self._random_sample(logits, temperature, top_k)
+            else:
+                # 当temperature=0时使用贪婪解码
+                if temperature < 0.1:
+                    idx_next = self._greedy_decode(logits)
+                else:
+                    # 低温度下的随机采样（接近贪婪）
+                    idx_next = self._random_sample(logits, temperature, top_k)
+
+            prev_finished = finished.clone()
+            if stop_id is not None:
+                if prev_finished.any():
+                    fill_token = pad_token_id if pad_token_id is not None else stop_id
+                    idx_next = torch.where(prev_finished[:, None], torch.full_like(idx_next, fill_token), idx_next)
+                finished = prev_finished | idx_next[:, 0].eq(stop_id)
+
+            # 将选择的token添加到序列中
+            idx = torch.cat((idx, idx_next), dim=1)
+            if attention_mask is not None:
+                next_mask = torch.ones((attention_mask.size(0), 1), dtype=attention_mask.dtype, device=attention_mask.device)
+                if prev_finished.any():
+                    next_mask[prev_finished] = False
+                attention_mask = torch.cat((attention_mask, next_mask), dim=1)
+
+            if stop_id is not None and finished.all():
+                break
 
         return idx[:, index:] # 只返回生成的token
 
@@ -421,6 +778,7 @@ if __name__ == '__main__':
     args = ModelConfig(
         dim=1024,
         n_layers=18,
+        pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0,
     )
     # 实例化LLaMA2Model
     model = Transformer(args=args)
@@ -442,4 +800,4 @@ if __name__ == '__main__':
     print("Y shape :", Y.shape)
 
     # 将输入张量传入模型
-    output = model(X, Y)
+    output = model(X, Y)

docs/chapter5/code/model_sample.py

@@ -8,7 +8,7 @@ import argparse
 
 class TextGenerator:
     def __init__(self, 
-                 checkpoint='out/SkyWork_pretrain_768_12_6144.pth',  # 模型检查点路径
+                 checkpoint='./base_model_215M/pretrain_1024_18_6144.pth',  # 模型检查点路径
                  tokenizer_model_path='./tokenizer_k/',  # 分词器模型路径
                  seed=42,  # 随机种子，确保可重复性
                  device=None,  # 设备，优先使用 CUDA，如果没有可用的 CUDA，则使用 CPU
@@ -33,10 +33,18 @@ class TextGenerator:
         # 根据 dtype 选择适当的自动混合精度上下文
         ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[self.dtype]
         self.ctx = nullcontext() if self.device_type == 'cpu' else torch.amp.autocast(device_type=self.device_type, dtype=ptdtype)
-        
+        # 初始化分词器
+        self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_model_path)  # 根据指定的路径加载分词器
+
         # 加载模型检查点文件
         checkpoint_dict = torch.load(self.checkpoint, map_location=self.device)  # 加载模型参数 # 初始化模型参数
-        self.model = Transformer(ModelConfig(dim=1024, n_layers=18))  # 实例化 Transformer 模型
+        self.model = Transformer(
+            ModelConfig(
+                dim=1024,
+                n_layers=18,
+                pad_token_id=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0
+            )
+        )  # 实例化 Transformer 模型
         sunwanted_prefix = '_orig_mod.'
         for k, v in list(checkpoint_dict.items()):
             if k.startswith(sunwanted_prefix):
@@ -50,12 +58,10 @@ class TextGenerator:
         self.model.eval()
         # 将模型放置到正确的设备上（GPU 或 CPU）
         self.model.to(self.device)
-        # 初始化分词器
-        self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_model_path)  # 根据指定的路径加载分词器
 
     def chat_template(self, prompt):
         message = [
-            {"role": "system", "content": "你是一个AI助手。"},
+            {"role": "system", "content": "你是一个AI助手，你的名字叫小明。"},
             {"role": "user", "content": prompt}
         ]
         return self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
@@ -126,18 +132,6 @@ class TextGenerator:
         return generated_texts  # 返回生成的文本样本
     
 if __name__ == "__main__":
-    print("\n ------------------- SFT Sample ------------------- \n")
-    sft_prompt_datas = [
-        '你好呀',
-        "中国的首都是哪里？",
-        "1+1等于多少？",
-    ]
-    generator = TextGenerator(checkpoint='./BeelGroup_sft_model_215M/sft_dim1024_layers18_vocab_size6144.pth')  # 初始化生成器
-    for i in range(len(sft_prompt_datas)):
-        samples = generator.sft_sample(start=sft_prompt_datas[i], num_samples=1, max_new_tokens=512, temperature=0.75)
-        print(f"\nSample {i+1}:\nQuestion: {sft_prompt_datas[i]} \nAI answer: {samples[0]}\n{'-'*20}")  # 打印生成的样本并用分隔线分割
-
-    
     print("------------------- Pretrain Sample ------------------- \n")
 
     pretrain_prompt_datas = [
@@ -145,7 +139,22 @@ if __name__ == "__main__":
         '<|im_start|>中国矿业大学（北京）地球科学与测绘工程学院',
     ]
 
-    generator = TextGenerator(checkpoint='./base_monkey_215M/pretrain_1024_18_6144.pth')  # 初始化生成器
+    generator = TextGenerator(checkpoint='./base_model_215M/pretrain_1024_18_6144.pth')  # 初始化生成器
     for i in range(len(pretrain_prompt_datas)):
-        samples = generator.pretrain_sample(start=pretrain_prompt_datas[i], num_samples=1, max_new_tokens=120, temperature=1.0)
-        print(f"\nSample {i+1}:\n{pretrain_prompt_datas[i]}{samples[0]}\n{'-'*20}")  # 打印生成的样本并用分隔线分割
+        samples = generator.pretrain_sample(start=pretrain_prompt_datas[i], num_samples=1, max_new_tokens=120, temperature=0.75)
+        print(f"\nSample {i+1}:\n{pretrain_prompt_datas[i]}{samples[0]}\n{'-'*20}")  # 打印生成的样本并用分隔线分割
+
+    print("\n ------------------- SFT Sample ------------------- \n")
+
+    sft_prompt_datas = [
+        '你好呀',
+        "中国的首都是哪里？",
+        "1+12等于多少？",
+        "你是谁？"
+    ]
+    generator = TextGenerator(checkpoint='./sft_model_215M/sft_dim1024_layers18_vocab_size6144.pth')  # 初始化生成器
+    for i in range(len(sft_prompt_datas)):
+        samples = generator.sft_sample(start=sft_prompt_datas[i], num_samples=1, max_new_tokens=128, temperature=0.6)
+        print(f"\nSample {i+1}:\nQuestion: {sft_prompt_datas[i]} \nAI answer: {samples[0]}\n{'-'*20}")  # 打印生成的样本并用分隔线分割
+
+    

docs/chapter5/code/windows_download_dataset.sh

@@ -0,0 +1,35 @@
+# Windows下载方式
+# 使用PowerShell下载
+# 暂时为当前PowerShell界面设置环境，关闭Powershell环境自动消失
+$env:HF_ENDPOINT = "https://hf-mirror.com"
+
+# 将\path\to\your\dataset替换成想要下载dataset目录地址
+$dataset_dir = "\path\to\your\dataset"
+
+# 需要预先安装modelscope，使用pip install modelscope安装
+modelscope download --dataset ddzhu123/seq-monkey mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2 --local_dir "$dataset_dir"
+
+tar -xvf "$dataset_dir\mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2" -C "$dataset_dir"
+
+huggingface-cli download `
+  --repo-type dataset `
+  --resume-download `
+  BelleGroup/train_3.5M_CN `
+  --local-dir "$dataset_dir\BelleGroup"
+
+# 使用CMD下载
+# 暂时为当前CMD界面设置环境，关闭CMD环境自动消失
+set HF_ENDPOINT=https://hf-mirror.com
+
+# 将\path\to\your\dataset替换成想要下载dataset目录地址
+set dataset_dir=\path\to\your\dataset
+
+modelscope download --dataset ddzhu123/seq-monkey mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2 --local_dir %dataset_dir%
+
+tar -xvf "%dataset_dir%\mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2" -C "%dataset_dir%"
+
+huggingface-cli download ^
+  --repo-type dataset ^
+  --resume-download ^
+  BelleGroup/train_3.5M_CN ^
+  --local-dir "%dataset_dir%\BelleGroup"

docs/chapter5/第五章 动手搭建大模型.md

@@ -1,14 +1,16 @@
 # 第五章 动手搭建大模型
 
+
+
 ## 5.1 动手实现一个 LLaMA2 大模型
 
-Meta（原Facebook）于2023年2月发布第一款基于Transformer结构的大型语言模型LLaMA，并于同年7月发布同系列模型LLaMA2。我们在第四章已经学习了解的了LLM，记忆如何训练LLM等等。那本小节我们就来学习，如何动手写一个LLaMA2模型。
+Meta（原Facebook）于2023年2月发布第一款基于Transformer结构的大型语言模型LLaMA，并于同年7月发布同系列模型LLaMA2。我们在第四章已经学习和了解了LLM，以及如何训练LLM等内容。本小节我们就来学习如何动手实现一个LLaMA2模型。
 
-LLaMA2 模型结构如下图5.0所示：
+LLaMA2 模型结构如下图5.1所示：
 
 <div align='center'>
-    <img src="../images/5-images/LLama2.png" alt="alt text" width="100%">
-    <p>图 5.0 LLaMA2结构</p>
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/5-images/LLama2.png" alt="alt text" width="100%">
+    <p>图 5.1 LLaMA2结构</p>
 </div>
 
 ### 5.1.1 定义超参数
@@ -51,6 +53,8 @@ class ModelConfig(PretrainedConfig):
         super().__init__(**kwargs)
 ```
 
+> 在以下代码中出现 `args` 时，即默认为以上 `ModelConfig` 参数配置。
+
 我们来看一下其中的一些超参数的含义，比如`dim`是模型维度，`n_layers`是Transformer的层数，`n_heads`是注意力机制的头数，`vocab_size`是词汇表大小，`max_seq_len`是输入的最大序列长度等等。上面的代码中也对每一个参数做了详细的注释，在后面的代码中我们会根据这些超参数来构建我们的模型。
 
 ### 5.1.2 构建 RMSNorm
@@ -95,7 +99,7 @@ class RMSNorm(nn.Module):
         return output * self.weight
 ```
 
-并且，我们可以用下面的代码来对`RMSNorm`模块进行测试，可以看到代码最终输出的形状为`torch.Size([1, 50, 288])`，与我们输入的形状一致，说明模块的实现是正确的，归一化并不会改变输入的形状。
+并且，我们可以用下面的代码来对`RMSNorm`模块进行测试，可以看到代码最终输出的形状为`torch.Size([1, 50, 768])`，与我们输入的形状一致，说明模块的实现是正确的，归一化并不会改变输入的形状。
 
 ```python
 norm = RMSNorm(args.dim, args.norm_eps)
@@ -104,13 +108,18 @@ output = norm(x)
 print(output.shape)
 
 out:
-orch.Size([1, 50, 288])
+torch.Size([1, 50, 768])
 ```
 
 ### 5.1.3 构建 LLaMA2 Attention
 
 在 LLaMA2 模型中，虽然只有 LLaMA2-70B模型使用了分组查询注意力机制（Grouped-Query Attention，GQA），但我们依然选择使用 GQA 来构建我们的 LLaMA Attention 模块，它可以提高模型的效率，并节省一些显存占用。
 
+<div align='center'>
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/5-images/llama2-attention.png" alt="alt text" width="50%">
+    <p>图 5.2 LLaMA2 Attention 结构</p>
+</div>
+
 #### 5.1.3.1 repeat_kv
 
 在 LLaMA2 模型中，我们需要将键和值的维度扩展到和查询的维度一样，这样才能进行注意力计算。我们可以通过如下代码实现`repeat_kv`：
@@ -547,20 +556,20 @@ class Transformer(PreTrainedModel):
         elif isinstance(module, nn.Embedding):
             torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
     
-    def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, **keyargs) -> torch.Tensor:
+    def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
         """
         - tokens: Optional[torch.Tensor], 输入 token 张量。
         - targets: Optional[torch.Tensor], 目标 token 张量。
         - kv_cache: bool, 是否使用键值缓存。
-        - keyargs: 其他关键字参数。
+        - kwargs: 其他关键字参数。
 
         - self.OUT: CausalLMOutputWithPast, 包含 logits 和损失。
         """
 
-        if 'input_ids' in keyargs:
-            tokens = keyargs['input_ids']
-        if 'attention_mask' in keyargs:
-            targets = keyargs['attention_mask']
+        if 'input_ids' in kwargs:
+            tokens = kwargs['input_ids']
+        if 'labels' in kwargs:
+            targets = kwargs['labels']
 
         # 前向传播函数
         _bsz, seqlen = tokens.shape
@@ -654,7 +663,7 @@ torch.Size([1, 1, 6144])
 
 在自然语言处理 (NLP) 中，Tokenizer 是一种将文本分解为较小单位（称为 token）的工具。这些 token 可以是词、子词、字符，甚至是特定的符号。Tokenization 是 NLP 中的第一步，直接影响后续处理和分析的效果。不同类型的 tokenizer 适用于不同的应用场景，以下是几种常见的 tokenizer 及其特点。
 
-### 5.3.1 Word-based Tokenizer
+### 5.2.1 Word-based Tokenizer
 
 **Word-based Tokenizer** 是最简单和直观的一种分词方法。它将文本按空格和标点符号分割成单词。这种方法的优点在于其简单和直接，易于实现，且与人类对语言的直觉相符。然而，它也存在一些明显的缺点，如无法处理未登录词（OOV，out-of-vocabulary）和罕见词，对复合词（如“New York”）或缩略词（如“don't”）的处理也不够精细。此外，Word-based Tokenizer 在处理不同语言时也会遇到挑战，因为一些语言（如中文、日文）没有显式的单词分隔符。
 
@@ -759,34 +768,26 @@ from typing import Generator
 
 #### Step 2: 加载训练数据
 
-我们使用 `datasets.load_dataset()` 库加载一个英文文本数据集，用于训练 BPE Tokenizer。这里我们使用 `wikitext` 数据集，包含了维基百科的文章文本。
+这里我们使用与预训练相同的数据集（出门问问序列猴子开源数据集）训练tokenizer，可使用`code/download_dataset.sh` 和 `code/deal_dataset.py` 下载和预处理数据集。
+
+> 注：由于数据集过大，可能会导致在训练过程中内存不足。因为本项目为学习目的，建议学习者手动分割小部分数据集用于训练验证，笔者也在 Github 仓库中存放了训练好的 tokenizer，可以直接使用。
 
 ```python
-dataset = load_dataset("wikitext", "wikitext-103-v1", split="train")
-
-# 准备训练数据
-def batch_iterator(batch_size=1000):
-    for i in range(0, len(dataset), batch_size):
-        yield dataset[i:i + batch_size]["text"]
-```
-
-如果你使用本地的文本数据集，可以将数据加载到一个列表中，然后传入 `batch_iterator()` 函数中。如下所示：
-
-```python
-def load_text_from_files(path_list):
-    text_data = []
-    for file_path in path_list:
-        with open(file_path, 'r', encoding='utf-8') as file:
-            text_data.extend(file.readlines())
-    return text_data
-
-def batch_iterator(text_data, batch_size=1000):
-    for i in range(0, len(text_data), batch_size):
-        yield text_data[i:i + batch_size]
-
-# 假设你的文件路径列表是
-path_list = ['text_data1.txt', 'text_data2.txt', 'text_data3.txt']
-text_data = load_text_from_files(path_list)
+def read_texts_from_jsonl(file_path: str) -> Generator[str, None, None]:
+    """读取JSONL文件并安全提取文本数据"""
+    with open(file_path, 'r', encoding='utf-8') as f:
+        for line_num, line in enumerate(f, 1):
+            try:
+                data = json.loads(line)
+                if 'text' not in data:
+                    raise KeyError(f"Missing 'text' field in line {line_num}")
+                yield data['text']
+            except json.JSONDecodeError:
+                print(f"Error decoding JSON in line {line_num}")
+                continue
+            except KeyError as e:
+                print(e)
+                continue
 ```
 
 #### Step 3: 创建配置文件
@@ -991,7 +992,7 @@ Special tokens preserved: False
 
 在前面的章节中，我们熟悉了各种大模型的模型结构，以及如如何训练Tokenizer。在本节中，我们将动手训练一个八千万参数的LLM。
 
-### 5.3.0 数据下载
+### 5.3.1 数据下载
 
 首先，我们需要下载预训练数据集。在这里，我们使用两个开源的数据集，包含了大量的中文对话数据，可以用于训练对话生成模型。
 
@@ -1045,7 +1046,7 @@ def convert_message(data):
     return message
 
 with open('BelleGroup_sft.jsonl', 'a', encoding='utf-8') as sft:
-    with open('BelleGroup/train_3.5M_CN.json', 'r') as f:
+    with open('BelleGroup/train_3.5M_CN.json', 'r', encoding='utf-8') as f:
         data = f.readlines()
         for item in tqdm(data, desc="Processing", unit="lines"):
             item = json.loads(item)
@@ -1053,7 +1054,7 @@ with open('BelleGroup_sft.jsonl', 'a', encoding='utf-8') as sft:
             sft.write(json.dumps(message, ensure_ascii=False) + '\n')
 ```
 
-### 5.3.1 训练Tokenize
+### 5.3.2 训练 Tokenizer
 
 首先，我们需要为文本处理训练一个Tokenizer。Tokenizer的作用是将文本转换为数字序列，以便模型能够理解和处理。我们使用的数据集是 [出门问问序列猴子开源数据集](https://www.modelscope.cn/datasets/ddzhu123/seq-monkey/files) ，这个数据集包含了大量的中文文本数据，可以用于训练Tokenizer。
 
@@ -1290,7 +1291,7 @@ Hello<|im_end|>
 Special tokens preserved: False
 ```
 
-### 5.3.2 Dataset
+### 5.3.3 Dataset
 
 #### PretrainDataset
 
@@ -1305,15 +1306,23 @@ class PretrainDataset(Dataset):
         self.data_path = data_path
         self.tokenizer = tokenizer
         self.max_length = max_length
-        self.padding = 0
-        with open(data_path, 'r', encoding='utf-8') as f:
-            self.data = f.readlines()
+        self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+        # 预计算每行的起始字节偏移量
+        self._offsets = []
+        with open(data_path, 'rb') as f:
+            self._offsets.append(0)
+            while f.readline():
+                self._offsets.append(f.tell())
+        self._total_lines = len(self._offsets) - 1  # 最后一个 tell() 是 EOF
 
     def __len__(self):
-        return len(self.data)
+        return self._total_lines
 
     def __getitem__(self, index: int):
-        sample = json.loads(self.data[index])
+        with open(self.data_path, 'rb') as f:
+            f.seek(self._offsets[index])
+            line = f.readline().decode('utf-8')
+        sample = json.loads(line)
         text = f"{self.tokenizer.bos_token}{sample['text']}"
         input_id = self.tokenizer(text).data['input_ids'][:self.max_length]
         text_len = len(input_id)
@@ -1330,11 +1339,11 @@ class PretrainDataset(Dataset):
         return torch.from_numpy(X), torch.from_numpy(Y), torch.from_numpy(loss_mask)
 ```
 
-在以上代码和图5.1可以看出，`Pretrain Dataset` 主要是将 `text` 通过 `tokenizer` 转换成 `input_id`，然后将 `input_id` 拆分成 `X` 和 `Y`，其中 `X` 为 `input_id` 的前 n-1 个元素，`Y` 为 `input_id` 的后 n-1 `个元素。loss_mask` 主要是用来标记哪些位置需要计算损失，哪些位置不需要计算损失。
+在以上代码和图5.3可以看出，`Pretrain Dataset` 主要是将 `text` 通过 `tokenizer` 转换成 `input_id`，然后将 `input_id` 拆分成 `X` 和 `Y`，其中 `X` 为 `input_id` 的前 n-1 个元素，`Y` 为 `input_id` 的后 n-1 `个元素。loss_mask` 主要是用来标记哪些位置需要计算损失，哪些位置不需要计算损失。
 
 <div align='center'>
-    <img src="../images/5-images/pretrain_dataset.png" alt="alt text" width="100%">
-    <p>图5.1 预训练损失函数计算</p>
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/5-images/pretrain_dataset.png" alt="alt text" width="100%">
+    <p>图5.3 预训练损失函数计算</p>
 </div>
 
 图中示例展示了当`max_length=9`时的处理过程：
@@ -1356,17 +1365,21 @@ class SFTDataset(Dataset):
         self.data_path = data_path
         self.tokenizer = tokenizer
         self.max_length = max_length
-        self.padding = 0
-        with open(data_path, 'r', encoding='utf-8') as f:
-            self.data = f.readlines()
+        self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+        self._offsets = []
+        with open(data_path, 'rb') as f:
+            self._offsets.append(0)
+            while f.readline():
+                self._offsets.append(f.tell())
+        self._total_lines = len(self._offsets) - 1 
 
     def __len__(self):
-        return len(self.data)
+        return self._total_lines
 
     def generate_loss_mask(self, input_ids):
         # 生成 loss mask, 0 表示不计算损失, 1 表示计算损失
         mask = [0] * len(input_ids)
-        a_sequence = [3, 1074, 537, 500, 203]  # <|im_start|>assistant\n
+        a_sequence = self.tokenizer("<|im_start|>assistant\n")['input_ids']  # <|im_start|>assistant\n
         a_length = len(a_sequence)
         n = len(input_ids)
         i = 0
@@ -1379,10 +1392,10 @@ class SFTDataset(Dataset):
                     match = False
                     break
             if match:
-                # 从子序列结束的位置开始查找第一个4, 4 为 <|im_end|> EOS id
+                # 从子序列结束的位置开始查找第一个 4 (eos_token_id)
                 j = None
                 for idx in range(i + a_length, n):
-                    if input_ids[idx] == 4:
+                    if input_ids[idx] == self.tokenizer.eos_token_id:
                         j = idx
                         break
                 if j is not None:
@@ -1400,7 +1413,10 @@ class SFTDataset(Dataset):
         return mask
 
     def __getitem__(self, index: int):
-        sample = json.loads(self.data[index])
+        with open(self.data_path, 'rb') as f:
+            f.seek(self._offsets[index])
+            line = f.readline().decode('utf-8')
+        sample = json.loads(line)
         text = self.tokenizer.apply_chat_template(sample, tokenize=False, add_generation_prompt=False)
         input_id = self.tokenizer(text).data['input_ids'][:self.max_length]
         text_len = len(input_id)
@@ -1417,17 +1433,17 @@ class SFTDataset(Dataset):
         return torch.from_numpy(X), torch.from_numpy(Y), torch.from_numpy(loss_mask)
 ```
 
-在 SFT 阶段，这里使用的是多轮对话数据集，所以就需要区分哪些位置需要计算损失，哪些位置不需要计算损失。在上面的代码中，我使用了一个 `generate_loss_mask` 函数来生成 `loss_mask`。这个函数主要是用来生成 `loss_mask`，其中 `loss_mask` 的生成规则是：当遇到 `|<im_start|>assistant\n` 时，就开始计算损失，直到遇到 `|<im_end|>` 为止。这样就可以保证我们的模型在 SFT 阶段只计算当前轮的对话内容，如图5.2所示。
+在 SFT 阶段，这里使用的是多轮对话数据集，所以就需要区分哪些位置需要计算损失，哪些位置不需要计算损失。在上面的代码中，我使用了一个 `generate_loss_mask` 函数来生成 `loss_mask`。这个函数主要是用来生成 `loss_mask`，其中 `loss_mask` 的生成规则是：当遇到 `|<im_start|>assistant\n` 时，就开始计算损失，直到遇到 `|<im_end|>` 为止。这样就可以保证我们的模型在 SFT 阶段只计算当前轮的对话内容，如图5.4所示。
 
 <div align='center'>
-    <img src="../images/5-images/sftdataset.png" alt="alt text" width="90%">
-    <p>图5.2 SFT 损失函数计算</p>
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/5-images/sftdataset.png" alt="alt text" width="90%">
+    <p>图5.4 SFT 损失函数计算</p>
 </div>
 
 可以看到，其实 SFT Dataset 和 Pretrain Dataset 的 `X` 和 `Y` 是一样的，只是在 SFT Dataset 中我们需要生成一个 `loss_mask` 来标记哪些位置需要计算损失，哪些位置不需要计算损失。 图中 `Input ids` 中的蓝色小方格就是AI的回答，所以是需要模型学习的地方。所以在 `loss_mask` 中，蓝色小方格对应的位置是黄色，其他位置是灰色。在代码 `loss_mask` 中的 1 对应的位置计算损失，0 对应的位置不计算损失。
 
 
-### 5.3.3 预训练
+### 5.3.4 预训练
 
 在数据预处理完成后，我们就可以开始训练模型了。我们使用的模型是一个和LLama2结构一样的 Decoder only Transformer模型，使用Pytorch实现。相关代码在`code/k_model.py`文件中。此处不再赘述，源码中有详细的中文注释，且我们在之前的文章中也有详细的介绍。
 
@@ -1480,174 +1496,298 @@ class SFTDataset(Dataset):
 
 ```python
 def get_lr(it, all):
-    warmup_iters = args.warmup_iters
-    lr_decay_iters = all
-    min_lr = args.learning_rate / 10
+    """
+    计算当前迭代的学习率，使用余弦退火调度策略
+    
+    学习率调度策略：
+    1. Warmup阶段：学习率从0线性增长到目标学习率
+    2. 余弦退火阶段：学习率按余弦函数衰减到最小学习率
+    3. 超出训练步数后：保持最小学习率
+    
+    Args:
+        it (int): 当前迭代步数
+        all (int): 总迭代步数
+        
+    Returns:
+        float: 当前步数对应的学习率
+    """
+    warmup_iters = args.warmup_iters  # 预热迭代次数
+    lr_decay_iters = all  # 学习率衰减的总迭代次数
+    min_lr = args.learning_rate / 10  # 最小学习率，为初始学习率的1/10
 
+    # Warmup阶段：线性增长
     if it < warmup_iters:
         return args.learning_rate * it / warmup_iters
     
+    # 超出训练步数：保持最小学习率
     if it > lr_decay_iters:
         return min_lr
     
+    # 余弦退火阶段
     decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
     assert 0 <= decay_ratio <= 1
-    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # 余弦系数
     return min_lr + coeff * (args.learning_rate - min_lr)
 
 def train_epoch(epoch):
-    start_time = time.time()
+    """
+    训练一个epoch的函数
+    
+    实现了完整的训练循环，包括：
+    1. 数据加载和设备转移
+    2. 动态学习率调整
+    3. 前向传播和损失计算
+    4. 梯度累积和反向传播
+    5. 梯度裁剪和优化器更新
+    6. 日志记录和模型保存
+    
+    Args:
+        epoch (int): 当前epoch编号
+    """
+    start_time = time.time()  # 记录开始时间
+    
+    # 遍历数据加载器中的每个batch
     for step, (X, Y, loss_mask) in enumerate(train_loader):
-        X = X.to(args.device)
-        Y = Y.to(args.device)
-        loss_mask = loss_mask.to(args.device)
+        # 将数据转移到指定设备（GPU/CPU）
+        X = X.to(args.device)  # 输入序列
+        Y = Y.to(args.device)  # 目标序列
+        loss_mask = loss_mask.to(args.device)  # 损失掩码，用于忽略padding token
 
+        # 计算当前步骤的学习率
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch)
+        # 更新优化器中所有参数组的学习率
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
+        # 使用混合精度训练上下文
         with ctx:
+            # 前向传播
             out = model(X, Y)
+            # 计算损失并除以累积步数（用于梯度累积）
             loss = out.last_loss / args.accumulation_steps
+            # 将loss_mask展平为一维
             loss_mask = loss_mask.view(-1)
+            # 应用掩码计算有效损失（忽略padding位置）
             loss = torch.sum(loss * loss_mask) / loss_mask.sum()
 
+        # 使用scaler进行混合精度的反向传播
         scaler.scale(loss).backward()
 
+        # 每accumulation_steps步执行一次优化器更新
         if (step + 1) % args.accumulation_steps == 0:
+            # 取消梯度缩放，准备梯度裁剪
             scaler.unscale_(optimizer)
+            # 梯度裁剪，防止梯度爆炸
             torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
 
+            # 执行优化器步骤
             scaler.step(optimizer)
+            # 更新scaler的缩放因子
             scaler.update()
 
+            # 清零梯度，set_to_none=True可以节省内存
             optimizer.zero_grad(set_to_none=True)
 
+        # 每log_interval步记录一次日志
         if step % args.log_interval == 0:
             spend_time = time.time() - start_time
+            # 打印训练进度信息
             Logger(
-                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
+                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min;'.format(
                     epoch + 1,
                     args.epochs,
                     step,
                     iter_per_epoch,
-                    loss.item() * args.accumulation_steps,
+                    loss.item() * args.accumulation_steps,  # 恢复真实的loss值
                     optimizer.param_groups[-1]['lr'],
                     spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
+            
+            # 如果启用SwanLab，记录训练指标
             if args.use_swanlab:
                 swanlab.log({
                     "loss": loss.item() * args.accumulation_steps,
                     "lr": optimizer.param_groups[-1]['lr']
                 })
 
+        # 每save_interval步保存一次模型
         if (step + 1) % args.save_interval == 0:
-            model.eval()
+            model.eval()  # 切换到评估模式
+            # 构建检查点文件名
             ckp = f'{args.save_dir}/pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}.pth'
 
-            # 处理多卡保存
+            # 处理多卡保存：如果是DataParallel模型，需要访问.module属性
             state_dict = model.module.state_dict() if isinstance(model, torch.nn.DataParallel) else model.state_dict()
             torch.save(state_dict, ckp)
-            model.train()
+            model.train()  # 切换回训练模式
         
+        # 每20000步保存一个带步数标记的检查点
         if (step + 1) % 20000 == 0:
             model.eval()
+            # 构建带步数的检查点文件名
             ckp = f'{args.save_dir}/pretrain_{lm_config.dim}_{lm_config.n_layers}_{lm_config.vocab_size}_step{step+1}.pth'
 
+            # 保存模型状态字典
             state_dict = model.module.state_dict() if isinstance(model, torch.nn.DataParallel) else model.state_dict()
             torch.save(state_dict, ckp)
             model.train()
 
 
 def init_model():
+    """
+    初始化模型和分词器
+    
+    功能包括：
+    1. 加载预训练的分词器
+    2. 创建Transformer模型
+    3. 设置多GPU并行训练（如果可用）
+    4. 将模型移动到指定设备
+    5. 统计并打印模型参数量
+    
+    Returns:
+        tuple: (model, tokenizer) 初始化后的模型和分词器
+    """
     def count_parameters(model):
+        """
+        统计模型中可训练参数的数量
+        
+        Args:
+            model: PyTorch模型
+            
+        Returns:
+            int: 可训练参数总数
+        """
         return sum(p.numel() for p in model.parameters() if p.requires_grad)
 
+    # 从本地路径加载预训练的分词器
     tokenizer = AutoTokenizer.from_pretrained('./tokenizer_k/')
 
+    # 根据配置创建Transformer模型
     model = Transformer(lm_config)
     
-    # 多卡初始化
+    # 多卡初始化：检查可用GPU数量并设置DataParallel
     num_gpus = torch.cuda.device_count()
     if num_gpus > 1:
         Logger(f"Using {num_gpus} GPUs with DataParallel!")
+        # 使用DataParallel包装模型以支持多GPU训练
         model = torch.nn.DataParallel(model)
     
+    # 将模型移动到指定设备（GPU或CPU）
     model = model.to(args.device)
+    
+    # 计算并打印模型参数量（以百万为单位）
     Logger(f'LLM总参数量：{count_parameters(model) / 1e6:.3f} 百万')
     return model, tokenizer
 
+
 if __name__ == "__main__":
+    # ==================== 命令行参数解析 ====================
     parser = argparse.ArgumentParser(description="Tiny-LLM Pretraining")
-    parser.add_argument("--out_dir", type=str, default="output", help="Output directory")
-    parser.add_argument("--epochs", type=int, default=1, help="Number of epochs")
-    parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
-    parser.add_argument("--learning_rate", type=float, default=2e-4, help="Learning rate")
-    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="Device to use")
-    parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type")
-    parser.add_argument("--use_swanlab", type=bool, default=True, help="Use Weights & Biases")
-    parser.add_argument("--num_workers", type=int, default=8, help="Number of workers for data loading")
-    parser.add_argument("--data_path", type=str, default="", help="Path to training data")
-    parser.add_argument("--accumulation_steps", type=int, default=8, help="Gradient accumulation steps")
-    parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold")
-    parser.add_argument("--warmup_iters", type=int, default=0, help="Number of warmup iterations")
-    parser.add_argument("--log_interval", type=int, default=100, help="Logging interval")
-    parser.add_argument("--save_interval", type=int, default=1000, help="Model saving interval")
-    # 添加多卡参数
-    parser.add_argument("--gpus", type=str, default='0,1', help="Comma-separated GPU IDs (e.g. '0,1,2')")
+    
+    # 基础训练参数
+    parser.add_argument("--out_dir", type=str, default="base_model_215M", help="模型输出目录")
+    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=64, help="批次大小")
+    parser.add_argument("--learning_rate", type=float, default=2e-4, help="学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型")
+    
+    # 实验跟踪和数据加载参数
+    parser.add_argument("--use_swanlab", action="store_true", help="是否使用SwanLab进行实验跟踪")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载的工作进程数")
+    parser.add_argument("--data_path", type=str, default="./seq_monkey_datawhale.jsonl", help="训练数据路径")
+    
+    # 训练优化参数
+    parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--warmup_iters", type=int, default=0, help="学习率预热迭代次数")
+    
+    # 日志和保存参数
+    parser.add_argument("--log_interval", type=int, default=100, help="日志记录间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
+    
+    # 多GPU训练参数
+    parser.add_argument("--gpus", type=str, default='0,1,2,3,4,5,6,7', help="使用的GPU ID，用逗号分隔 (例如: '0,1,2')")
 
     args = parser.parse_args()
 
-    # 设置可见GPU
+    # ==================== GPU环境设置 ====================
+    # 设置可见的GPU设备
     if args.gpus is not None:
         os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
-        # 自动设置主设备为第一个GPU
+        # 自动设置主设备为第一个可用GPU
         if torch.cuda.is_available():
             args.device = "cuda:0"
         else:
             args.device = "cpu"
 
+    # ==================== 实验跟踪初始化 ====================
     if args.use_swanlab:
-        swanlab.login(api_key='your key')
+        # 注意：使用前需要先登录 swanlab.login(api_key='your key')
         run = swanlab.init(
-            project="Tiny-LLM",
-            experiment_name="Pretrain-215M",
-            config=args,
+            project="Happy-LLM",  # 项目名称
+            experiment_name="Pretrain-215M",  # 实验名称
+            config=args,  # 保存所有超参数
         )
 
+    # ==================== 模型配置 ====================
+    # 定义语言模型的配置参数
     lm_config = ModelConfig(
-        dim=1024,
-        n_layers=18,
+        dim=1024,      # 模型维度
+        n_layers=18,   # Transformer层数
     )
-    max_seq_len = lm_config.max_seq_len
-    args.save_dir = os.path.join(args.out_dir)
-    os.makedirs(args.save_dir, exist_ok=True)
+
+    # ==================== 训练环境设置 ====================
+    max_seq_len = lm_config.max_seq_len  # 最大序列长度
+    args.save_dir = os.path.join(args.out_dir)  # 模型保存目录
+    
+    # 创建必要的目录
     os.makedirs(args.out_dir, exist_ok=True)
+    
+    # 设置随机种子以确保结果可复现
     torch.manual_seed(42)
+    
+    # 确定设备类型（用于选择合适的上下文管理器）
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
+    # 设置混合精度训练的上下文管理器
+    # CPU训练时使用nullcontext，GPU训练时使用autocast
     ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
 
+    # ==================== 模型和数据初始化 ====================
+    # 初始化模型和分词器
     model, tokenizer = init_model()
     
+    # 创建训练数据集
     train_ds = PretrainDataset(args.data_path, tokenizer, max_length=max_seq_len)
+    
+    # 创建数据加载器
     train_loader = DataLoader(
         train_ds,
-        batch_size=args.batch_size,
-        pin_memory=True,
-        drop_last=False,
-        shuffle=True,
-        num_workers=args.num_workers
+        batch_size=args.batch_size,  # 批次大小
+        pin_memory=True,             # 将数据加载到固定内存中，加速GPU传输
+        drop_last=False,             # 不丢弃最后一个不完整的批次
+        shuffle=True,                # 随机打乱数据
+        num_workers=args.num_workers # 数据加载的并行工作进程数
     )
 
+    # ==================== 优化器和训练组件初始化 ====================
+    # 初始化混合精度训练的梯度缩放器
+    # 只有在使用float16或bfloat16时才启用
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    
+    # 初始化Adam优化器
     optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
 
+    # ==================== 开始训练 ====================
+    # 计算每个epoch的迭代次数
     iter_per_epoch = len(train_loader)
+    
+    # 开始训练循环
     for epoch in range(args.epochs):
         train_epoch(epoch)
 ```
 
-### 5.3.4 SFT 训练
+### 5.3.5 SFT 训练
 
 SFT 训练和预训练的代码基本一样，只是导入的 Dataset 不一样。在这里我们使用的是 SFTDataset，用于多轮对话的训练。
 
@@ -1671,13 +1811,18 @@ from dataset import SFTDataset
 
 import swanlab
 
+# 忽略警告
 warnings.filterwarnings('ignore')
 
 
 def Logger(content):
+    """日志记录器"""
     print(content)
 
 def get_lr(it, all):
+    """获取学习率"""
+    # 1) linear warmup for warmup_iters steps
+    # 1) 预热迭代的线性预热
     warmup_iters = args.warmup_iters
     lr_decay_iters = all
     min_lr = args.learning_rate / 10
@@ -1685,33 +1830,42 @@ def get_lr(it, all):
     if it < warmup_iters:
         return args.learning_rate * it / warmup_iters
     
+    # 2) if it > lr_decay_iters, return min learning rate
+    # 2) 如果迭代次数超过学习率衰减迭代次数，则返回最小学习率
     if it > lr_decay_iters:
         return min_lr
     
+    # 3) in between, use cosine decay down to min learning rate
+    # 3) 在两者之间，使用余弦衰减至最小学习率
     decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
     assert 0 <= decay_ratio <= 1
     coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
     return min_lr + coeff * (args.learning_rate - min_lr)
 
 def train_epoch(epoch):
+    """训练一个epoch"""
     start_time = time.time()
     for step, (X, Y, loss_mask) in enumerate(train_loader):
         X = X.to(args.device)
         Y = Y.to(args.device)
         loss_mask = loss_mask.to(args.device)
 
+        # 获取学习率并更新优化器
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
+        # 前向传播
         with ctx:
             out = model(X, Y)
             loss = out.last_loss / args.accumulation_steps
             loss_mask = loss_mask.view(-1)
             loss = torch.sum(loss * loss_mask) / loss_mask.sum()
 
+        # 反向传播
         scaler.scale(loss).backward()
 
+        # 更新权重
         if (step + 1) % args.accumulation_steps == 0:
             scaler.unscale_(optimizer)
             torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
@@ -1721,6 +1875,7 @@ def train_epoch(epoch):
 
             optimizer.zero_grad(set_to_none=True)
 
+        # 打印日志
         if step % args.log_interval == 0:
             spend_time = time.time() - start_time
             Logger(
@@ -1738,6 +1893,7 @@ def train_epoch(epoch):
                     "lr": optimizer.param_groups[-1]['lr']
                 })
 
+        # 保存模型
         if (step + 1) % args.save_interval == 0:
             model.eval()
             ckp = f'{args.save_dir}/sft_dim{lm_config.dim}_layers{lm_config.n_layers}_vocab_size{lm_config.vocab_size}.pth'
@@ -1747,6 +1903,7 @@ def train_epoch(epoch):
             torch.save(state_dict, ckp)
             model.train()
         
+        # 定期保存模型
         if (step + 1) % 20000 == 0:
             model.eval()
             ckp = f'{args.save_dir}/sft_dim{lm_config.dim}_layers{lm_config.n_layers}_vocab_size{lm_config.vocab_size}_step{step+1}.pth'
@@ -1757,14 +1914,19 @@ def train_epoch(epoch):
 
 
 def init_model():
+    """初始化模型"""
     def count_parameters(model):
+        """计算模型参数量"""
         return sum(p.numel() for p in model.parameters() if p.requires_grad)
 
+    # 加载分词器
     tokenizer = AutoTokenizer.from_pretrained('./tokenizer_k/')
 
+    # 初始化模型
     model = Transformer(lm_config)
 
-    ckp = './base_monkey_215M/pretrain_1024_18_6144.pth'
+    # 加载预训练权重
+    ckp = './base_model_215M/pretrain_1024_18_6144.pth'
     state_dict = torch.load(ckp, map_location=args.device)
     unwanted_prefix = '_orig_mod.'
     for k, v in list(state_dict.items()):
@@ -1785,22 +1947,22 @@ def init_model():
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Tiny-LLM Pretraining")
-    parser.add_argument("--out_dir", type=str, default="output", help="Output directory")
-    parser.add_argument("--epochs", type=int, default=1, help="Number of epochs")
-    parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
-    parser.add_argument("--learning_rate", type=float, default=2e-4, help="Learning rate")
-    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="Device to use")
-    parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type")
-    parser.add_argument("--use_swanlab", type=bool, default=True, help="Use Weights & Biases")
-    parser.add_argument("--num_workers", type=int, default=4, help="Number of workers for data loading")
-    parser.add_argument("--data_path", type=str, default="", help="Path to training data")
-    parser.add_argument("--accumulation_steps", type=int, default=4, help="Gradient accumulation steps")
-    parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold")
-    parser.add_argument("--warmup_iters", type=int, default=0, help="Number of warmup iterations")
-    parser.add_argument("--log_interval", type=int, default=100, help="Logging interval")
-    parser.add_argument("--save_interval", type=int, default=1000, help="Model saving interval")
+    parser.add_argument("--out_dir", type=str, default="sft_model_215M", help="输出目录")
+    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=64, help="批处理大小")
+    parser.add_argument("--learning_rate", type=float, default=2e-4, help="学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="使用的设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型")
+    parser.add_argument("--use_swanlab", action="store_true", help="是否使用SwanLab进行实验跟踪")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载的工作进程数")
+    parser.add_argument("--data_path", type=str, default="./BelleGroup_sft.jsonl", help="训练数据路径")
+    parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--warmup_iters", type=int, default=0, help="预热迭代次数")
+    parser.add_argument("--log_interval", type=int, default=100, help="日志记录间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
     # 添加多卡参数
-    parser.add_argument("--gpus", type=str, default='0,1', help="Comma-separated GPU IDs (e.g. '0,1,2')")
+    parser.add_argument("--gpus", type=str, default='0,1,2,3,4,5,6,7', help="逗号分隔的GPU ID (例如 '0,1,2')")
 
     args = parser.parse_args()
 
@@ -1813,29 +1975,32 @@ if __name__ == "__main__":
         else:
             args.device = "cpu"
 
+    # 初始化swanlab
     if args.use_swanlab:
-        swanlab.login(api_key='your key')
         run = swanlab.init(
-            project="Tiny-LLM",
-            experiment_name="BelleGropu-sft-215M",
+            project="Happy-LLM",
+            experiment_name="SFT-215M",
             config=args,
         )
 
+    # 模型配置
     lm_config = ModelConfig(
         dim=1024,
         n_layers=18,
     )
     max_seq_len = lm_config.max_seq_len
     args.save_dir = os.path.join(args.out_dir)
-    os.makedirs(args.save_dir, exist_ok=True)
     os.makedirs(args.out_dir, exist_ok=True)
     torch.manual_seed(42)
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
+    # 上下文管理器
     ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
 
+    # 初始化模型和分词器
     model, tokenizer = init_model()
     
+    # 创建数据集和数据加载器
     train_ds = SFTDataset(args.data_path, tokenizer, max_length=max_seq_len)
     train_loader = DataLoader(
         train_ds,
@@ -1846,16 +2011,18 @@ if __name__ == "__main__":
         num_workers=args.num_workers
     )
 
+    # 缩放器和优化器
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
-    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
 
+    # 开始训练
     iter_per_epoch = len(train_loader)
     for epoch in range(args.epochs):
         train_epoch(epoch)
 ```
 
 
-### 5.3.4 使用模型生成文本
+### 5.3.6 使用模型生成文本
 
 在模型训练完成后，会在`output`目录下生成模型文件，这个文件就是我们训练好的模型。我们可以使用以下命令生成文本。
 
@@ -1866,9 +2033,17 @@ python model_sample.py
 我们来看下`model_sample.py`文件中的代码，这个文件中定义了一个`TextGenerator`类，用于生成文本。
 
 ```python
+import os
+import pickle
+from contextlib import nullcontext
+import torch
+from k_model import ModelConfig, Transformer
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import argparse
+
 class TextGenerator:
     def __init__(self, 
-                 checkpoint='out/SkyWork_pretrain_768_12_6144.pth',  # 模型检查点路径
+                 checkpoint='./base_model_215M/pretrain_1024_18_6144.pth',  # 模型检查点路径
                  tokenizer_model_path='./tokenizer_k/',  # 分词器模型路径
                  seed=42,  # 随机种子，确保可重复性
                  device=None,  # 设备，优先使用 CUDA，如果没有可用的 CUDA，则使用 CPU
@@ -1915,7 +2090,7 @@ class TextGenerator:
 
     def chat_template(self, prompt):
         message = [
-            {"role": "system", "content": "你是一个AI助手。"},
+            {"role": "system", "content": "你是一个AI助手，你的名字叫小明。"},
             {"role": "user", "content": prompt}
         ]
         return self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
@@ -1984,6 +2159,33 @@ class TextGenerator:
                     generated_texts.append(self.tokenizer.decode(y[0].tolist()))  # 解码生成的 token 序列为可读文本
         
         return generated_texts  # 返回生成的文本样本
+    
+if __name__ == "__main__":
+    print("------------------- Pretrain Sample ------------------- \n")
+
+    pretrain_prompt_datas = [
+        '<|im_start|>北京大学是',
+        '<|im_start|>中国矿业大学（北京）地球科学与测绘工程学院',
+    ]
+
+    generator = TextGenerator(checkpoint='./base_model_215M/pretrain_1024_18_6144.pth')  # 初始化生成器
+    for i in range(len(pretrain_prompt_datas)):
+        samples = generator.pretrain_sample(start=pretrain_prompt_datas[i], num_samples=1, max_new_tokens=120, temperature=0.75)
+        print(f"\nSample {i+1}:\n{pretrain_prompt_datas[i]}{samples[0]}\n{'-'*20}")  # 打印生成的样本并用分隔线分割
+
+    print("\n ------------------- SFT Sample ------------------- \n")
+
+    sft_prompt_datas = [
+        '你好呀',
+        "中国的首都是哪里？",
+        "1+12等于多少？",
+        "你是谁？"
+    ]
+    generator = TextGenerator(checkpoint='./sft_model_215M/sft_dim1024_layers18_vocab_size6144.pth')  # 初始化生成器
+    for i in range(len(sft_prompt_datas)):
+        samples = generator.sft_sample(start=sft_prompt_datas[i], num_samples=1, max_new_tokens=128, temperature=0.6)
+        print(f"\nSample {i+1}:\nQuestion: {sft_prompt_datas[i]} \nAI answer: {samples[0]}\n{'-'*20}")  # 打印生成的样本并用分隔线分割
+
 ```
 
 最后我们来看一下模型输出的结果：
@@ -2023,9 +2225,15 @@ Sample 2:
 --------------------
 ```
 
-到这里，我们的模型就训绽完成了，恭喜你训练了一个属于你自己的大模型。
+到这里，我们的模型就训练完成了，恭喜你训练了一个属于你自己的大模型。
+
+> 大家在训练的时候可以将 batch 调的低一些，这样可以减少显存的占用，避免显存不足的问题。当然这样会增加训练时间，可以根据自己的显卡显存大小来调整 batch 的大小。实测 Pretrain batch 为 4 的情况下只需要 7G 显存，训练时长预计 533 小时。作者是在 8卡4090 上进行训练的，预训练一共耗时 46 小时，SFT 阶段在 BelleGroup 350万条中文指令训练 24 小时。
+
+作者也在魔搭平台上传了本章节训来的模型，如果大家的设备不足以训练大模型，大家也可以在魔搭平台下载模型进行调试和模型体验。模型下载地址如下：
+
+> *ModelScope 模型下载地址：[🤖 ModelScope](https://www.modelscope.cn/collections/Happy-LLM-e98b91b10b684a)*    
+> *ModelScope 创空间体验地址：[🤖 创空间](https://www.modelscope.cn/studios/kmno4zx/happy_llm_215M_sft)*
 
-> 大家在训练的时候可以将 batch 调的低一些，这样可以减少显存的占用，避免显存不足的问题。当然这样会增加训练时间，可以根据自己的显卡显存大小来调整 batch 的大小。实测 Pretrain batch 为 4 的情况下只需要 7G 显存，训练时长预计 533 小时。作者是在 4卡A100上进行训练的，预训练一共耗时26小时，SFT 阶段在 BelleGroup 350万条中文指令训练 4 小时。
 
 **参考资料**
 
@@ -2041,4 +2249,4 @@ Sample 2:
 
 [6] Jingyao Gong. (2023). *minimind: Minimalist LLM implementation*. GitHub repository. https://github.com/jingyaogong/minimind  
 
-[7] Mobvoi. (2023). *seq-monkey-data: Llama2 training/inference data*. GitHub repository. https://github.com/mobvoi/seq-monkey-data
+[7] Mobvoi. (2023). *seq-monkey-data: Llama2 training/inference data*. GitHub repository. https://github.com/mobvoi/seq-monkey-data

docs/chapter6/code/finetune.py

@@ -33,7 +33,7 @@ from transformers import (
 import datetime
 from transformers.testing_utils import CaptureLogger
 from transformers.trainer_utils import get_last_checkpoint
-import wandb
+import swanlab
 from tqdm import tqdm
 
 
@@ -183,8 +183,8 @@ def main():
     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
     model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    # 初始化 WandB
-    wandb.init(project="sft", name="qwen-1.5b")
+    # 初始化 SwanLab
+    swanlab.init(project="sft", experiment_name="qwen-1.5b")
     
     # 设置日志
     logging.basicConfig(

docs/chapter6/code/finetune.sh

@@ -22,6 +22,6 @@ deepspeed finetune.py \
     --bf16 \
     --gradient_checkpointing \
     --deepspeed ./ds_config_zero2.json \
-    --report_to wandb
+    --report_to swanlab
     
     # --resume_from_checkpoint ${output_model}/checkpoint-20400 \

docs/chapter6/code/pretrain.py

@@ -30,7 +30,7 @@ from transformers import (
 import datetime
 from transformers.testing_utils import CaptureLogger
 from transformers.trainer_utils import get_last_checkpoint
-import wandb
+import swanlab
 
 
 logger = logging.getLogger(__name__)
@@ -95,8 +95,8 @@ def main():
     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
     model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-    # 初始化 WandB
-    wandb.init(project="pretrain", name="from_scrach")
+    # 初始化 SwanLab
+    swanlab.init(project="pretrain", experiment_name="from_scrach")
     
     # 设置日志
     logging.basicConfig(

docs/chapter6/code/pretrain.sh

@@ -24,6 +24,6 @@ deepspeed pretrain.py \
     --bf16 \
     --gradient_checkpointing \
     --deepspeed ./ds_config_zero2.json \
-    --report_to wandb
+    --report_to swanlab
     
     # --resume_from_checkpoint ${output_model}/checkpoint-20400 \

docs/chapter6/code/requirements.txt

@@ -0,0 +1,7 @@
+transformers
+datasets
+torch
+torchdata==0.9.0
+deepspeed
+pandas
+swanlab

docs/chapter6/第六章 大模型训练流程实践.md

@@ -1,5 +1,7 @@
 # 第六章 大模型训练流程实践
 
+第五章中，我们从零开始动手搭建了 LLaMA2 模型，并完整实现了其预训练和微调的全流程。在本章中，我们将深入探讨大模型的训练流程实践，重点介绍如何利用主流的大模型框架高效地进行模型训练和性能优化。
+
 ## 6.1 模型预训练
 
 在上一章，我们逐步拆解了 LLM 的模型结构及训练过程，从零手写实现了 LLaMA 模型结构及 Pretrain、SFT 全流程，更深入地理解了 LLM 的模型原理及训练细节。但是，在实际应用中，手写实现的 LLM 训练存在以下问题：
@@ -15,7 +17,7 @@
 Transformers 是由 Hugging Face 开发的 NLP 框架，通过模块化设计实现了对 BERT、GPT、LLaMA、T5、ViT 等上百种主流模型架构的统一支持。通过使用 Transformers，开发者无需重复实现基础网络结构，通过 AutoModel 类即可一键加载任意预训练，图6.1 为 Hugging Face Transformers 课程首页：
 
 <div align='center'>
-    <img src="../images/6-images/1-1.png" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/6-images/1-1.png" alt="alt text" width="90%">
     <p>图6.1 Hugging Face Transformers</p>
 </div>
 
@@ -24,7 +26,7 @@ Transformers 是由 Hugging Face 开发的 NLP 框架，通过模块化设计实
 对 LLM 时代的 NLP 研究者更为重要的是，HuggingFace 基于 Transformers 框架搭建了其庞大的 AI 社区，开放了数亿个预训练模型参数、25万+不同类型数据集，通过 Transformers、Dataset、Evaluate 等多个框架实现对预训练模型、数据集及评估函数的集成，从而帮助开发者可以便捷地使用任一预训练模型，在开源模型及数据集的基础上便捷地实现个人模型的开发与应用。
 
 <div align='center'>
-    <img src="../images/6-images/1-2.png" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/6-images/1-2.png" alt="alt text" width="90%">
     <p>图6.2 Hugging Face Transformers 模型社区</p>
 </div>
 
@@ -35,14 +37,14 @@ Transformers 是由 Hugging Face 开发的 NLP 框架，通过模块化设计实
 我们可以使用 transformers 的 AutoModel 类来直接初始化已经实现好的模型。对于任意预训练模型，其参数中都包含有模型的配置信息。如果是想要从头训练一个 LLM，可以使用一个已有的模型架构来直接初始化。这里，我们以 [Qwen-2.5-1.5B](https://huggingface.co/Qwen/Qwen2.5-1.5B/tree/main)的模型架构为例：
 
 <div align='center'>
-    <img src="../images/6-images/1-3.png" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/6-images/1-3.png" alt="alt text" width="90%">
     <p>图6.3 Qwen-2.5-1.5B</p>
 </div>
 
 该界面即为 HuggingFace 社区中的 Qwen-2.5-1.5B 模型参数，其中的 `config.json` 文件即是模型的配置信息，包括了模型的架构、隐藏层大小、模型层数等，如图6.4所示：
 
 <div align='center'>
-    <img src="../images/6-images/1-4.png" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/6-images/1-4.png" alt="alt text" width="90%">
     <p>图6.4 Qwen-2.5-1.5B config.json 文件</p>
 </div>
 
@@ -59,7 +61,7 @@ os.system('huggingface-cli download --resume-download Qwen/Qwen2.5-1.5B --local-
 如图6.5，此处的 “Qwen/Qwen2.5-1.5B”即为要下载模型的标识符，对于其他模型，可以直接复制 HuggingFace 上的模型名即可：
 
 <div align='center'>
-    <img src="../images/6-images/1-5.png" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/6-images/1-5.png" alt="alt text" width="90%">
     <p>图6.5 模型下载标识</p>
 </div>
 
@@ -87,7 +89,7 @@ model = AutoModelForCausalLM.from_config(config,trust_remote_code=True)
 由于 LLM 一般都是 CausalLM 架构，此处使用了 AutoModelForCausalLM 类进行加载。如果是用于分类任务训练，可使用 AutoModelForSequenceClassification 类来加载。查看该 model，图6.6可以看到其架构和定义的配置文件相同：
 
 <div align='center'>
-    <img src="../images/6-images/1-6.png" alt="alt text" width="70%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/6-images/1-6.png" alt="alt text" width="70%">
     <p>图6.6 模型结构输出结果</p>
 </div>
 
@@ -101,7 +103,7 @@ model = AutoModelForCausalLM.from_pretrained(model_name_or_path,trust_remote_cod
 
 类似的，直接使用 from_pretrained 方法加载即可，此处的 model_name_or_path 即为下载好的参数的本地路径。
 
-我们还需要初始化一个 tokenizer。此处，我们直接使用 Qwen-2.5-1.5B 对应的 tokenzier 参数即可：
+我们还需要初始化一个 tokenizer。此处，我们直接使用 Qwen-2.5-1.5B 对应的 tokenizer 参数即可：
 
 ```python
 # 加载一个预训练好的 tokenizer
@@ -130,7 +132,7 @@ ds["train"][0]
 ```
 
 <div align='center'>
-    <img src="../images/6-images/1-7.png" alt="alt text" width="100%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/6-images/1-7.png" alt="alt text" width="100%">
     <p>图6.7 数据集展示</p>
 </div>
 
@@ -250,7 +252,7 @@ trainer = Trainer(
 trainer.train()
 ```
 
-> 注：上述代码存放于 `./code/pretrian.ipynb` 文件中。
+> 注：上述代码存放于 `./code/pretrain.ipynb` 文件中。
 
 ### 6.1.5 使用 DeepSpeed 实现分布式训练
 
@@ -289,7 +291,7 @@ from transformers import (
 import datetime
 from transformers.testing_utils import CaptureLogger
 from transformers.trainer_utils import get_last_checkpoint
-import wandb
+import swanlab
 ```
 
 首先需要定义几个超参的类型，用于处理 sh 脚本中设定的超参值。由于 transformers 本身有 TraingingArguments 类，其中包括了训练的一些必备超参数。我们这里只需定义 TrainingArguments 中未包含的超参即可，主要包括模型相关的超参（定义在 ModelArguments）和数据相关的超参（定义在 DataTrainingArguments）：
@@ -456,14 +458,14 @@ trainer.save_model()
 ```
 注意，由于上文检测了是否存在 checkpoint，此处使用 resume_from_checkpoint 来实现从 checkpoint 恢复训练的功能。
 
-由于在大规模训练中监测训练进度、loss 下降趋势尤为重要，在脚本中，我们使用了 wandb 作为训练检测的工具。在脚本开始进行了 wandb 的初始化：
+由于在大规模训练中监测训练进度、loss 下降趋势尤为重要，在脚本中，我们使用了 swanlab 作为训练检测的工具。在脚本开始进行了 swanlab 的初始化：
 
 ```python
-# 初始化 WandB
-wandb.init(project="pretrain", name="from_scrach")
+# 初始化 SwanLab
+swanlab.init(project="pretrain", experiment_name="from_scrach")
 ```
 
-在启动训练后，终端会输出 wandb 监测的 url，点击即可观察训练进度。此处不再赘述 wandb 的使用细节，欢迎读者查阅相关的资料说明。
+在启动训练后，终端会输出 swanlab 监测的 url，点击即可观察训练进度。此处不再赘述 swanlab 的使用细节，欢迎读者查阅相关的资料说明。
 
 完成上述代码后，我们使用一个 sh 脚本（`./code/pretrain.sh`）定义超参数的值，并通过 Deepspeed 启动训练，从而实现高效的多卡分布式训练：
 
@@ -495,7 +497,7 @@ deepspeed pretrain.py \
     --bf16 \
     --gradient_checkpointing \
     --deepspeed ./ds_config_zero2.json \
-    --report_to wandb
+    --report_to swanlab
     # --resume_from_checkpoint ${output_model}/checkpoint-20400 \
 ```
 在安装了 Deepspeed 第三方库后，可以直接通过 Deepspeed 命令来启动多卡训练。上述脚本命令主要是定义了各种超参数的值，可参考使用。在第四章中，我们介绍了 DeepSpeed 分布式训练的原理和 ZeRO 阶段设置，在这里，我们使用 ZeRO-2 进行训练。此处加载了 `ds_config_zero.json` 作为 DeepSpeed 的配置参数：
@@ -690,8 +692,8 @@ class SupervisedDataset(Dataset):
 parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
 model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
-# 初始化 WandB
-wandb.init(project="sft", name="qwen-1.5b")
+# 初始化 SwanLab
+swanlab.init(project="sft", experiment_name="qwen-1.5b")
 
 # 设置日志
 logging.basicConfig(
@@ -741,7 +743,7 @@ logger.info(f"继承一个预训练模型 - Total size={n_params/2**20:.2f}M par
 
 # 初始化 Tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
-logger.info("完成 tokenzier 加载")
+logger.info("完成 tokenizer 加载")
 
 # 加载微调数据
 with open(data_args.train_files) as f:
@@ -788,7 +790,7 @@ trainer.save_model()
 具体而言，其在预训练模型每层中插入用于下游任务的参数，即 Adapter 模块，在微调时冻结模型主体，仅训练特定于任务的参数，如图6.8所示。
 
 <div align='center'>
-    <img src="../images/6-images/3-1.png" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/6-images/3-1.png" alt="alt text" width="90%">
     <p>图6.8 Adapt Tuning</p>
 </div>
 
@@ -840,7 +842,7 @@ $$h = W_0 x + \Delta W x = W_0 x + B A x$$
 训练思路如图6.9所示：
 
 <div align='center'>
-    <img src="../images/6-images/3-2.jpg" alt="alt text" width="90%">
+    <img src="https://raw.githubusercontent.com/datawhalechina/happy-llm/main/docs/images/6-images/3-2.jpg" alt="alt text" width="90%">
     <p>图6.9 LoRA</p>
 </div>
 

docs/chapter7/Agent/demo.py

@@ -1,5 +1,5 @@
 from src.core import Agent
-from src.tools import add, count_letter_in_string, compare, get_current_datetime
+from src.tools import add, count_letter_in_string, compare, get_current_datetime, search_wikipedia, get_current_temperature
 
 from openai import OpenAI
 
@@ -13,7 +13,7 @@ if __name__ == "__main__":
     agent = Agent(
         client=client,
         model="Qwen/Qwen2.5-32B-Instruct",
-        tools=[get_current_datetime, add, compare, count_letter_in_string],
+        tools=[get_current_datetime, search_wikipedia, get_current_temperature],
     )
 
     while True:

docs/chapter7/Agent/requirements.txt

@@ -1,4 +1,55 @@
-json
-openai
-datetime
-pprint
+altair==5.5.0
+annotated-types==0.7.0
+anyio==4.9.0
+attrs==25.3.0
+beautifulsoup4==4.13.4
+blinker==1.9.0
+cachetools==6.1.0
+certifi==2025.6.15
+charset-normalizer==3.4.2
+click==8.2.1
+datetime==5.5
+distro==1.9.0
+gitdb==4.0.12
+gitpython==3.1.44
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+idna==3.10
+jinja2==3.1.6
+jiter==0.10.0
+jsonschema==4.24.0
+jsonschema-specifications==2025.4.1
+markupsafe==3.0.2
+narwhals==1.43.1
+numpy==2.3.0
+openai==1.88.0
+packaging==25.0
+pandas==2.3.0
+pillow==11.2.1
+protobuf==6.31.1
+pyarrow==20.0.0
+pydantic==2.11.7
+pydantic-core==2.33.2
+pydeck==0.9.1
+python-dateutil==2.9.0.post0
+pytz==2025.2
+referencing==0.36.2
+requests==2.32.4
+rpds-py==0.25.1
+setuptools==80.9.0
+six==1.17.0
+smmap==5.0.2
+sniffio==1.3.1
+soupsieve==2.7
+streamlit==1.46.0
+tenacity==9.1.2
+toml==0.10.2
+tornado==6.5.1
+tqdm==4.67.1
+typing-extensions==4.14.0
+typing-inspection==0.4.1
+tzdata==2025.2
+urllib3==2.5.0
+wikipedia==1.4.0
+zope-interface==7.2

docs/chapter7/Agent/src/core.py

@@ -2,11 +2,11 @@ from openai import OpenAI
 import json
 from typing import List, Dict, Any
 from src.utils import function_to_json
-from src.tools import get_current_datetime, add, compare, count_letter_in_string
+from src.tools import get_current_datetime, add, compare, count_letter_in_string, search_wikipedia, get_current_temperature
 
 import pprint
 
-SYSREM_PROMPT = """
+SYSTEM_PROMPT = """
 你是一个叫不要葱姜蒜的人工智能助手。你的输出应该与用户的语言保持一致。
 当用户的问题需要调用工具时，你可以从提供的工具列表中调用适当的工具函数。
 """
@@ -17,7 +17,7 @@ class Agent:
         self.tools = tools
         self.model = model
         self.messages = [
-            {"role": "system", "content": SYSREM_PROMPT},
+            {"role": "system", "content": SYSTEM_PROMPT},
         ]
         self.verbose = verbose
 
@@ -51,7 +51,24 @@ class Agent:
             stream=False,
         )
         if response.choices[0].message.tool_calls:
-            self.messages.append({"role": "assistant", "content": response.choices[0].message.content})
+            # 将包含 tool_calls 的完整 assistant 消息添加到历史中
+            assistant_message = {
+                "role": "assistant",
+                "content": response.choices[0].message.content,
+                "tool_calls": [
+                    {
+                        "id": tool_call.id,
+                        "type": "function",
+                        "function": {
+                            "name": tool_call.function.name,
+                            "arguments": tool_call.function.arguments
+                        }
+                    }
+                    for tool_call in response.choices[0].message.tool_calls
+                ]
+            }
+            self.messages.append(assistant_message)
+
             # 处理工具调用
             tool_list = []
             for tool_call in response.choices[0].message.tool_calls:

docs/chapter7/Agent/src/tools.py

@@ -1,12 +1,14 @@
-from datetime import datetime
+import datetime
+import wikipedia
+import requests
 
 # 获取当前日期和时间
 def get_current_datetime() -> str:
     """
-    获取当前日期和时间。
+    获取真实的当前日期和时间。
     :return: 当前日期和时间的字符串表示。
     """
-    current_datetime = datetime.now()
+    current_datetime = datetime.datetime.now()
     formatted_datetime = current_datetime.strftime("%Y-%m-%d %H:%M:%S")
     return formatted_datetime
 
@@ -54,3 +56,76 @@ def count_letter_in_string(a: str, b: str):
     
     count = string.count(letter)
     return(f"The letter '{letter}' appears {count} times in the string.")
+
+def search_wikipedia(query: str) -> str:
+    """
+    在维基百科中搜索指定查询的前三个页面摘要。
+    :param query: 要搜索的查询字符串。
+    :return: 包含前三个页面摘要的字符串。
+    """
+    page_titles = wikipedia.search(query)
+    summaries = []
+    for page_title in page_titles[: 3]:  # 取前三个页面标题
+        try:
+            # 使用 wikipedia 模块的 page 函数，获取指定标题的维基百科页面对象。
+            wiki_page = wikipedia.page(title=page_title, auto_suggest=False)
+            # 获取页面摘要
+            summaries.append(f"页面: {page_title}\n摘要: {wiki_page.summary}")
+        except (
+                wikipedia.exceptions.PageError,
+                wikipedia.exceptions.DisambiguationError,
+        ):
+            pass
+    if not summaries:
+        return "维基百科没有搜索到合适的结果"
+    return "\n\n".join(summaries)
+
+
+def get_current_temperature(latitude: float, longitude: float) -> str:
+    """
+    获取指定经纬度位置的当前温度。
+    :param latitude: 纬度坐标。
+    :param longitude: 经度坐标。
+    :return: 当前温度的字符串表示。
+    """
+
+    # Open Meteo API 的URL
+    open_meteo_url = "https://api.open-meteo.com/v1/forecast"
+
+    # 请求参数
+    params = {
+        'latitude': latitude,
+        'longitude': longitude,
+        'hourly': 'temperature_2m',
+        'forecast_days': 1,
+    }
+
+    # 发送 API 请求
+    response = requests.get(open_meteo_url, params=params)
+
+    # 检查响应状态码
+    if response.status_code == 200:
+        # 解析 JSON 响应
+        results = response.json()
+    else:
+        # 处理请求失败的情况
+        raise Exception(f"API Request failed with status code: {response.status_code}")
+
+    # 获取当前 UTC 时间
+    current_utc_time = datetime.datetime.now(datetime.UTC)
+
+    # 将时间字符串转换为 datetime 对象
+    time_list = [datetime.datetime.fromisoformat(time_str).replace(tzinfo=datetime.timezone.utc) for time_str in
+                 results['hourly']['time']]
+
+    # 获取温度列表
+    temperature_list = results['hourly']['temperature_2m']
+
+    # 找到最接近当前时间的索引
+    closest_time_index = min(range(len(time_list)), key=lambda i: abs(time_list[i] - current_utc_time))
+
+    # 获取当前温度
+    current_temperature = temperature_list[closest_time_index]
+
+    # 返回当前温度的字符串形式
+    return f'现在温度是 {current_temperature}°C'

docs/chapter7/Agent/web_demo.py

@@ -0,0 +1,62 @@
+import streamlit as st
+from src.core import Agent
+from src.tools import add, count_letter_in_string, compare, get_current_datetime, search_wikipedia, get_current_temperature
+from openai import OpenAI
+
+# --- 页面配置 ---
+st.set_page_config(
+    page_title="Tiny Agent Demo",  # 页面标题
+    page_icon="🤖",  # 页面图标
+    layout="centered",  # 页面布局
+    initial_sidebar_state="auto",  # 侧边栏初始状态
+)
+
+# --- OpenAI客户端初始化 ---
+client = OpenAI(
+    api_key="your siliconflow api key",
+    base_url="https://api.siliconflow.cn/v1",  
+)
+
+# --- Agent初始化 ---
+@st.cache_resource
+def load_agent():
+    """创建并缓存Agent实例。"""
+    return Agent(
+        client=client,
+        model="Qwen/Qwen2.5-32B-Instruct",  # 使用的模型
+        tools=[get_current_datetime, search_wikipedia, get_current_temperature],  # Agent可以使用的工具
+    )
+
+agent = load_agent()  # 加载Agent
+
+# --- UI组件 ---
+st.title("🤖 Happy-LLM Tiny Agent")  # 设置页面标题
+st.markdown("""欢迎来到 Tiny Agent web 界面！
+
+在下方输入您的提示，查看 Agent 的实际操作。
+""")  # 显示Markdown格式的欢迎信息
+
+# 初始化聊天记录
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# 在应用重新运行时显示历史聊天记录
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# 响应用户输入
+if prompt := st.chat_input("我能为您做些什么？"):
+    # 在聊天消息容器中显示用户消息
+    st.chat_message("user").markdown(prompt)
+    # 将用户消息添加到聊天记录中
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    with st.spinner('思考中...'):
+        response = agent.get_completion(prompt)  # 获取Agent的响应
+        
+    # 在聊天消息容器中显示助手响应
+    with st.chat_message("assistant"):
+        st.markdown(response)
+    # 将助手响应添加到聊天记录中
+    st.session_state.messages.append({"role": "assistant", "content": response})

docs/chapter7/RAG/.env_example

@@ -0,0 +1,4 @@
+# 此处默认使用国内可访问的硅基流动平台 https://cloud.siliconflow.cn/
+
+OPENAI_API_KEY='your api key'
+OPENAI_BASE_URL='https://api.siliconflow.cn/v1' 

docs/chapter7/RAG/Embeddings.py

@@ -1,10 +1,10 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 '''
-@File    :   Embeddings.py
-@Time    :   2024/02/10 21:55:39
+@File    :   Embedding.py
+@Time    :   2025/06/20 13:50:47
 @Author  :   不要葱姜蒜
-@Version :   1.0
+@Version :   1.1
 @Desc    :   None
 '''
 
@@ -12,8 +12,8 @@ import os
 from copy import copy
 from typing import Dict, List, Optional, Tuple, Union
 import numpy as np
+from openai import OpenAI
 
-os.environ['CURL_CA_BUNDLE'] = ''
 from dotenv import load_dotenv, find_dotenv
 _ = load_dotenv(find_dotenv())
 
@@ -23,21 +23,59 @@ class BaseEmbeddings:
     Base class for embeddings
     """
     def __init__(self, path: str, is_api: bool) -> None:
+        """
+        初始化嵌入基类
+        Args:
+            path (str): 模型或数据的路径
+            is_api (bool): 是否使用API方式。True表示使用在线API服务，False表示使用本地模型
+        """
         self.path = path
         self.is_api = is_api
     
     def get_embedding(self, text: str, model: str) -> List[float]:
+        """
+        获取文本的嵌入向量表示
+        Args:
+            text (str): 输入文本
+            model (str): 使用的模型名称
+        Returns:
+            List[float]: 文本的嵌入向量
+        Raises:
+            NotImplementedError: 该方法需要在子类中实现
+        """
         raise NotImplementedError
     
     @classmethod
     def cosine_similarity(cls, vector1: List[float], vector2: List[float]) -> float:
         """
-        calculate cosine similarity between two vectors
+        计算两个向量之间的余弦相似度
+        Args:
+            vector1 (List[float]): 第一个向量
+            vector2 (List[float]): 第二个向量
+        Returns:
+            float: 两个向量的余弦相似度，范围在[-1,1]之间
         """
-        dot_product = np.dot(vector1, vector2)
-        magnitude = np.linalg.norm(vector1) * np.linalg.norm(vector2)
-        if not magnitude:
-            return 0
+        # 将输入列表转换为numpy数组，并指定数据类型为float32
+        v1 = np.array(vector1, dtype=np.float32)
+        v2 = np.array(vector2, dtype=np.float32)
+
+        # 检查向量中是否包含无穷大或NaN值
+        if not np.all(np.isfinite(v1)) or not np.all(np.isfinite(v2)):
+            return 0.0
+
+        # 计算向量的点积
+        dot_product = np.dot(v1, v2)
+        # 计算向量的范数（长度）
+        norm_v1 = np.linalg.norm(v1)
+        norm_v2 = np.linalg.norm(v2)
+        
+        # 计算分母（两个向量范数的乘积）
+        magnitude = norm_v1 * norm_v2
+        # 处理分母为0的特殊情况
+        if magnitude == 0:
+            return 0.0
+            
+        # 返回余弦相似度
         return dot_product / magnitude
     
 
@@ -48,70 +86,18 @@ class OpenAIEmbedding(BaseEmbeddings):
     def __init__(self, path: str = '', is_api: bool = True) -> None:
         super().__init__(path, is_api)
         if self.is_api:
-            from openai import OpenAI
             self.client = OpenAI()
+            # 从环境变量中获取 硅基流动 密钥
             self.client.api_key = os.getenv("OPENAI_API_KEY")
+            # 从环境变量中获取 硅基流动 的基础URL
             self.client.base_url = os.getenv("OPENAI_BASE_URL")
     
-    def get_embedding(self, text: str, model: str = "text-embedding-3-large") -> List[float]:
+    def get_embedding(self, text: str, model: str = "BAAI/bge-m3") -> List[float]:
+        """
+        此处默认使用轨迹流动的免费嵌入模型 BAAI/bge-m3
+        """
         if self.is_api:
             text = text.replace("\n", " ")
             return self.client.embeddings.create(input=[text], model=model).data[0].embedding
         else:
             raise NotImplementedError
-
-class JinaEmbedding(BaseEmbeddings):
-    """
-    class for Jina embeddings
-    """
-    def __init__(self, path: str = 'jinaai/jina-embeddings-v2-base-zh', is_api: bool = False) -> None:
-        super().__init__(path, is_api)
-        self._model = self.load_model()
-        
-    def get_embedding(self, text: str) -> List[float]:
-        return self._model.encode([text])[0].tolist()
-    
-    def load_model(self):
-        import torch
-        from transformers import AutoModel
-        if torch.cuda.is_available():
-            device = torch.device("cuda")
-        else:
-            device = torch.device("cpu")
-        model = AutoModel.from_pretrained(self.path, trust_remote_code=True).to(device)
-        return model
-
-class ZhipuEmbedding(BaseEmbeddings):
-    """
-    class for Zhipu embeddings
-    """
-    def __init__(self, path: str = '', is_api: bool = True) -> None:
-        super().__init__(path, is_api)
-        if self.is_api:
-            from zhipuai import ZhipuAI
-            self.client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY")) 
-    
-    def get_embedding(self, text: str) -> List[float]:
-        response = self.client.embeddings.create(
-        model="embedding-2",
-        input=text,
-        )
-        return response.data[0].embedding
-
-class DashscopeEmbedding(BaseEmbeddings):
-    """
-    class for Dashscope embeddings
-    """
-    def __init__(self, path: str = '', is_api: bool = True) -> None:
-        super().__init__(path, is_api)
-        if self.is_api:
-            import dashscope
-            dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")
-            self.client = dashscope.TextEmbedding
-
-    def get_embedding(self, text: str, model: str='text-embedding-v1') -> List[float]:
-        response = self.client.call(
-            model=model,
-            input=text
-        )
-        return response.output['embeddings'][0]['embedding']

docs/chapter7/RAG/LLM.py

@@ -2,37 +2,33 @@
 # -*- coding: utf-8 -*-
 '''
 @File    :   LLM.py
-@Time    :   2024/02/12 13:50:47
+@Time    :   2025/06/20 13:50:47
 @Author  :   不要葱姜蒜
-@Version :   1.0
+@Version :   1.1
 @Desc    :   None
 '''
 import os
 from typing import Dict, List, Optional, Tuple, Union
+from openai import OpenAI
 
-PROMPT_TEMPLATE = dict(
-    RAG_PROMPT_TEMPALTE="""使用以上下文来回答用户的问题。如果你不知道答案，就说你不知道。总是使用中文回答。
-        问题: {question}
-        可参考的上下文：
-        ···
-        {context}
-        ···
-        如果给定的上下文无法让你做出回答，请回答数据库中没有这个内容，你不知道。
-        有用的回答:""",
-    InternLM_PROMPT_TEMPALTE="""先对上下文进行内容总结,再使用上下文来回答用户的问题。如果你不知道答案，就说你不知道。总是使用中文回答。
-        问题: {question}
-        可参考的上下文：
-        ···
-        {context}
-        ···
-        如果给定的上下文无法让你做出回答，请回答数据库中没有这个内容，你不知道。
-        有用的回答:"""
-)
+from dotenv import load_dotenv, find_dotenv
+_ = load_dotenv(find_dotenv())
+
+RAG_PROMPT_TEMPLATE="""
+使用以上下文来回答用户的问题。如果你不知道答案，就说你不知道。总是使用中文回答。
+问题: {question}
+可参考的上下文：
+···
+{context}
+···
+如果给定的上下文无法让你做出回答，请回答数据库中没有这个内容，你不知道。
+有用的回答:
+"""
 
 
 class BaseModel:
-    def __init__(self, path: str = '') -> None:
-        self.path = path
+    def __init__(self, model) -> None:
+        self.model = model
 
     def chat(self, prompt: str, history: List[dict], content: str) -> str:
         pass
@@ -41,73 +37,18 @@ class BaseModel:
         pass
 
 class OpenAIChat(BaseModel):
-    def __init__(self, path: str = '', model: str = "gpt-3.5-turbo-1106") -> None:
-        super().__init__(path)
+    def __init__(self, model: str = "Qwen/Qwen2.5-32B-Instruct") -> None:
         self.model = model
 
     def chat(self, prompt: str, history: List[dict], content: str) -> str:
-        from openai import OpenAI
         client = OpenAI()
         client.api_key = os.getenv("OPENAI_API_KEY")   
         client.base_url = os.getenv("OPENAI_BASE_URL")
-        history.append({'role': 'user', 'content': PROMPT_TEMPLATE['RAG_PROMPT_TEMPALTE'].format(question=prompt, context=content)})
+        history.append({'role': 'user', 'content': RAG_PROMPT_TEMPLATE.format(question=prompt, context=content)})
         response = client.chat.completions.create(
             model=self.model,
             messages=history,
-            max_tokens=150,
+            max_tokens=2048,
             temperature=0.1
         )
         return response.choices[0].message.content
-
-class InternLMChat(BaseModel):
-    def __init__(self, path: str = '') -> None:
-        super().__init__(path)
-        self.load_model()
-
-    def chat(self, prompt: str, history: List = [], content: str='') -> str:
-        prompt = PROMPT_TEMPLATE['InternLM_PROMPT_TEMPALTE'].format(question=prompt, context=content)
-        response, history = self.model.chat(self.tokenizer, prompt, history)
-        return response
-
-
-    def load_model(self):
-        import torch
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-        self.tokenizer = AutoTokenizer.from_pretrained(self.path, trust_remote_code=True)
-        self.model = AutoModelForCausalLM.from_pretrained(self.path, torch_dtype=torch.float16, trust_remote_code=True).cuda()
-
-class DashscopeChat(BaseModel):
-    def __init__(self, path: str = '', model: str = "qwen-turbo") -> None:
-        super().__init__(path)
-        self.model = model
-
-    def chat(self, prompt: str, history: List[Dict], content: str) -> str:
-        import dashscope
-        dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")
-        history.append({'role': 'user', 'content': PROMPT_TEMPLATE['RAG_PROMPT_TEMPALTE'].format(question=prompt, context=content)})
-        response = dashscope.Generation.call(
-            model=self.model,
-            messages=history,
-            result_format='message',
-            max_tokens=150,
-            temperature=0.1
-        )
-        return response.output.choices[0].message.content
-    
-
-class ZhipuChat(BaseModel):
-    def __init__(self, path: str = '', model: str = "glm-4") -> None:
-        super().__init__(path)
-        from zhipuai import ZhipuAI
-        self.client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))
-        self.model = model
-
-    def chat(self, prompt: str, history: List[Dict], content: str) -> str:
-        history.append({'role': 'user', 'content': PROMPT_TEMPLATE['RAG_PROMPT_TEMPALTE'].format(question=prompt, context=content)})
-        response = self.client.chat.completions.create(
-            model=self.model,
-            messages=history,
-            max_tokens=150,
-            temperature=0.1
-        )
-        return response.choices[0].message

docs/chapter7/RAG/VectorBase.py

@@ -2,16 +2,16 @@
 # -*- coding: utf-8 -*-
 '''
 @File    :   VectorBase.py
-@Time    :   2024/02/12 10:11:13
+@Time    :   2025/06/20 10:11:13
 @Author  :   不要葱姜蒜
-@Version :   1.0
+@Version :   1.1
 @Desc    :   None
 '''
 
 import os
 from typing import Dict, List, Optional, Tuple, Union
 import json
-from RAG.Embeddings import BaseEmbeddings, OpenAIEmbedding, JinaEmbedding, ZhipuEmbedding
+from Embeddings import BaseEmbeddings, OpenAIEmbedding
 import numpy as np
 from tqdm import tqdm
 

docs/chapter7/RAG/demo.py

@@ -0,0 +1,19 @@
+from VectorBase import VectorStore
+from utils import ReadFiles
+from LLM import OpenAIChat
+from Embeddings import OpenAIEmbedding
+
+# 没有保存数据库
+docs = ReadFiles('./data').get_content(max_token_len=600, cover_content=150) # 获得data目录下的所有文件内容并分割
+vector = VectorStore(docs)
+embedding = OpenAIEmbedding() # 创建EmbeddingModel
+vector.get_vector(EmbeddingModel=embedding)
+vector.persist(path='storage') # 将向量和文档内容保存到storage目录下，下次再用就可以直接加载本地的数据库
+
+# vector.load_vector('./storage') # 加载本地的数据库
+
+question = 'RAG的原理是什么？'
+
+content = vector.query(question, EmbeddingModel=embedding, k=1)[0]
+chat = OpenAIChat(model='Qwen/Qwen2.5-32B-Instruct')
+print(chat.chat(question, [], content))

docs/chapter7/RAG/requirements.txt

@@ -0,0 +1,28 @@
+annotated-types==0.7.0
+anyio==4.9.0
+beautifulsoup4==4.13.4
+bs4==0.0.2
+certifi==2025.6.15
+charset-normalizer==3.4.2
+distro==1.9.0
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+idna==3.10
+jiter==0.10.0
+markdown==3.8.2
+numpy==2.3.0
+openai==1.88.0
+pydantic==2.11.7
+pydantic-core==2.33.2
+pypdf2==3.0.1
+python-dotenv==1.1.0
+regex==2024.11.6
+requests==2.32.4
+sniffio==1.3.1
+soupsieve==2.7
+tiktoken==0.9.0
+tqdm==4.67.1
+typing-extensions==4.14.0
+typing-inspection==0.4.1
+urllib3==2.5.0

docs/chapter7/RAG/utils.py

@@ -2,9 +2,9 @@
 # -*- coding: utf-8 -*-
 '''
 @File    :   utils.py
-@Time    :   2024/02/11 09:52:26
+@Time    :   2025/06/20 13:50:47
 @Author  :   不要葱姜蒜
-@Version :   1.0
+@Version :   1.1
 @Desc    :   None
 '''
 
@@ -13,7 +13,6 @@ from typing import Dict, List, Optional, Tuple, Union
 
 import PyPDF2
 import markdown
-import html2text
 import json
 from tqdm import tqdm
 import tiktoken
@@ -69,37 +68,65 @@ class ReadFiles:
         lines = text.splitlines()  # 假设以换行符分割文本为行
 
         for line in lines:
-            line = line.replace(' ', '')
+            # 保留空格，只移除行首行尾空格
+            line = line.strip()
             line_len = len(enc.encode(line))
+            
             if line_len > max_token_len:
                 # 如果单行长度就超过限制，则将其分割成多个块
-                num_chunks = (line_len + token_len - 1) // token_len
-                for i in range(num_chunks):
-                    start = i * token_len
-                    end = start + token_len
-                    # 避免跨单词分割
-                    while not line[start:end].rstrip().isspace():
-                        start += 1
-                        end += 1
-                        if start >= line_len:
-                            break
-                    curr_chunk = curr_chunk[-cover_content:] + line[start:end]
+                # 先保存当前块（如果有内容）
+                if curr_chunk:
                     chunk_text.append(curr_chunk)
-                # 处理最后一个块
-                start = (num_chunks - 1) * token_len
-                curr_chunk = curr_chunk[-cover_content:] + line[start:end]
-                chunk_text.append(curr_chunk)
+                    curr_chunk = ''
+                    curr_len = 0
                 
-            if curr_len + line_len <= token_len:
+                # 将长行按token长度分割
+                line_tokens = enc.encode(line)
+                num_chunks = (len(line_tokens) + token_len - 1) // token_len
+                
+                for i in range(num_chunks):
+                    start_token = i * token_len
+                    end_token = min(start_token + token_len, len(line_tokens))
+                    
+                    # 解码token片段回文本
+                    chunk_tokens = line_tokens[start_token:end_token]
+                    chunk_part = enc.decode(chunk_tokens)
+                    
+                    # 添加覆盖内容（除了第一个块）
+                    if i > 0 and chunk_text:
+                        prev_chunk = chunk_text[-1]
+                        cover_part = prev_chunk[-cover_content:] if len(prev_chunk) > cover_content else prev_chunk
+                        chunk_part = cover_part + chunk_part
+                    
+                    chunk_text.append(chunk_part)
+                
+                # 重置当前块状态
+                curr_chunk = ''
+                curr_len = 0
+                
+            elif curr_len + line_len + 1 <= token_len:  # +1 for newline
+                # 当前行可以加入当前块
+                if curr_chunk:
+                    curr_chunk += '\n'
+                    curr_len += 1
                 curr_chunk += line
-                curr_chunk += '\n'
                 curr_len += line_len
-                curr_len += 1
             else:
-                chunk_text.append(curr_chunk)
-                curr_chunk = curr_chunk[-cover_content:]+line
-                curr_len = line_len + cover_content
+                # 当前行无法加入当前块，开始新块
+                if curr_chunk:
+                    chunk_text.append(curr_chunk)
+                
+                # 开始新块，添加覆盖内容
+                if chunk_text:
+                    prev_chunk = chunk_text[-1]
+                    cover_part = prev_chunk[-cover_content:] if len(prev_chunk) > cover_content else prev_chunk
+                    curr_chunk = cover_part + '\n' + line
+                    curr_len = len(enc.encode(cover_part)) + 1 + line_len
+                else:
+                    curr_chunk = line
+                    curr_len = line_len
 
+        # 添加最后一个块（如果有内容）
         if curr_chunk:
             chunk_text.append(curr_chunk)