DOCSEARCHbyAlgolia
Submit Docs

CodeT5 and CodeT5+

repository·main·Indexed 25 days ago

https://github.com/salesforce/codet5

Specialized large language models from Salesforce Research optimized for code understanding and generation tasks, including text-to-code, autocompletion, and summarization. The repository provides tools for inference using Hugging Face transformers, fine-tuning via run_exp.py and DeepSpeed, and extracting code embeddings with CodeT5+ 110M. It supports various tasks such as code translation, defect detection, and clone detection across multiple languages.

Tokens
5.1K
Snippets
11
Records
21
Agent score
85%

What's inside CodeT5

  1. Overview of CodeT5 and CodeT5+

    main

    CodeT5 and CodeT5+ are open-source large language models developed by Salesforce Research designed for Code Understanding and Generation. These models can be used to build AI-powered coding assistants with capabilities such as:

    • Text-to-code generation: Generating code from natural language descriptions.
    • Code autocompletion: Completing entire functions based on a target function name.
    • Code summarization: Generating natural language summaries of code functions.

  3. Perform Text-to-Code Retrieval with CodeT5+ Bimodal model

    main

    The CodeT5+ 220M bimodal model improves retrieval performance by using a matching decoder to rerank the top k candidates from the initial embedding retrieval.

    Execution: Navigate to code_retrieval and run eval_match_retrieval.py. Use the --top_k flag to control the number of candidates to rerank.

    cd code_retrieval

# Example configuration for Ruby with reranking
LANG=ruby
BS=256
CODE_LEN=360
TEXT_LEN=64
TOPK=32
MODEL_NAME=Salesforce/codet5p-220m-bimodal
DATA_DIR=/path/to/data

TRG_DIR=saved_models/${LANG}/codet5p_220m_bimodal_TL${TEXT_LEN}_CL${CODE_LEN}_top${TOPK}
mkdir -p $TRG_DIR

python eval_match_retrieval.py --model_name $MODEL_NAME --lang $LANG --output_dir $TRG_DIR \
  --data_dir $DATA_DIR --max_text_len $TEXT_LEN --max_code_len $CODE_LEN --batch_size $BS --top_k $TOPK

  4. Fine-tune on a custom task and dataset

    main

    To add a custom dataset and task to CodeT5, follow these steps:

    1. Register Task: Add your new task and sub_task definitions in configs.py.
    2. Data Loading:
      • Add your data path and a reading function in utils.py.
      • Implement the reading function in _utils.py (refer to existing implementations for guidance).
    3. Execution Logic:
      • For generation tasks: Reuse or customize run_gen.py.
      • For understanding tasks: Refer to run_defect.py or run_clone.py for implementation patterns.

  5. Run fine-tuning experiments with run_exp.py

    main

    You can run experiments using run_exp.py by passing specific arguments for the model, task, and sub-task.

    Setup Step: Before running, navigate to the sh folder and set the WORKDIR variable in exp_with_args.sh to the absolute path of your cloned CodeT5 repository.

    Supported Models (--model_tag): roberta, codebert, bart_base, codet5_small, codet5_base

    Supported Tasks and Sub-tasks:

    --task--sub_taskDescription
    summarizeruby/javascript/go/python/java/phpCode summarization on CodeSearchNet
    concodenoneText-to-code generation on Concode
    translatejava-cs/cs-javaJava-to-C# or C#-to-Java translation
    refinesmall/mediumCode refinement on repair data
    defectnoneCode defect detection in C/C++
    clonenoneCode clone detection in Java
    multi_tasknoneMulti-task training

    Additional Arguments:

    • --model_dir: Directory to save fine-tuning checkpoints.
    • --res_dir: Directory to save performance results.
    • --summary_dir: Directory to save training curves (can be visualized with tensorboard).
    • --data_num: Number of data instances to use (default -1 uses full data).
    • --gpu: Index of the GPU to use.

  6. Load CodeT5+ generative models

    main

    CodeT5+ models (including 2B, 6B, 16B, and InstructCodeT5+ 16B) can be loaded using Hugging Face's AutoModelForSeq2SeqLM and AutoTokenizer.

    Important Requirements:

    • For models 2B and larger, you must set trust_remote_code=True when loading the model because the model class is defined within the Hugging Face repository.
    • To improve generation performance, it is recommended to pass additional prompts to the decoder via decoder_input_ids (e.g., by cloning the input_ids).
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch

checkpoint = "Salesforce/instructcodet5p-16b"
device = "cuda" # for GPU usage or "cpu" for CPU usage

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint,
                                              torch_dtype=torch.float16,
                                              low_cpu_mem_usage=True,
                                              trust_remote_code=True).to(device)

