/CodeGen2

CodeGen2 models for program synthesis

Primary LanguagePythonApache License 2.0Apache-2.0

CodeGen2

Official research release for the CodeGen2 models (1B, 3B, 7B, 16B) for Program Synthesis as presented in ICLR 2023:

Title: CodeGen2: Lessons for Training LLMs on Programming and Natural Languages

Authors: Erik Nijkamp*, Hiroaki Hayashi*, Caiming Xiong, Silvio Savarese, and Yingbo Zhou (* indicates equal contribution)

Hugging Face Integration

Model checkpoints are published at Hugging Face Hub.

Model cards outline how to use the model for causal and infill sampling.

Sampling

Program synthesis in the form of auto-regressive sampling can be performed as follows:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen2-7B")
model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen2-7B", trust_remote_code=True, revision="main")
inputs = tokenizer("# this function prints hello world", return_tensors="pt")
sample = model.generate(**inputs, max_length=128)
print(tokenizer.decode(sample[0], truncate_before_pattern=[r"\n\n^#", "^'''", "\n\n\n"]))

Citation

@article{Nijkamp2023codegen2,
  title={CodeGen2: Lessons for Training LLMs on Programming and Natural Languages},
  author={Nijkamp, Erik and Hayashi, Hiroaki and Xiong, Caiming and Savarese, Silvio and Zhou, Yingbo},
  journal={arXiv preprint},
  year={2023}
}