随着领域特定语言模型的兴起，数据的重要性日益凸显。对于代码生成模型而言，预训练阶段所使用的源代码质量、多样性和纯净度，直接影响其下游任务的性能表现。本文将探讨如何从零开始，构建一个高质量的代码数据集，并以此训练出一个强大的代码生成大模型。
本文主要参考了一篇英文文章，该文章详细描述了作者 Wasi Ullah 如何利用 GitHub API，构建一个可扩展、去重且定制化的源代码数据集，并基于此预训练了一个基于 GPT 的代码模型。目标是为构建一个开放的代码模型奠定坚实的基础，使其能够实现智能的代码补全和生成。
一、 构建高质量代码数据集：数据是关键
“模型的好坏取决于训练它的数据”。许多开源代码数据集存在冗余、不完整、包含大量二进制文件、格式不佳或存在与现代代码库的领域偏差等问题。因此，构建自定义数据集至关重要，这能确保：
- 更高的质量: 从真实的 GitHub 仓库中获取高质量的代码。
- 更强的控制: 严格的应用筛选和去重标准。
- 更好的可追溯性: 每一步都具有可重复性和可审计性。
作者没有依赖现成的开源数据集，而是选择直接利用 GitHub API，从高质量的 GitHub 仓库中获取数据，并通过严格的数据清洗流程，构建了一个定制化的代码数据集。
二、 GitHub API：获取高质量源代码的入口
GitHub 提供了强大的 REST API，可以访问数百万个公共仓库。为了获取高质量的预训练数据，作者首先使用 GitHub Search API 检索了 star 数最高的 100 个公共仓库。
为什么选择 star 数最高的仓库？因为 star 数是社区信任、代码质量和积极维护的信号。通过以下 Python 代码可以实现这一目标：
import requests
GITHUB_TOKEN = "your_github_token"
HEADERS = {'Authorization': f'token {GITHUB_TOKEN}'}
def get_top_repositories(pages=4):
repos = []
for page in range(1, pages + 1):
url = f'https://api.github.com/search/repositories?q=stars:>1&sort=stars&order=desc&per_page=25&page={page}'
resp = requests.get(url, headers=HEADERS)
repos += resp.json().get("items", [])
return repos
这段代码：
- 使用 GitHub token 进行安全身份验证。
- 使用分页来获取前 100 个仓库。
- 确保获取的是最新、高质量的源代码。
直接使用 API 可以对收集的数据类型和范围进行精细控制，这对于可重复性和领域对齐至关重要。
三、 代码仓库下载：确保数据的完整性
一旦确定了要使用的仓库，就可以使用 GitHub API 从其默认分支下载整个仓库。以下代码展示了如何下载仓库的 zip 文件：
from zipfile import ZipFile
def download_repo_zip(repo):
owner, name, branch = repo['owner']['login'], repo['name'], repo['default_branch']
zip_url = f"https://api.github.com/repos/{owner}/{name}/zipball/{branch}"
response = requests.get(zip_url, headers=HEADERS, stream=True)
zip_path = f"{owner}_{name}.zip"
with open(zip_path, 'wb') as f:
for chunk in response.iter_content(1024):
f.write(chunk)
return zip_path
最佳实践是添加具有指数退避的重试逻辑，以处理网络错误、超时或 GitHub 速率限制。这样可以确保数据的完整性，即使在网络不稳定或达到 API 速率限制的情况下，也能成功下载所有必要的数据。
四、 代码预处理：数据清洗的三大步骤
在将下载的仓库用于预训练之前，需要对其进行处理，主要包括三个关键步骤：文件过滤、Tokenization 和去重。
-
文件过滤： 并非所有文件都有用。为了确保只有紧凑且有意义的代码才能通过，需要进行过滤：
- 忽略大于 1MB 的文件（跳过二进制文件、大型日志、数据转储、压缩后的 JavaScript）。
- 过滤掉少于 100 个 token 的文件（对于建模来说太短）。
以下代码演示了如何过滤文件：
import os def is_valid_file(path): return os.path.getsize(path) < 1 * 1024 * 1024 # max 1MB这种过滤方法确保只处理紧凑且有意义的代码，跳过大型二进制文件和日志，并丢弃对于有效建模来说太短的文件。例如，许多 GitHub 仓库包含大量的
node_modules文件夹，其中包含许多压缩后的 JavaScript 文件和图片等二进制文件，这些文件对于代码生成模型的训练没有帮助，反而会干扰模型的学习，因此需要过滤掉。 -
Tokenization：测量代码密度 使用 Python 的内置
tokenize模块，计算每个文件的词汇 token 数量，以估计代码的丰富程度。import tokenize, io def count_tokens(code): try: return len(list(tokenize.generate_tokens(io.StringIO(code).readline))) except: return 0Token 密度比行数或字符数更能代表语义内容，并且在一定程度上与语言无关。例如，一个包含大量注释或空行的 Python 文件，其代码密度较低，而一个包含大量复杂逻辑和函数调用的 Python 文件，其代码密度较高。代码密度高的文件通常包含更多的信息和知识，更有助于模型学习代码的语法和语义。
-
去重：避免数据泄露和模型污染 训练数据中的重复项可能导致：
- 过拟合
- 不公平的基准测试
- 测试数据中泄露的答案
为了防止这种情况，我们使用 SHA-256 对每个文件的内容进行哈希处理。仅保留唯一的哈希值。
import hashlib hashes = set() def is_duplicate(code): code_hash = hashlib.sha256(code.encode('utf-8')).hexdigest() if code_hash in hashes: return True hashes.add(code_hash) return False这种去重方法符合研究人员遵循的最佳实践，以防止过拟合和评估基准的污染。例如，如果训练集中包含大量的重复代码片段，模型可能会过度拟合这些片段，从而导致在训练集上表现良好，但在测试集上表现较差。此外，如果测试集中包含与训练集中相同的代码片段，模型可能会直接“记住”这些片段的答案，从而导致评估结果不准确。因此，去重是构建高质量数据集的关键步骤。
作者的这个方法受到了 PolyCoder 论文的启发，该论文也采用了类似的方法来构建代码数据集。
五、 数据存储与压缩：为模型训练做好准备
每个有效、token 丰富且唯一的文件都写入一个新的目录，保留带有哈希前缀的文件名元数据。
save_path = os.path.join(processed_dir, f"{code_hash}_{os.path.basename(file)}")
with open(save_path, 'w') as f:
f.write(code)
处理后，将数据集压缩并上传到 Google Drive，以便安全保存并在模型训练中使用。
import shutil
shutil.make_archive('/content/drive/MyDrive/processed_code_dataset', 'zip', processed_dir)
最终的数据集包含来自排名前 100 的 GitHub 仓库的干净、去重、token 丰富的代码，输出路径为：/content/drive/MyDrive/processed_code_dataset.zip。
六、 使用 GPT-2 和自定义 GitHub API 数据集预训练代码语言模型
完成了高质量的数据集构建之后，就可以开始预训练代码语言模型了。作者选择使用 GPT-2 模型，并对其进行微调，使其能够更好地理解和生成代码。
作者首先将源代码文件转换为 token 化的 512-token 块，并逐行保存在文本文件中，而不是在运行时对数据进行 token 化（这非常昂贵）。为了有效地加载这些数据，作者使用 PyTorch 的 Dataset 接口实现了一个自定义的 CodeDataset 类：
import torch
from torch.utils.data import Dataset
import linecache
import json
class CodeDataset(Dataset):
def __init__(self, file_path):
self.file_path = file_path
self.num = sum(1 for _ in open(file_path))
def __len__(self):
return self.num
def __getitem__(self, i):
line = linecache.getline(self.file_path, i+1).strip()
return {"input_ids": torch.tensor(json.loads(line), dtype=torch.long)}
这种方法的优点在于：
- 高效的内存使用（一次只加载一个块）。
- 完全兼容 HuggingFace Trainer。
- 易于迭代并拆分为训练/验证子集。
使用基于块的策略进行 Tokenization，可以最大限度地保持 token 上下文（512 个 token）：
def tokenize_and_save(tokenizer, input_dir, output_file, block_size):
for file in input_dir:
tokens = tokenizer.encode(text, add_special_tokens=True)
for i in range(0, len(tokens), block_size):
block = tokens[i:i+block_size]
if len(block) < block_size:
block += [tokenizer.pad_token_id] * (block_size - len(block))
fout.write(json.dumps(block) + "\n")
基于块的 Tokenization 最大限度地提高了训练吞吐量，同时最大限度地减少了 token 截断。例如，如果直接将整个文件作为输入，可能会导致 token 数量超过模型的最大限制，从而需要截断部分内容。而使用基于块的 Tokenization，可以将文件分成多个块，每个块都包含固定数量的 token，从而避免了截断的问题。
七、 预训练设置：Trainer、Arguments 和策略
作者使用了 HuggingFace 的 Trainer API，并由 TrainerArguments 配置提供支持。关键决策包括：
- 模型：gpt2
- 梯度累积：32（模拟更大的批次大小）
- Checkpointing：每 500 步保存一次
- Eval Steps：每 20,000 步评估一次
- Training Steps：总共 60,000 步（在 Google Colab 上训练约 24 小时）
以下是 TrainingArguments 的关键配置：
TrainingArguments(
output_dir=checkpoint_dir,
save_strategy="steps",
save_steps=500,
eval_strategy="steps",
eval_steps=20000,
gradient_accumulation_steps=32,
report_to="wandb",
run_name="DevScribe-GPT2M"
)
为什么要使用梯度累积？这允许在不超出内存限制的情况下进行有效的批次大小缩放——这对于在有限的硬件上进行训练至关重要。
作者还设计了从最新保存的检查点恢复训练的逻辑，确保中断不会丢失进度。
def get_latest_checkpoint():
ckpts = [d for d in os.listdir(checkpoint_dir) if d.startswith("checkpoint")]
latest = sorted(ckpts, key=lambda x: int(x.split("-")[1]))[-1] if ckpts else None
return os.path.join(checkpoint_dir, latest)
if a checkpoint is found, training continues from it:
trainer.train(resume_from_checkpoint=latest_ckpt)
八、 使用 Weights & Biases (W&B) 进行实时监控
为了跟踪训练指标，作者集成了 Weights & Biases，使用了 HuggingFace 的内置支持：
import wandb
wandb.login()
run = wandb.init(project="huggingface", name="DevScribe-GPT2M", reinit=True)
通过 W&B，可以实时监控训练过程中的损失、学习率、梯度范数等指标，从而更好地了解模型的训练情况。
九、 Prompt 评估：模型的效果如何？
训练完成后，作者评估了模型在各种提示上的表现。以下是两个突出的完成示例，展示了其学习到的行为。
Prompt 1：GitHub 用户身份验证
Prompt:
def authenticate_github_user(token):
headers = {"Authorization": f"Bearer {token}"}
url = "https://api.github.com/user"
response = requests.get(url, headers=headers)
Generated:
if response.statusCode != 200:
raise ValueError("Invalid response")
return response
def _process_authenticate_token(self):
headers = {"Authorization": f"Bearer {token}"}
url = "https://api.github.com/user"
response = requests.get(url, headers=headers)
if response.statusCode != 200:
raise ValueError("Invalid response")
return response
Prompt 2：获取仓库 README
Prompt:
import requests
def get_repo_readme(owner, repo):
url = f"https://api.github.com/repos/{owner}/{repo}/readme"
headers = {"Accept": "application/vnd.github.v3+json"}
response = requests.get(url, headers=headers)
Generated:
if response is not None:
logger.warningf("Couldn't get repo readme response from %s\n", response.message)
return
return response
def get_repo_file(owner, repo):
url = f"https://api.github.com/repos/{owner}/{repo}/file"
headers = {"Accept": "application/vnd.github.v3+json"}
response = requests.get(url, headers=headers)
if response is not None:
logger.warningf("Couldn't get repo readme response from %s\n", response.message)
return
return response
这些完成示例表明，该模型学习了错误处理约定和 API 调用验证——这是后端编码中的一个重要模式。模型能够根据给定的 prompt，生成符合语法规则和语义逻辑的代码，并且能够处理一些常见的错误情况，例如无效的响应或无法获取 README 文件。
十、 模型导出：部署代码生成大模型
训练完成后，保存模型和 tokenizer：
final_dir = os.path.join(output_dir, "final_model")
trainer.save_model(final_dir)
tokenizer.save_pretrained(final_dir)
最终得到一个专门针对代码生成任务进行微调的可部署 GPT-2 模型。它的训练利用了直接从 GitHub 仓库获取的广泛的真实世界数据。
结论：从数据集到 DevScribe
本文主要参考的文章展示了如何使用开放工具，通过定制的数据集和高效的管道来引导强大的、领域特定的模型。 从收集真实世界数据到构建和训练专门的语言模型，这个过程不仅仅是关于代码——而是关于从头开始创造有目的的东西。
关键要点包括：
- 构建了一个强大的 REST-API 管道，用于获取、过滤和去重排名最高的 GitHub 仓库。
- 将原始代码转换为固定长度的 token 块，并将其包装在 PyTorch Dataset 类 (CodeDataset) 中，以便进行内存高效的训练和与 HuggingFace 的 Trainer 无缝集成。
- 利用梯度累积、检查点恢复和基于块的训练来预训练 GPT2 模型。
- 集成了 Weights & Biases 以进行实时指标跟踪，并展示了模型的能力。
“预测未来的最好方法就是创造未来。” 从收集真实世界数据到构建和训练专门的语言模型，这个过程不仅仅是关于代码——而是关于从头开始创造有目的的东西。 随着继续推进这个项目，期待分享更多的见解，并构建推动开放 AI 发展边界的工具。
感谢您的阅读！