<a href="https://arxiv.org/pdf/2306.00029.pdf">技术报告</a>, <a href="https://opensource.salesforce.com/CodeTF/latest/index.html">文档</a>, <a href="https://github.com/salesforce/CodeTF/tree/main/test_inference">示例</a>,

CodeTF - 一站式最先进的代码大语言模型Transformer库

</div>

简介

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<br>Multi: 350M, 2B, 6B<br>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}
}