你是否好奇过 GitHub Copilot、CodeLlama 这些代码生成模型是如何工作的？本文将带你从零开始，一步步实现一个专注于 Python 代码生成的小型语言模型。通过这个项目，你将深入理解 Transformer 架构、代码 tokenization、以及如何让模型学会”写代码”。

为什么要自己实现一个代码模型？

市面上已经有很多优秀的代码生成模型，但自己动手实现一个有几个独特的价值：

  1. 深入理解原理：纸上得来终觉浅，只有亲手实现才能真正理解每个组件的作用
  2. 定制化需求：你可以针对特定的代码风格或领域进行优化
  3. 资源可控：小模型可以在消费级 GPU 上训练和运行
  4. 学习路径：这是进入 AI 领域的绝佳实践项目

我们的目标是训练一个约 50M 参数的模型，能够：

  • 根据函数签名和注释生成 Python 函数体
  • 补全未完成的代码片段
  • 理解基本的 Python 语法和常用库

整体架构概览

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
┌─────────────────────────────────────────────────────────────┐
│                    Code Generation LLM                       │
├─────────────────────────────────────────────────────────────┤
│  1. 数据收集与预处理                                          │
│     └── Python 代码语料库 → 清洗 → Tokenization              │
├─────────────────────────────────────────────────────────────┤
│  2. 模型架构                                                  │
│     └── Decoder-only Transformer (GPT-style)                 │
├─────────────────────────────────────────────────────────────┤
│  3. 训练流程                                                  │
│     └── Next Token Prediction + Causal Language Modeling     │
├─────────────────────────────────────────────────────────────┤
│  4. 推理与代码生成                                            │
│     └── Temperature Sampling + Top-k/Top-p                   │
└─────────────────────────────────────────────────────────────┘

Step 1: 数据收集与预处理

1.1 收集 Python 代码数据

高质量的训练数据是模型成功的基础。我们可以从以下来源获取 Python 代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
import os
import ast
from pathlib import Path
from typing import List, Optional
from dataclasses import dataclass

@dataclass
class CodeSample:
    """表示一个代码样本"""
    source: str           # 原始代码
    file_path: str        # 文件路径
    is_valid: bool        # 是否语法正确
    functions: List[str]  # 提取的函数列表

def collect_python_files(root_dir: str) -> List[Path]:
    """

    Args:

    Returns:
        Python 文件路径列表
    """
    python_files = []
    for path in Path(root_dir).rglob("*.py"):
        # 跳过测试文件和虚拟环境
        if "test" not in str(path).lower() and "venv" not in str(path):
            python_files.append(path)
    return python_files

def validate_python_syntax(code: str) -> bool:
    """
    检查代码是否是有效的 Python 语法

    Args:
        code: Python 代码字符串

    Returns:
        语法是否有效
    """
    try:
        ast.parse(code)
        return True
    except SyntaxError:
        return False

def extract_functions(code: str) -> List[str]:
    """
    从代码中提取所有函数定义

    Args:
        code: Python 代码字符串

    Returns:
        函数代码列表
    """
    functions = []
    try:
        tree = ast.parse(code)
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                # 获取函数的源代码
                func_source = ast.get_source_segment(code, node)
                if func_source:
                    functions.append(func_source)
    except SyntaxError:
        pass
    return functions

