codesage/codesage-large

1---2license: apache-2.03datasets:4- bigcode/the-stack-dedup5library_name: transformers6language:7- code8---9 10## CodeSage-Large11 12### Updates13* [12/2024] <span style="color:blue">We are excited to announce the release of the CodeSage V2 model family with largely improved performance and flexible embedding dimensions!</span> Please check out our [models](https://huggingface.co/codesage) and [blogpost](https://code-representation-learning.github.io/codesage-v2.html) for more details.14* [11/2024] You can now access CodeSage models through SentenceTransformer.15 16### Model description17CodeSage is a new family of open code embedding models with an encoder architecture that support a wide range of source code understanding tasks. It is introduced in the paper:18 19[Code Representation Learning At Scale by 20Dejiao Zhang*, Wasi Uddin Ahmad*, Ming Tan, Hantian Ding, Ramesh Nallapati, Dan Roth, Xiaofei Ma, Bing Xiang](https://arxiv.org/abs/2402.01935) (* indicates equal contribution).21 22### Pretraining data23This checkpoint is trained on the Stack data (https://huggingface.co/datasets/bigcode/the-stack-dedup). Supported languages (9 in total) are as follows: c, c-sharp, go, java, javascript, typescript, php, python, ruby.24 25### Training procedure26This checkpoint is first trained on code data via masked language modeling (MLM) and then on bimodal text-code pair data. Please refer to the paper for more details.27 28### How to Use29This checkpoint consists of an encoder (1.3B model), which can be used to extract code embeddings of 1024 dimension. 30 311. Accessing CodeSage via HuggingFace: it can be easily loaded using the AutoModel functionality and employs the [Starcoder Tokenizer](https://arxiv.org/pdf/2305.06161.pdf).32   33```34from transformers import AutoModel, AutoTokenizer35 36checkpoint = "codesage/codesage-large"37device = "cuda"  # for GPU usage or "cpu" for CPU usage38 39# Note: CodeSage requires adding eos token at the end of each tokenized sequence 40 41tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True, add_eos_token=True)42 43model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True).to(device)44 45inputs = tokenizer.encode("def print_hello_world():\tprint('Hello World!')", return_tensors="pt").to(device)46embedding = model(inputs)[0]47```48 492. Accessing CodeSage via SentenceTransformer50```51from sentence_transformers import SentenceTransformer52model = SentenceTransformer("codesage/codesage-large", trust_remote_code=True)53```54 55### BibTeX entry and citation info56```57@inproceedings{58    zhang2024codesage,59    title={CodeSage: Code Representation Learning At Scale},60    author={Dejiao Zhang* and Wasi Ahmad* and Ming Tan and Hantian Ding and Ramesh Nallapati and Dan Roth and Xiaofei Ma and Bing Xiang},61    booktitle={The Twelfth International Conference on Learning Representations},62    year={2024},63    url={https://openreview.net/forum?id=vfzRRjumpX}64}65```