目录
简介
CodeTF是一个一站式的基于Python的Transformer库,用于代码大语言模型(Code LLMs)和代码智能,为代码智能任务(如代码摘要、翻译、代码生成等)的训练和推理提供了无缝接口。它旨在促进将最先进的CodeLLMs轻松集成到实际应用中。
除了核心LLMs的代码功能外,CodeTF还提供了跨各种语言的代码操作实用工具,包括轻松提取代码属性。使用tree-sitter作为其核心AST解析器,它能够解析诸如函数名、注释和变量名等属性。提供了许多语言的预构建库,无需复杂的解析器设置。因此,CodeTF确保了一个用户友好且易于使用的代码智能任务环境。
该库的当前版本提供:
- 快速模型服务:我们支持一个易于使用的接口,可以使用预量化模型(int8、int16、float16)进行快速推理。CodeTF处理所有设备管理方面的问题,因此用户无需担心这方面的问题。如果您的模型很大,我们提供了高级功能,如跨GPU的权重分片,以更快地提供模型服务。
- 微调您自己的模型:我们提供了一个API,可以使用分布式环境中的最先进参数高效微调技术(HuggingFace PEFT)快速微调您自己的代码LLMs。
- 支持的任务:自然语言到代码、代码摘要、代码补全、代码翻译、代码优化、克隆检测、缺陷预测。
- 数据集+:我们已预处理了知名基准测试(Human-Eval、MBPP、CodeXGLUE、APPS等),并为这些数据集提供了易于加载的功能。
- 模型评估器:我们提供了一个接口,可以在知名基准测试(例如Human-Eval)上使用流行的指标(例如pass@k)评估模型,只需很少的努力(约15行代码)。
- 预训练模型:我们提供了最先进的代码基础语言模型(CodeBERT、CodeT5、CodeGen、CodeT5+、Incoder、StarCoder等)的预训练检查点。
- 微调模型:我们提供了8+个下游任务的微调检查点。
- 源代码操作实用工具:我们提供了实用工具,可以轻松操作源代码,例如用户友好的AST解析器(基于tree-sitter),支持15+种编程语言,以提取重要的代码特征,如函数名、标识符等。
下表显示了支持的模型及其大小和模型支持的任务。这是一项持续的努力,我们正在努力进一步扩大列表。
模型 | 大小 | 任务 |
---|---|---|
CodeT5 | Base, Base-multi-sum, Base-translate-cs, Base-translate-java, Base-sum, Base-clone, Base-defect | 预训练、自然语言到代码、优化、翻译(C#到Java、Java到C#)、摘要(Python、Go、PHP、JavaScript、Java、Ruby)、克隆检测、缺陷预测 |
CodeT5+ | Plus-instruct-16B, Plus-16B, Plus-6B, Plus-2B, Plus-770M-python, Plus-770M, Plus-220M | 预训练、自然语言到代码、优化、缺陷预测 |
CodeGen | Mono: 350M, 2B, 6B, 1B, 3.7B, 7B, 16B Multi: 350M, 2B, 6B NL: 350M, 2B | 预训练 |
StarCoder | 15.5B | 预训练 |
SantaCoder | 1.1B | 预训练 |
GPT-NeoX | 20B | 预训练 |
GPT-Neo | 1.3B | 预训练 |
GPT-J | 6B | 预训练 |
Incoder | 6B | 预训练 |
CodeParrot | Small-python (110M), Small-multi(110M), 1.5B | 预训练 |
CodeBERT | CodeBERT-base, UnixCoder-base, CodeBERTa-small | 预训练 |
安装指南
- (可选)创建conda环境
conda create -n codetf python=3.8
conda activate codetf
- 从PyPI安装:
pip install salesforce-codetf
- 或者,从源代码构建CodeTF:
git clone https://github.com/salesforce/CodeTF.git
cd CodeTF
pip install -e .
此外,为确保量化功能正常工作,还需安装以下依赖项:
pip install -q -U git+https://github.com/huggingface/transformers.git
pip install -q -U git+https://github.com/huggingface/peft.git
pip install -q -U git+https://github.com/huggingface/accelerate.git
对于某些模型,如StarCoder,需要登录Huggingface。请获取HuggingFace令牌并登录:
huggingface-cli login
入门指南
推理流程
使用我们的模型加载流程函数load_model_pipeline()
,可以简单快速地开始使用CodeTF。以下示例展示了如何加载codet5+模型并对代码生成任务进行推理:
from codetf.models import load_model_pipeline
code_generation_model = load_model_pipeline(model_name="codet5", task="pretrained",
model_type="plus-770M-python", is_eval=True,
load_in_8bit=True, load_in_4bit=False, weight_sharding=False)
result = code_generation_model.predict(["def print_hello_world():"])
print(result)
需要考虑以下几个重要参数:
model_name
:模型名称,目前支持codet5
和causal-lm
。model_type
:每个模型名称的模型类型,如base
、codegen-350M-mono
、j-6B
等。load_in_8bit
和load_in_4bit
:继承自Huggingface量化的动态量化功能。weight_sharding
:我们的高级功能,利用HuggingFace分片检查点将大型模型分割成几个较小的分片,分布在不同的GPU上。如果您处理大型模型,请考虑使用此功能。
模型库
您可能想查看所有支持的模型。为此,可以使用model_zoo()
:
from codetf.models import model_zoo
print(model_zoo)
# ============================================================================================================
# Architectures Types Tasks
# ============================================================================================================
# causallm codegen-350M-mono pretrained
# codegen-350M-multi pretrained
# codegen-350M-nl pretrained
# codegen-2B-mono pretrained
# codegen-2B-multi pretrained
# codegen-2B-nl pretrained
# codegen-6B-mono pretrained
# codegen-6B-nl pretrained
# codegen-6B-multi pretrained
# starcoder-15.5B pretrained
# gpt-neox-20B pretrained
# gpt-neo-1.3B pretrained
# gpt-j-6B pretrained
# incoder-6B pretrained
# codegen2-1B pretrained
# codegen2-3.7B pretrained
# codegen2-7B pretrained
# codegen2-16B pretrained
# codet5 base-multi-sum pretrained
# base nl2code
# base refine
# base translate_cs_java
# base translate_java_cs
# base sum_python
# base sum_go
# base sum_php
# base sum_javascript
# base sum_java
# base sum_ruby
# base clone
# base defect
# plus-instruct-16B pretrained
# plus-16B pretrained
# plus-6B pretrained
# plus-2B pretrained
# plus-770M-python pretrained
# plus-770M pretrained
# plus-220M pretrained
# bert codebert-base pretrained
# unixcoder-base pretrained
# codeberta-small pretrained
微调流程
想要为代码训练自定义LLM吗?我们为您提供了解决方案。以下是使用Seq2SeqTrainer
微调CodeT5+预训练模型的示例,结合我们的数据集工具,可以轻松使用CodeXGLUE数据集微调您的模型。以下是一个示例:
from codetf.trainer.codet5_trainer import CodeT5Seq2SeqTrainer
from codetf.data_utility.codexglue_dataset import CodeXGLUEDataset
from codetf.models import load_model_pipeline
from codetf.performance.evaluation_metric import EvaluationMetric
from codetf.data_utility.base_dataset import CustomDataset
model_class = load_model_pipeline(model_name="codet5", task="pretrained",
model_type="plus-220M", is_eval=True)
dataset = CodeXGLUEDataset(tokenizer=model_class.get_tokenizer())
train, test, validation = dataset.load(subset="text-to-code")
train_dataset= CustomDataset(train[0], train[1])
test_dataset= CustomDataset(test[0], test[1])
val_dataset= CustomDataset(validation[0], validation[1])
evaluator = EvaluationMetric(metric="bleu", tokenizer=model_class.tokenizer)
# peft可以是["lora", "prefixtuning"]
trainer = CodeT5Seq2SeqTrainer(train_dataset=train_dataset,
validation_dataset=val_dataset,
peft="lora",
pretrained_model_or_path=model_class.get_model(),
tokenizer=model_class.tokenizer)
trainer.train()
与StarCoder的这个脚本相比,该脚本需要~300行代码来微调模型,而我们只需要14行代码就能完成相同的任务!!!
在知名基准测试上进行评估
计划复现Human-Eval
等知名基准测试的结果,但难以达到原始论文中报告的相同数字?担心复杂的评估过程?别担心,我们为您提供了一个直观、易用的接口。以下是一个示例代码片段,展示了如何使用pass@k(k=[1,10,100])作为指标评估Human Eval:
from codetf.models import load_model_pipeline
from codetf.data_utility.human_eval_dataset import HumanEvalDataset
from codetf.performance.model_evaluator import ModelEvaluator
os.environ["HF_ALLOW_CODE_EVAL"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "true"
model_class = load_model_pipeline(model_name="causal-lm", task="pretrained",
model_type="codegen-350M-mono", is_eval=True,
load_in_8bit=True, weight_sharding=False)
dataset = HumanEvalDataset(tokenizer=model_class.get_tokenizer())
prompt_token_ids, prompt_attention_masks, references= dataset.load()
problems = TensorDataset(prompt_token_ids, prompt_attention_masks)
evaluator = ModelEvaluator(model_class)
avg_pass_at_k = evaluator.evaluate_pass_k(problems=problems, unit_tests=references)
print("Pass@k: ", avg_pass_at_k)
与HuggingFace的这个脚本相比,该脚本需要~230行代码来评估pass@k,而我们只需要14行代码就能完成相同的任务!!!
加载预处理数据
CodeTF为多个知名数据集提供了数据集工具,如CodeXGLUE、Human Eval、MBPP和APPS。以下是加载CodeXGLUE数据集的示例:
from codetf.data_utility.codexglue_dataset import CodeXGLUEDataset
from transformers import RobertaTokenizer
tokenizer = RobertaTokenizer.from_pretrained("Salesforce/codet5-base", use_fast=True)
dataset = CodeXGLUEDataset(tokenizer=tokenizer)
train, test, validation = dataset.load(subset="text-to-code")
"train"、"test"、"validation" 以 PyTorch 张量的形式返回,为用户提供灵活性,以便将其封装到更高级的包装器中用于自己的用例。
代码工具
除了为 LLM 提供工具外,CodeTF 还为用户配备了有效操作源代码的工具。这在代码智能管道中至关重要,因为通常需要执行诸如将代码解析为抽象语法树(AST)或提取代码属性(如函数名或标识符)等操作(CodeT5)。这些任务可能具有挑战性,特别是在需要设置和多语言支持时。我们的代码工具接口提供了一个简化的解决方案,便于跨 15 种以上语言轻松解析和提取代码属性。
多语言 AST 解析器
CodeTF 包含与众多编程语言兼容的 AST 解析器。以下是将 Apex 代码解析为 AST 的示例:
from codetf.code_utility.apex.apex_code_utility import ApexCodeUtility
apex_code_utility = ApexCodeUtility()
sample_code = """
public class SampleClass {
public Integer myNumber;
**
* This is a method that returns the value of myNumber.
* @return An integer value
*/
public Integer getMyNumber() {
// Return the current value of myNumber
return this.myNumber;
}
}
"""
ast = apex_code_utility.parse(sample_code)
# 这将打印 tree-sitter AST 对象
print(ast)
然后,您可以使用 py-tree-sitter 的接口遍历树
root_node = ast.root_node
assert root_node.type == 'module'
assert root_node.start_point == (1, 0)
assert root_node.end_point == (3, 13)
还有其他用于 Java、Python 等的工具,可以执行相同的操作。
提取代码属性
CodeTF 提供了一个接口,可以轻松提取代码属性。以下是提取 Python 函数名的示例:
code_attributes = apex_code_utility.get_code_attributes(sample_code)
print(code_attributes)
这将打印:
{'class_names': ['AccountWithContacts'], 'method_names': ['getAccountsWithContacts'], 'comments': [], 'variable_names': ['acc', 'accounts', 'con', 'System', 'debug', 'Contacts', 'Id', 'Name', 'Account', 'Email', 'LastName']}
移除注释
还有其他现有的工具,如从代码中移除注释:
new_code_snippet = apex_code_utility.remove_comments(sample_code)
print(new_code_snippet)
这将打印:
public class SampleClass {
public Integer myNumber;
public Integer getMyNumber() {
return this.myNumber;
}
}
请注意,这是一个持续的过程,我们将在未来添加更多功能来提取复杂的代码属性。更多示例可以在这里找到。
更多示例
您可以找到每个用例的更多示例:
注意事项
- CodeTF 旨在补充和增强 HuggingFace Transformers 的功能,而非取代它。它作为一个专门为代码智能任务定制的专用层,如使用代码特定功能微调语言模型,以及在知名代码智能基准上进行评估。如果用户需要更多定制,鼓励他们从头开始编写自己的训练代码。
- CodeTF 利用 Accelerate 提供的强大功能进行推理和训练。通过 Accelerate,用户无需手动管理大多数操作的 GPU 或 CPU 设备,从而实现流畅高效的工作流程。
道德和负责任的使用
尽管 CodeTF 功能强大,但不能保证代码智能能力绝对无误。用户可能会遇到不准确或偏差,可能导致误解或不良行为。风险包括生成不安全的代码、传播糟糕的编码实践,或无意中泄露敏感数据。我们强烈建议用户在实际采用前检查预训练模型和系统。CodeTF 促进有效的代码分析、预测和调试,推动可重复的研究和开发。我们鼓励负责任地使用它来提高软件质量和开发者生产力。
然而,滥用可能导致非法代码操作、隐私泄露或不安全的编码实践等不道德后果。用户在使用 CodeTF 之前应熟悉负责任的 AI 指南。我们致力于通过识别和缓解潜在偏见和不当行为来不断完善该库。用户应在实际实施前审查模型和系统,并为完善库做出贡献,以确保道德使用。
技术报告和引用 CodeTF
您可以在我们的技术报告中找到更多详细信息。
如果您在研究或应用中使用 CodeTF,请使用以下 BibTeX 进行引用:
@misc{nghi2023codetf,
title={CodeTF: A Transformer-based Library for CodeLLM & Code Intelligence},
author={Nghi D. Q. Bui, Henry Le, Yue Wang, Akhilesh Deepak Gotmare, Junnan Li, Steven Hoi.},
year={2023},
eprint={2209.09019},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
联系我们
如果您有任何问题、意见或建议,请随时通过 codetf@salesforce.com 与我们联系。