updated readme

.ipynb_checkpoints/config-checkpoint.json

@@ -1,25 +0,0 @@
-{
-    "_name_or_path": "codesage/codesage-small-v2",
-    "architectures": [
-        "CodeSage"
-    ],
-    "auto_map": {
-        "AutoConfig": "config_codesage.CodeSageConfig",
-        "AutoTokenizer": "tokenization_codesage.CodeSageTokenizer",
-        "AutoModel": "modeling_codesage.CodeSageModel",
-        "AutoModelForMaskedLM": "modeling_codesage.CodeSageForMaskedLM",
-        "AutoModelForSequenceClassification": "modeling_codesage.CodeSageForSequenceClassification"
-    },
-    "activation_function": "gelu_new",
-    "attention_dropout_prob": 0.1,
-    "embedding_dropout_prob": 0.1,
-    "initializer_range": 0.02,
-    "layer_norm_epsilon": 1e-05,
-    "hidden_size": 1024,
-    "num_attention_heads": 8,
-    "num_hidden_layers": 6,
-    "intermediate_size": 4096,
-    "max_position_embeddings": 2048,
-    "residual_dropout_prob": 0.1,
-    "vocab_size": 49154
-}

.ipynb_checkpoints/config_codesage-checkpoint.py

@@ -1,52 +0,0 @@
-#!/usr/bin/env python
-# coding=utf-8
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-
-from transformers.configuration_utils import PretrainedConfig
-
-CODESAGE_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "codesage/codesage-small-v2": "https://huggingface.co/codesage/codesage-small-v2/resolve/main/config.json",
-    "codesage/codesage-base-v2": "https://huggingface.co/codesage/codesage-base-v2/resolve/main/config.json",
-    "codesage/codesage-large-v2": "https://huggingface.co/codesage/codesage-large-v2/resolve/main/config.json",
-}
-
-
-class CodeSageConfig(PretrainedConfig):
-    model_type = "codesage"
-
-    def __init__(
-            self,
-            vocab_size=50257,
-            max_position_embeddings=1024,
-            hidden_size=1024,
-            num_hidden_layers=24,
-            num_attention_heads=8,
-            intermediate_size=4096,
-            activation_function="gelu_new",
-            residual_dropout_prob=0.1,
-            embedding_dropout_prob=0.1,
-            attention_dropout_prob=0.1,
-            layer_norm_epsilon=1e-5,
-            initializer_range=0.02,
-            position_embedding_type='absolute',
-            bos_token_id=0,
-            eos_token_id=0,
-            pad_token_id=49153,
-            **kwargs
-    ):
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.intermediate_size = intermediate_size
-        assert 'gelu' in activation_function
-        self.activation_function = activation_function
-        self.residual_dropout_prob = residual_dropout_prob
-        self.embedding_dropout_prob = embedding_dropout_prob
-        self.attention_dropout_prob = attention_dropout_prob
-        self.layer_norm_epsilon = layer_norm_epsilon
-        self.initializer_range = initializer_range
-        self.position_embedding_type = position_embedding_type
-
-        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)

.ipynb_checkpoints/modeling_codesage-checkpoint.py