def process_code_file(file_path: Path) -> Optional[CodeSample]:
    """
    处理单个代码文件

    Args:
        file_path: 文件路径

    Returns:
        CodeSample 对象或 None
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            code = f.read()

        # 跳过过短或过长的文件
        if len(code) < 100 or len(code) > 100000:
            return None

        is_valid = validate_python_syntax(code)
        functions = extract_functions(code) if is_valid else []

        return CodeSample(
            source=code,
            file_path=str(file_path),
            is_valid=is_valid,
            functions=functions
        )
    except Exception:
        return None

# generated by hugo's coding agent

1.2 代码 Tokenizer

代码的 tokenization 与自然语言有所不同。我们需要保留缩进、特殊符号等对代码语义至关重要的信息。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
import re
from collections import Counter
from typing import Dict, List, Tuple

class CodeTokenizer:
    """
    专为 Python 代码设计的 Tokenizer

    使用 Byte-Pair Encoding (BPE) 算法，但对代码特殊处理：
    - 保留缩进信息
    - 识别关键字和运算符
    - 处理字符串和注释
    """

    # Python 关键字作为特殊 token
    PYTHON_KEYWORDS = [
        'def', 'class', 'if', 'else', 'elif', 'for', 'while', 'try',
        'except', 'finally', 'with', 'as', 'import', 'from', 'return',
        'yield', 'raise', 'pass', 'break', 'continue', 'lambda', 'and',
        'or', 'not', 'in', 'is', 'None', 'True', 'False', 'async', 'await'
    ]

    # 特殊 token
    SPECIAL_TOKENS = {
        '<PAD>': 0,
        '<UNK>': 1,
        '<BOS>': 2,  # Beginning of sequence
        '<EOS>': 3,  # End of sequence
        '<INDENT>': 4,  # 缩进增加
        '<DEDENT>': 5,  # 缩进减少
        '<NEWLINE>': 6,
    }

    def __init__(self, vocab_size: int = 8000):
        self.vocab_size = vocab_size
        self.token_to_id: Dict[str, int] = {}
        self.id_to_token: Dict[int, str] = {}
        self.merges: Dict[Tuple[str, str], str] = {}

    def _pre_tokenize(self, code: str) -> List[str]:
        """
        预分词：将代码分割成基本单元

        处理策略：
        1. 保留完整的字符串
        2. 分离运算符和标点
        3. 保留空白用于缩进处理
        """
        tokens = []

        # 正则表达式匹配不同类型的 token
        pattern = r'''
            ("[^"]*"|'[^']*')  |  # 字符串
            (\#[^\n]*)         |  # 注释
            (\d+\.?\d*)        |  # 数字
            ([a-zA-Z_]\w*)     |  # 标识符
            ([ ]{4}|\t)        |  # 缩进单位
            (\n)               |  # 换行
            ([^\s\w])             # 其他符号
        '''

        for match in re.finditer(pattern, code, re.VERBOSE):
            token = match.group()
            if token.strip() or token in ('\n', '    ', '\t'):
                tokens.append(token)

        return tokens

    def _process_indentation(self, tokens: List[str]) -> List[str]:
        """
        处理缩进，转换为 INDENT/DEDENT token
        """
        processed = []
        indent_stack = [0]
        current_indent = 0
        at_line_start = True

        for token in tokens:
            if token == '\n':
                processed.append('<NEWLINE>')
                at_line_start = True
                current_indent = 0
            elif at_line_start and token in ('    ', '\t'):
                current_indent += 1
            elif at_line_start:
                # 处理缩进变化
                while current_indent < indent_stack[-1]:
                    processed.append('<DEDENT>')
                    indent_stack.pop()
                if current_indent > indent_stack[-1]:
                    processed.append('<INDENT>')
                    indent_stack.append(current_indent)
                at_line_start = False
                processed.append(token)
            else:
                processed.append(token)

        return processed

    def train(self, code_samples: List[str], num_merges: int = 5000):
        """
        训练 BPE tokenizer

        Args:
            code_samples: 代码样本列表
            num_merges: BPE 合并次数
        """
        # 初始化词表
        self.token_to_id = dict(self.SPECIAL_TOKENS)
        next_id = len(self.SPECIAL_TOKENS)

        # 添加 Python 关键字
        for keyword in self.PYTHON_KEYWORDS:
            self.token_to_id[keyword] = next_id
            next_id += 1

        # 统计所有字符
        all_tokens = []
        for code in code_samples:
            tokens = self._pre_tokenize(code)
            tokens = self._process_indentation(tokens)
            all_tokens.extend(tokens)

        # 将每个 token 拆分成字符
        words = [list(t) + ['</w>'] for t in all_tokens if t not in self.SPECIAL_TOKENS]

        # BPE 训练
        for _ in range(num_merges):
            pairs = Counter()
            for word in words:
                for i in range(len(word) - 1):
                    pairs[(word[i], word[i + 1])] += 1

            if not pairs:
                break

            best_pair = max(pairs, key=pairs.get)
            new_token = ''.join(best_pair)

            if new_token not in self.token_to_id:
                self.token_to_id[new_token] = next_id
                next_id += 1
                self.merges[best_pair] = new_token

            # 应用合并
            new_words = []
            for word in words:
                new_word = []
                i = 0
                while i < len(word):
                    if i < len(word) - 1 and (word[i], word[i + 1]) == best_pair:
                        new_word.append(new_token)
                        i += 2
                    else:
                        new_word.append(word[i])
                        i += 1
                new_words.append(new_word)
            words = new_words

        # 构建反向映射
        self.id_to_token = {v: k for k, v in self.token_to_id.items()}

    def encode(self, code: str) -> List[int]:
        """将代码编码为 token ID 序列"""
        tokens = self._pre_tokenize(code)
        tokens = self._process_indentation(tokens)

        ids = [self.SPECIAL_TOKENS['<BOS>']]

        for token in tokens:
            if token in self.token_to_id:
                ids.append(self.token_to_id[token])
            elif token in self.SPECIAL_TOKENS:
                ids.append(self.SPECIAL_TOKENS[token])
            else:
                # 应用 BPE
                word = list(token) + ['</w>']
                while len(word) > 1:
                    pairs = [(word[i], word[i + 1]) for i in range(len(word) - 1)]
                    mergeable = [p for p in pairs if p in self.merges]
                    if not mergeable:
                        break
                    pair = min(mergeable, key=lambda p: list(self.merges.keys()).index(p))
                    new_word = []
                    i = 0
                    while i < len(word):
                        if i < len(word) - 1 and (word[i], word[i + 1]) == pair:
                            new_word.append(self.merges[pair])
                            i += 2
                        else:
                            new_word.append(word[i])
                            i += 1
                    word = new_word

                for subtoken in word:
                    if subtoken in self.token_to_id:
                        ids.append(self.token_to_id[subtoken])
                    else:
                        ids.append(self.SPECIAL_TOKENS['<UNK>'])

        ids.append(self.SPECIAL_TOKENS['<EOS>'])
        return ids

    def decode(self, ids: List[int]) -> str:
        """将 token ID 序列解码为代码"""
        tokens = []
        for id in ids:
            if id in self.id_to_token:
                token = self.id_to_token[id]
                if token not in ('<PAD>', '<BOS>', '<EOS>'):
                    tokens.append(token)

        # 重建代码
        code = ''
        indent_level = 0

        for token in tokens:
            if token == '<NEWLINE>':
                code += '\n' + '    ' * indent_level
            elif token == '<INDENT>':
                indent_level += 1
                code += '    '
            elif token == '<DEDENT>':
                indent_level = max(0, indent_level - 1)
                code = code.rstrip('    ')
            else:
                token = token.replace('</w>', '')
                code += token

        return code

# generated by hugo's coding agent

Step 2: 构建 Transformer 模型

2.1 模型配置

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
from dataclasses import dataclass

@dataclass
class CodeLLMConfig:
    """模型配置"""
    vocab_size: int = 8000       # 词表大小
    max_seq_len: int = 1024      # 最大序列长度
    d_model: int = 512           # 模型维度
    n_heads: int = 8             # 注意力头数
    n_layers: int = 6            # Transformer 层数
    d_ff: int = 2048             # 前馈网络维度
    dropout: float = 0.1         # Dropout 概率

    @property
    def n_params(self) -> int:
        """估算参数量"""
        # Embedding
        embed_params = self.vocab_size * self.d_model
        # Attention (Q, K, V, O projections per layer)
        attn_params = 4 * self.d_model * self.d_model * self.n_layers
        # FFN (2 linear layers per layer)
        ffn_params = 2 * self.d_model * self.d_ff * self.n_layers
        # Layer norms
        ln_params = 4 * self.d_model * self.n_layers
        # Output projection
        out_params = self.d_model * self.vocab_size

        return embed_params + attn_params + ffn_params + ln_params + out_params

# 50M 参数的配置
config = CodeLLMConfig()
print(f"Estimated parameters: {config.n_params / 1e6:.1f}M")
# Output: Estimated parameters: 51.2M

# generated by hugo's coding agent

2.2 核心组件实现

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional

class RotaryPositionalEmbedding(nn.Module):
    """
    旋转位置编码 (RoPE)

    相比传统的正弦位置编码，RoPE 有更好的长度外推能力，
    且能更好地编码相对位置信息。
    """

    def __init__(self, d_model: int, max_seq_len: int = 2048, base: float = 10000.0):
        super().__init__()
        self.d_model = d_model
        self.max_seq_len = max_seq_len

        # 计算频率
        inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
        self.register_buffer('inv_freq', inv_freq)

        # 预计算 cos 和 sin
        self._build_cache(max_seq_len)

    def _build_cache(self, seq_len: int):
        t = torch.arange(seq_len, device=self.inv_freq.device)
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer('cos_cached', emb.cos())
        self.register_buffer('sin_cached', emb.sin())

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        seq_len = x.shape[1]
        if seq_len > self.max_seq_len:
            self._build_cache(seq_len)

        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]
def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """将张量的后半部分旋转到前面并取负"""
    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
    return torch.cat([-x2, x1], dim=-1)
def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor,
                         cos: torch.Tensor, sin: torch.Tensor) -> tuple:
    """应用旋转位置编码到 Q 和 K"""
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed
class MultiHeadAttention(nn.Module):
    """
    多头自注意力机制

    使用 RoPE 位置编码和 KV Cache 优化推理速度
    """

    def __init__(self, config: CodeLLMConfig):
        super().__init__()
        self.n_heads = config.n_heads
        self.d_model = config.d_model
        self.head_dim = config.d_model // config.n_heads

        assert self.head_dim * self.n_heads == config.d_model

        self.q_pro
---
原文链接: [Step-by-Step 实现一个能编程的大模型](https://hugozhu.site/post/2026/120-build-code-llm-from-scratch/)