encoding = tokenizer("def print_hello_world():", return_tensors="pt").to(device)
encoding['decoder_input_ids'] = encoding['input_ids'].clone()
outputs = model.generate(**encoding, max_length=15)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

  7. Perform Text-to-Code Retrieval with CodeT5+ Embedding model

    main

    Use the CodeT5+ 110M embedding model for contrastive retrieval tasks.

    Prerequisites: Download and preprocess the CSN (6 PLs), AdvTest, or cosqa datasets following the UniXcoder instructions.

    Execution: Navigate to code_retrieval and run eval_contrast_retrieval.py.

    Supported Languages (LANG): ruby, javascript, go, python, java, php, AdvTest, cosqa.

    cd code_retrieval

# Example configuration for Ruby
LANG=ruby
BS=256
CODE_LEN=360
TEXT_LEN=64
MODEL_NAME=Salesforce/codet5p-110m-embedding
DATA_DIR=/path/to/data

TRG_DIR=saved_models/${LANG}/codet5p_110m_embedding_TL${TEXT_LEN}_CL${CODE_LEN}
mkdir -p $TRG_DIR

python eval_contrast_retrieval.py --model_name $MODEL_NAME --lang $LANG --output_dir $TRG_DIR \
  --data_dir $DATA_DIR --max_text_len $TEXT_LEN --max_code_len $CODE_LEN --batch_size $BS

  8. Finetune CodeT5+ using custom Seq2Seq data

    main

    You can finetune CodeT5+ models on any Seq2Seq LM task (e.g., Python code summarization) using the tune_codet5p_seq2seq.py script.

    Steps:

    1. Install transformers and datasets libraries.
    2. Prepare your data in the Hugging Face datasets format.
    3. Run the script, passing your data path to --cache-data.

    Key Arguments:

    • --load: Select the specific model to finetune from (e.g., Salesforce/codet5p-220m).
    • --cache-data: Path to your customized dataset.
    • --fp16: Enable mixed-precision training to save memory.
    • --deepspeed: Path to a DeepSpeed config file for optimization.
    • Other tunable arguments: --epochs, --lr, --lr-warmup-steps, --max-source-len, --max-target-len, --batch-size-per-replica, --grad-acc-steps.

  9. Reproduce results using released fine-tuned checkpoints

    main

    To evaluate existing fine-tuned checkpoints instead of training from scratch:

    1. Modify the script in sh/exp_with_args.sh: Remove --do_train --do_eval --do_eval_bleu and keep only --do_test.
    2. Update the checkpoint path in run_gen.py to point to your downloaded .bin file (e.g., file = "CodeT5/finetuned_models/summarize_python_codet5_base.bin").
    3. Execute the experiment using run_exp.py with the appropriate task arguments.

  10. Evaluate Pass@k for CodeT5+ models

    main

    After generating programs, evaluate their functional correctness using the humaneval directory scripts.

    To evaluate your own generated predictions:

    1. Process the predictions into a .jsonl format.
    2. Run the functional correctness evaluation.

    To quickly verify results using the provided InstructCodeT5+ 16B reference file, run evaluate_functional_correctness directly on the released .jsonl file.

    cd humaneval

# Evaluate generated predictions
output_path=preds/instructcodet5p-16b_T0.2_N200
python process_preds.py --path ${output_path} --out_path ${output_path}.jsonl
evaluate_functional_correctness ${output_path}.jsonl

# Or evaluate the released reference file directly
evaluate_functional_correctness humaneval/instructcodet5p-16b_T0.2_N200.jsonl

  11. Calculate CodeBLEU scores using calc_code_bleu.py

    main
    Use the calc_code_bleu.py script to evaluate code generation models by comparing candidate files against reference files. You must specify the language (e.g., java or c_sharp) and provide paths to the reference and hypothesis files. You can optionally tune the weights for the four components of CodeBLEU using the --params flag.

  12. Download CodeT5 pretrained and fine-tuned checkpoints

    main

    To use the models for fine-tuning or reproduction, you must download the checkpoints and data from Google Cloud Storage using gsutil. Ensure you have gsutil installed via pip first.

    # pip install gsutil
cd your-cloned-codet5-path

gsutil -m cp -r "gs://sfr-codet5-data-research/pretrained_models" .
gsutil -m cp -r "gs://sfr-codet5-data-research/data" .
gsutil -m cp -r "gs://sfr-codet5-data-research/finetuned_models" .