@@ -1,426 +0,0 @@
-#!/usr/bin/env python
-# coding=utf-8
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-
-import math
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
-from transformers.activations import ACT2FN
-from transformers.modeling_utils import Conv1D, PreTrainedModel
-from transformers.utils import logging
-from .config_codesage import CodeSageConfig
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPooling,
-    MaskedLMOutput,
-    SequenceClassifierOutput
-)
-
-logger = logging.get_logger(__name__)
-
-CODESAGE_PRETRAINED_MODEL_ARCHIVE_LIST = [
-    "codesage/codesage-small-v2",
-    "codesage/codesage-base-v2",
-    "codesage/codesage-large-v2",
-    # See all CodeSage models at https://huggingface.co/models?filter=codesage
-]
-
-
-class CodeSageAttention(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.head_dim = config.hidden_size // self.num_heads
-        if self.head_dim * self.num_heads != config.hidden_size:
-            raise ValueError(
-                f"`hidden_size` must be divisible by num_heads "
-                f"(got `hidden_size`: {config.hidden_size} and `num_heads`: {self.num_heads})."
-            )
-
-        self.c_attn = Conv1D(3 * self.hidden_size, self.hidden_size)
-        self.c_proj = Conv1D(self.hidden_size, self.hidden_size)
-
-        self.attention_dropout = nn.Dropout(config.attention_dropout_prob)
-        self.residual_dropout = nn.Dropout(config.residual_dropout_prob)
-
-    def attn(self, query, key, value, attention_mask=None, head_mask=None):
-        attn_weights = torch.matmul(query, key.transpose(-1, -2))
-        attn_weights = attn_weights / math.sqrt(self.head_dim)
-        if attention_mask is not None:
-            attn_weights = attn_weights + attention_mask
-
-        attn_weights = nn.Softmax(dim=-1)(attn_weights)
-        attn_weights = self.attention_dropout(attn_weights)
-        if head_mask is not None:
-            attn_weights = attn_weights * head_mask
-
-        attn_output = torch.matmul(attn_weights, value)
-        return attn_output, attn_weights
-
-    def split_heads(self, tensor, num_heads, attn_head_size):
-        """
-        Splits hidden_size dim into attn_head_size and num_heads
-        """
-        new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
-        tensor = tensor.view(*new_shape)
-        return tensor.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)
-
-    def merge_heads(self, tensor, num_heads, attn_head_size):
-        """
-        Merges attn_head_size dim and num_attn_heads dim into hidden_size
-        """
-        tensor = tensor.permute(0, 2, 1, 3).contiguous()
-        new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
-        return tensor.view(new_shape)
-
-    def forward(
-            self,
-            hidden_states,
-            attention_mask=None,
-            head_mask=None,
-            output_attentions=False,
-    ):
-        query, key, value = self.c_attn(hidden_states).split(self.hidden_size, dim=2)
-        query = self.split_heads(query, self.num_heads, self.head_dim)
-        key = self.split_heads(key, self.num_heads, self.head_dim)
-        value = self.split_heads(value, self.num_heads, self.head_dim)
-
-        attn_output, attn_weights = self.attn(query, key, value, attention_mask, head_mask)
-
-        attn_output = self.merge_heads(attn_output, self.num_heads, self.head_dim)
-        attn_output = self.c_proj(attn_output)
-        attn_output = self.residual_dropout(attn_output)
-
-        outputs = (attn_output, attn_weights) if output_attentions else (attn_output,)
-        return outputs  # a, present, (attentions)
-
-
-class CodeSageMLP(nn.Module):
-    def __init__(self, intermediate_size, config):
-        super().__init__()
-
-        self.c_fc = Conv1D(intermediate_size, config.hidden_size)
-        self.act = ACT2FN[config.activation_function]
-        self.c_proj = Conv1D(config.hidden_size, intermediate_size)
-        self.dropout = nn.Dropout(config.residual_dropout_prob)
-
-    def forward(self, hidden_states):
-        hidden_states = self.c_fc(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.c_proj(hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        return hidden_states
-
-
-class CodeSageBlock(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        hidden_size = config.hidden_size
-        inner_dim = config.intermediate_size if config.intermediate_size is not None else 4 * hidden_size
-        self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-        self.attn = CodeSageAttention(config)
-        self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-        self.mlp = CodeSageMLP(inner_dim, config)
-
-    def forward(
-            self,
-            hidden_states,
-            attention_mask=None,
-            head_mask=None,
-            output_attentions=False,
-    ):
-        residual = hidden_states
-        hidden_states = self.ln_1(hidden_states)
-        attn_outputs = self.attn(
-            hidden_states,
-            attention_mask=attention_mask,
-            head_mask=head_mask,
-            output_attentions=output_attentions
-        )
-        attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)
-        outputs = attn_outputs[1:]
-        hidden_states = attn_output + residual
-
-        residual = hidden_states
-        hidden_states = self.ln_2(hidden_states)
-        feed_forward_hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + feed_forward_hidden_states
-
-        outputs = (hidden_states,) + outputs[1:]
-        return outputs  # hidden_states, present, (attentions)
-
-
-class CodeSagePreTrainedModel(PreTrainedModel):
-    config_class = CodeSageConfig
-    base_model_prefix = "transformer"
-
-    def _init_weights(self, module):
-        """Initialize the weights."""
-        if isinstance(module, (nn.Linear, Conv1D)):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-
-class CodeSageModel(CodeSagePreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.wte = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.wpe = nn.Embedding(config.max_position_embeddings, config.hidden_size)
-
-        self.drop = nn.Dropout(config.embedding_dropout_prob)
-        self.h = nn.ModuleList([CodeSageBlock(config) for _ in range(config.num_hidden_layers)])
-        self.ln_f = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
-
-        self.init_weights()
-
-    def get_input_embeddings(self):
-        return self.wte
-
-    def set_input_embeddings(self, new_embeddings: torch.Tensor):
-        self.wte = new_embeddings
-
-    def forward(
-            self,
-            input_ids=None,
-            attention_mask=None,
-            position_ids=None,
-            head_mask=None,
-            inputs_embeds=None,
-            output_attentions=None,
-            output_hidden_states=None,
-            return_dict=None
-    ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
-        if input_ids is not None:
-            input_shape = input_ids.size()
-        elif inputs_embeds is not None:
-            input_shape = inputs_embeds.size()[:-1]
-        else:
-            raise ValueError("You have to specify either input_ids or inputs_embeds")
-
-        device = input_ids.device if input_ids is not None else inputs_embeds.device
-        if position_ids is None:
-            position_ids = torch.arange(input_shape[-1], dtype=torch.long, device=device)
-            position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])
-        else:
-            position_ids = position_ids.view(-1, input_shape[-1])
-
-        extended_attention_mask = None
-        if attention_mask is not None:
-            assert attention_mask.dim() == 2
-            extended_attention_mask = attention_mask[:, None, None, :]
-            extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
-            extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
-
-        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-        if inputs_embeds is None:
-            inputs_embeds = self.wte(input_ids)
-
-        position_embeds = self.wpe(position_ids)
-        hidden_states = inputs_embeds + position_embeds
-
-        hidden_states = self.drop(hidden_states)
-        output_shape = input_shape + (hidden_states.size(-1),)
-
-        all_self_attentions = () if output_attentions else None
-        all_hidden_states = () if output_hidden_states else None
-        for i, block in enumerate(self.h):
-            if output_hidden_states:
-                all_hidden_states = all_hidden_states + (hidden_states,)
-
-            outputs = block(
-                hidden_states,
-                attention_mask=extended_attention_mask,
-                head_mask=head_mask[i],
-                output_attentions=output_attentions,
-            )
-
-            hidden_states = outputs[0]
-            if output_attentions:
-                all_self_attentions = all_self_attentions + (outputs[1],)
-
-        hidden_states = self.ln_f(hidden_states)
-        hidden_states = hidden_states.view(*output_shape)
-        if output_hidden_states:
-            all_hidden_states = all_hidden_states + (hidden_states,)
-
-        pooled_output = None  # max-pooled output
-        if attention_mask is not None:
-            pooled_output = (hidden_states * attention_mask[:, :, None]).sum(1) / attention_mask.sum(1)[:, None]
-
-        if not return_dict:
-            return tuple(
-                v
-                for v in [hidden_states, pooled_output, all_hidden_states, all_self_attentions]
-                if v is not None
-            )
-
-        return BaseModelOutputWithPooling(
-            last_hidden_state=hidden_states,
-            pooler_output=pooled_output,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attentions
-        )
-
-
-class CodeSageForMaskedLM(CodeSagePreTrainedModel):
-    _tied_weights_keys = ["lm_head.weight"]
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.transformer = CodeSageModel(config)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
-        self.init_weights()
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
-
-    def forward(
-            self,
-            input_ids=None,
-            attention_mask=None,
-            position_ids=None,
-            head_mask=None,
-            inputs_embeds=None,
-            labels=None,
-            output_attentions=None,
-            output_hidden_states=None,
-            return_dict=None
-    ):
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        transformer_outputs = self.transformer(
-            input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict
-        )
-        hidden_states = transformer_outputs[0]
-        lm_logits = self.lm_head(hidden_states)
-
-        masked_lm_loss = None
-        if labels is not None:
-            loss_fct = CrossEntropyLoss()
-            masked_lm_loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
-
-        if not return_dict:
-            output = (lm_logits,) + transformer_outputs[1:]
-            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
-        return MaskedLMOutput(
-            loss=masked_lm_loss,
-            logits=lm_logits,
-            hidden_states=transformer_outputs.hidden_states,
-            attentions=transformer_outputs.attentions,
-        )
-
-
-class CodeSageForSequenceClassification(CodeSagePreTrainedModel):
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.num_labels = config.num_labels
-        self.config = config
-
-        self.transformer = CodeSageModel(config)
-        classifier_dropout = (
-            config.classifier_dropout 
-            if hasattr(config, 'classifier_dropout') and config.classifier_dropout is not None 
-            else config.residual_dropout_prob
-        )
-        self.dropout = nn.Dropout(classifier_dropout)
-        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def forward(
-            self,
-            input_ids=None,
-            attention_mask=None,
-            position_ids=None,
-            head_mask=None,
-            inputs_embeds=None,
-            labels=None,
-            output_attentions=None,
-            output_hidden_states=None,
-            return_dict=None,
-    ):
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        assert attention_mask is not None, "attention_mask is needed to perform max-pooling"
-
-        outputs = self.transformer(
-            input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        pooled_output = outputs[1]
-        pooled_output = self.dropout(pooled_output)
-        logits = self.classifier(pooled_output)
-
-        loss = None
-        if labels is not None:
-            if self.config.problem_type is None:
-                if self.num_labels == 1:
-                    self.config.problem_type = "regression"
-                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
-                    self.config.problem_type = "single_label_classification"
-                else:
-                    self.config.problem_type = "multi_label_classification"
-
-            if self.config.problem_type == "regression":
-                loss_fct = MSELoss()
-                if self.num_labels == 1:
-                    loss = loss_fct(logits.squeeze(), labels.squeeze())
-                else:
-                    loss = loss_fct(logits, labels)
-            elif self.config.problem_type == "single_label_classification":
-                loss_fct = CrossEntropyLoss()
-                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-            elif self.config.problem_type == "multi_label_classification":
-                loss_fct = BCEWithLogitsLoss()
-                loss = loss_fct(logits, labels)
-
-        if not return_dict:
-            output = (logits,) + outputs[2:]
-            return ((loss,) + output) if loss is not None else output
-
-        return SequenceClassifierOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )

.ipynb_checkpoints/tokenizer_config-checkpoint.json

@@ -1,34 +0,0 @@
-{
-  "add_prefix_space": false,
-  "additional_special_tokens": [
-    "<|endoftext|>",
-    "<fim_prefix>",
-    "<fim_middle>",
-    "<fim_suffix>",
-    "<fim_pad>",
-    "<filename>",
-    "<gh_stars>",
-    "<issue_start>",
-    "<issue_comment>",
-    "<issue_closed>",
-    "<jupyter_start>",
-    "<jupyter_text>",
-    "<jupyter_code>",
-    "<jupyter_output>",
-    "<empty_output>",
-    "<commit_before>",
-    "<commit_msg>",
-    "<commit_after>",
-    "<reponame>"
-  ],
-  "bos_token": "<|endoftext|>",
-  "eos_token": "<|endoftext|>",
-  "add_eos_token": true,
-  "model_max_length": 1000000000000000019884624838656,
-  "unk_token": "<|endoftext|>",
-  "vocab_size": 49152,
-  "tokenizer_class": "CodeSageTokenizer",
-  "auto_map": {
-    "AutoTokenizer": ["tokenization_codesage.CodeSageTokenizer", null]
-  }
-}

README.md

@@ -1,3 +1,90 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+datasets:
+- bigcode/the-stack-dedup
+- bigcode/the-stack-v2
+
+library_name: transformers
+language:
+- code
+---
+
+## CodeSage-Base-v2
+
+### Model description
+CodeSage is a family of open code embedding models with an encoder architecture that supports a wide range of source code understanding tasks. It was initially introduced in the paper:
+
+[Code Representation Learning At Scale by Dejiao Zhang*, Wasi Uddin Ahmad*, et al.](https://arxiv.org/abs/2402.01935)  
+
+For this V2 model, we enhanced semantic search performance by improving the quality of the contrastive learning data through [consistency filtering](https://arxiv.org/abs/2209.11755). Starting from the pretrained checkpoint (trained with both Masked Language Modeling (MLM) and deobfuscation [Section 3.1](https://arxiv.org/abs/2402.01935)) from our V1 model training (Zhang et al., 2023), we applied contrastive learning with the filtered data. Unlike the V1 model, we extracted the initial set of (text, code) pairs—specifically, summaries and function/class bodies—from [The Stack V2](https://huggingface.co/datasets/bigcode/the-stack-v2) data instead of using the [V1](https://huggingface.co/datasets/bigcode/the-stack-dedup) data. We employed simple rule-based filtering as detailed in our previous work. We then applied consistency filtering to further refine the data. While using The Stack V2 resulted in minor performance boosts on downstream tasks, the majority of the performance improvements came from the consistency filtering.
+
+### Model Performance
+
+#### 1.Code2Code Search
+| Model Name          | # Params | Embd Dim | Python | Java  | JS    | TS     | C#     | C      | Ruby   | PhP    | GO     | AVG    |
+|---------------------|----------|----------|--------|-------|-------|--------|--------|--------|--------|--------|--------|--------|
+| OpenAI-Code-01      | NA       | 3072     | 21.92  | 8.90  | 4.90  | 5.70   | 3.15   | 11.58  | 26.25  | 16.60  | 9.40   | 12.04  |
+| OpenAI-Ada-002      | NA       | 1536     | 35.91  | 25.13 | 19.01 | 21.86  | 10.17  | 29.15  | 40.85  | 40.47  | 23.43  | 27.33  |
+| OpenAI-Text-3-Small | NA       | 1536     | 25.18  | 12.61 | 8.00  | 9.44   | 5.46   | 15.86  | 30.70  | 23.33  | 11.20  | 15.57  |
+| OpenAI-Text-3-Large | NA       | 3072     | 40.57  | 25.33 | 20.09 | 22.00  | 11.84  | 31.90  | 42.54  | 41.84  | 21.75  | 28.65  |
+| CodeSage-Small      | 130M     | 1024     | 36.31  | 23.97 | 26.60 | 29.90  | 11.84  | 22.84  | 29.06  | 34.64  | 19.56  | 26.08  |
+| CodeSage-Base       | 356M     | 1024     | 47.52  | 22.84 | 28.70 | 31.95  | 13.37  | 30.99  | 44.86  | 51.13  | 25.15  | 32.95  |
+| CodeSage-Large      | 1.3B     | 2048     | 46.70  | 33.13 | 37.16 | 41.18  | 16.81  | 32.89  | 54.12  | 52.13  | 32.48  | 38.51  |
+| CodeSage-v2-Small   | 130M     | 1024     | 45.60  | 33.65 | 39.96 | 47.78  | 19.19  | 30.55  | 40.12  | 55.39  | 30.96  | 38.13  |
+| CodeSage-v2-Base    | 356M     | 1024     | 55.86  | 42.89 | 45.29 | 54.58  | 23.90  | 38.52  | 56.02  | 64.56  | 42.88  | 47.17  |
+| CodeSage-v2-Large   | 1.3B     | 2048     | 61.11  | 47.09 | 51.18 | 60.67  | 28.04  | 43.40  | 60.74  | 67.87  | 43.86  | 51.55  |
+
+
+#### 2. NL2Code Search
+| Model Name          | # Params | CoSQA | AdvTest | Python | Java  | JS    | PhP    | GO     | Ruby   | Avg    |
+|---------------------|----------|-------|---------|--------|-------|-------|--------|--------|--------|--------|
+| OpenAI-Code-01      | NA       | 52.20 | 36.03   | 63.13  | 67.85 | 62.30 | 57.47  | 85.22  | 69.28  | 61.69  |
+| OpenAI-Ada-002      | NA       | 44.23 | 38.08   | 68.02  | 71.49 | 67.50 | 60.62  | 85.63  | 74.20  | 63.72  |
+| OpenAI-Text-3-Small | NA       | 52.48 | 34.10   | 62.62  | 65.87 | 60.28 | 54.85  | 81.96  | 67.57  | 59.97  |
+| OpenAI-Text-3-Large | NA       | 55.21 | 46.83   | 70.81  | 72.89 | 68.12 | 59.58  | 87.60  | 75.22  | 67.03  |
+| CodeSage-Small      | 130M     | 49.93 | 41.05   | 64.26  | 63.19 | 59.87 | 54.65  | 77.60  | 63.18  | 59.22  |
+| CodeSage-Base       | 356M     | 48.50 | 48.87   | 67.81  | 68.00 | 66.87 | 58.13  | 83.17  | 68.00  | 63.67  |
+| CodeSage-Large      | 1.3B     | 47.49 | 52.35   | 70.64  | 70.20 | 69.54 | 61.31  | 83.60  | 71.88  | 65.88  |
+| CodeSage-v2-Small   | 130M     | 52.39 | 47.28   | 68.79  | 68.13 | 65.77 | 60.20  | 80.26  | 72.46  | 64.41  |
+| CodeSage-v2-Base    | 356M     | 50.74 | 52.00   | 70.46  | 70.89 | 69.61 | 62.81  | 82.37  | 73.71  | 66.57  |
+| CodeSage-v2-Large   | 1.3B     | 53.18 | 56.31   | 74.18  | 72.33 | 72.49 | 65.26  | 84.67  | 76.61  | 69.38  |
+
+
+
+
+
+
+
+### Training Data
+This pretrained checkpoint is the same as those used by our V1 model ([codesage/codesage-small](https://huggingface.co/codesage/codesage-small), which is trained on [The Stack](https://huggingface.co/datasets/bigcode/the-stack-dedup) data. The constative learning data are extracted from [The Stack V2](https://huggingface.co/datasets/bigcode/the-stack-v2). Same as our V1 model, we supported nine languages as follows: c, c-sharp, go, java, javascript, typescript, php, python, ruby.
+
+### How to use
+This checkpoint consists of an encoder (130M model), which can be used to extract code embeddings of 1024 dimension. It can be easily loaded using the AutoModel functionality and employs the [Starcoder Tokenizer](https://arxiv.org/pdf/2305.06161.pdf).
+
+```
+from transformers import AutoModel, AutoTokenizer
+
+checkpoint = "codesage/codesage-base-v2"
+device = "cuda"  # for GPU usage or "cpu" for CPU usage
+
+# Note: CodeSage requires adding eos token at the end of each tokenized sequence 
+
+tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True, add_eos_token=True)
+
+model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True).to(device)
+
+inputs = tokenizer.encode("def print_hello_world():\tprint('Hello World!')", return_tensors="pt").to(device)
+embedding = model(inputs)[0]
+```
+
+### BibTeX entry and citation info
+```
+@inproceedings{
+    zhang2024code,
+    title={{CODE} {REPRESENTATION} {LEARNING} {AT} {SCALE}},
+    author={Dejiao Zhang and Wasi Uddin Ahmad and Ming Tan and Hantian Ding and Ramesh Nallapati and Dan Roth and Xiaofei Ma and Bing Xiang},
+    booktitle={The Twelfth International Conference on Learning Representations},
+    year={2024},
+    url={https://openreview.net/forum?id=vfzRRjumpX}
+}
+```