Compare commits

...

No commits in common. "main" and "refs/deployment/triton" have entirely different histories.

19 changed files with 207 additions and 148183 deletions

35
.gitattributes vendored

@ -1,35 +0,0 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.npy filter=lfs diff=lfs merge=lfs -text
*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.pickle filter=lfs diff=lfs merge=lfs -text
*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text

170
1/model.py Normal file

@ -0,0 +1,170 @@
import os
import torch
import json
import numpy as np
import triton_python_backend_utils as pb_utils
from transformers import AutoModel, AutoTokenizer, GenerationConfig
from peft import PeftModel, PeftConfig
class TritonPythonModel:
def initialize(self, args):
"""
모델이 로드될 번만 호출됩니다.
"""
self.logger = pb_utils.Logger
self.model_config = json.loads(args["model_config"])
self.model_name = args["model_name"]
self.base_model_path = self._get_config_parameter("base_model_path")
self.sub_dir_path = self._get_config_parameter("sub_dir_path")
self.logger.log_info(f"base_model_path: {self.base_model_path}")
self.logger.log_info(f"sub_dir_path: {self.sub_dir_path}") # sub_dir_path 로깅 추가
if self.sub_dir_path and os.path.isdir(self.sub_dir_path):
try:
file_list = os.listdir(self.sub_dir_path)
self.logger.log_info(f"'{self.sub_dir_path}' 경로의 파일 목록:\n{file_list}")
# date-w-locale.txt 파일 내용 로깅
target_file_path = os.path.join(self.sub_dir_path, "date-w-locale.txt")
if os.path.exists(target_file_path):
with open(target_file_path, 'r', encoding='utf-8') as f:
file_content = f.read()
self.logger.log_info(f"'{target_file_path}' 파일 내용:\n{file_content}")
else:
self.logger.log_warn(f"'{target_file_path}' 파일을 찾을 수 없습니다.")
except Exception as e:
self.logger.log_error(f"파일 시스템 접근 중 오류 발생: {e}")
elif self.sub_dir_path:
self.logger.log_warn(f"지정된 경로 '{self.sub_dir_path}'가 유효한 디렉토리가 아니거나 존재하지 않습니다.")
self.load_model()
def load_model(self):
torch_dtype = torch.float16
# AutoModel로 모델 로드
self.model = AutoModel.from_pretrained(
pretrained_model_name_or_path=self.base_model_path,
local_files_only=True,
trust_remote_code=True
)
# 모델을 평가 모드로 설정
self.model.eval()
self.tokenizer = AutoTokenizer.from_pretrained(
self.base_model_path,
trust_remote_code=True,
add_eos_token=True
)
self.logger.log_info(f"'{self.model_name}' 모델 초기화 완료 (Code Embedding Mode)")
def execute(self, requests):
"""
Triton이 추론 요청에 대해 호출하는 실행 함수입니다.
Generation 대신 Embedding 생성을 수행하도록 수정합니다.
"""
responses = []
# 각 추론 요청을 순회하며 처리합니다.
for request in requests:
# Triton 입력 파싱: 텍스트 입력만 처리합니다.
input_text = self._get_input_value(request, "text_input")
# CodeSage는 대화 형식이 아닌 일반 텍스트 (코드)를 입력으로 받으므로
# JSON 파싱 로직과 Chat 템플릿 로직을 제거합니다.
text = input_text
self.logger.log_info(f"입력 text 출력:\n{text}")
# 입력 텍스트를 토큰화합니다.
# add_eos_token=True가 load_model에서 설정되었으므로 토큰화 시 자동으로 추가됩니다.
inputs = self.tokenizer(
text,
return_tensors="pt").to(device=self.model.device)
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]
# # CodeSage는 텍스트 생성이 아닌 임베딩 생성을 수행합니다.
# gened = self.model.generate(...)
# **임베딩 생성**
# CodeSage 모델은 (임베딩, 히든 스테이트, 어텐션)을 반환하며, 첫 번째 요소가 임베딩입니다.
with torch.no_grad():
# inputs에는 input_ids와 attention_mask가 모두 포함되어 전달됩니다.
outputs = self.model(**inputs)
# outputs[0]에는 임베딩 벡터가 포함되어 있습니다.
# 임베딩은 일반적으로 첫 번째 토큰 (CLS 토큰 또는 문맥 임베딩)을 사용합니다.
# CodeSage의 경우, 모델 카드 예시를 보면 outputs[0] 전체를 사용합니다.
# 여기서는 [batch_size, sequence_length, hidden_size] 형태의 임베딩 중 첫 번째 토큰 임베딩을 사용합니다.
# 임베딩 사용법은 모델의 목적에 따라 달라질 수 있습니다. CodeSage는 주로 전체 시퀀스 임베딩을 사용합니다.
# 여기서는 예시와 같이 첫 번째 요소 (last_hidden_state)를 가져옵니다.
# 임베딩 크기: [1, seq_len, hidden_size]
embeddings = outputs[0]
# 임베딩을 NumPy 배열로 변환
# CPU로 옮기고, NumPy로 변환
embedding_np = embeddings.squeeze().cpu().numpy()
# 출력 텐서 생성 (데이터 타입은 float32 또는 float16이 적합)
# CodeSage는 단일 문장 입력만 처리하므로, 배치 차원 없이 [seq_len, hidden_size]로 가정합니다.
# 실제 사용 목적에 따라 풀링 로직을 추가하여 [hidden_size] 벡터로 만들 수도 있습니다.
output_tensor = pb_utils.Tensor("embedding_output", embedding_np.astype(np.float32))
self.logger.log_info(f"모델이 생성한 임베딩 Shape:\n{embedding_np.shape}")
# 응답 객체를 생성하고 출력 텐서를 추가합니다.
responses.append(pb_utils.InferenceResponse(output_tensors=[output_tensor]))
return responses
def _get_config_parameter(self, parameter_name):
"""
모델 설정(config.pbtxt)에서 특정 파라미터의 문자열 값을 가져옵니다.
"""
self.parameters = self.model_config.get('parameters', {})
parameter_dict = self.parameters.get(parameter_name)
if isinstance(parameter_dict, dict) and 'string_value' in parameter_dict:
return parameter_dict['string_value']
return None
def _get_input_value(self, request, input_name: str, default=None):
"""
Triton 추론 요청에서 특정 이름의 입력 텐서 값을 가져옵니다.
"""
tensor_value = pb_utils.get_input_tensor_by_name(request, input_name)
if tensor_value is None:
return default
return self._np_decoder(tensor_value.as_numpy()[0])
def _np_decoder(self, obj):
"""
NumPy 객체의 데이터 타입을 확인하고 Python 기본 타입으로 변환합니다.
"""
if isinstance(obj, bytes):
return obj.decode('utf-8')
if np.issubdtype(obj, np.integer):
return int(obj)
if np.issubdtype(obj, np.floating):
return round(float(obj), 3)
if isinstance(obj, np.bool_):
return bool(obj)
def finalize(self):
"""
모델 실행이 완료된 Triton 서버가 종료될 호출되는 함수입니다.
"""
pass

@ -1,7 +0,0 @@
{
"word_embedding_dimension": 2048,
"pooling_mode_cls_token": false,
"pooling_mode_mean_tokens": true,
"pooling_mode_max_tokens": false,
"pooling_mode_mean_sqrt_len_tokens": false
}

100
README.md

@ -1,100 +0,0 @@
---
license: apache-2.0
datasets:
- bigcode/the-stack-dedup
- bigcode/the-stack-v2
library_name: transformers
language:
- code
---
## CodeSage-Large-v2
### Updates
* 1. <span style="color:blue">We are excited to announce the release of the CodeSage V2 model family with largely improved performance!</span> <strong>Please check out our [blogpost](https://code-representation-learning.github.io/codesage-v2.html) for more details</strong>.
* 2. In addition to retrieval performance boost, V2 models also support flexible embedding sizes (thanks to Matryoshka Representation Learning (MRL)).
* 3. You can access CodeSage v2 models through both HF and SentenceTransformer.
### 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 (1.3B model), which can be used to extract code embeddings of 2048 dimension.
1. 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).
```
from transformers import AutoModel, AutoTokenizer
checkpoint = "codesage/codesage-large-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]
```
2. Accessing CodeSage via SentenceTransformer
```
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("codesage/codesage-large-v2", trust_remote_code=True)
```
### 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}
}
```

@ -1,4 +0,0 @@
{
"<mask>": 49152,
"<pad>": 49153
}

@ -1,25 +0,0 @@
{
"_name_or_path": "codesage/codesage-large-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": 2048,
"num_attention_heads": 16,
"num_hidden_layers": 24,
"intermediate_size": 8192,
"max_position_embeddings": 2048,
"residual_dropout_prob": 0.1,
"vocab_size": 49154
}

37
config.pbtxt Normal file

@ -0,0 +1,37 @@
# Triton Backend for TransformerLLM.
backend: "python"
max_batch_size: 0
# Triton should expect as input a single string
# input of variable length named 'text_input'
input [
{
name: "text_input"
data_type: TYPE_STRING
dims: [ -1 ]
}
]
# output of variable length named 'embedding_outputt'
output [
{
name: "embedding_output"
data_type: TYPE_FP32 # 또는 TYPE_FP16
dims: [ -1, -1 ] # [seq_len, hidden_size]
}
]
parameters: [
{
key: "base_model_path",
value: {string_value: "/cheetah/input/model/groupuser/codesage-large-v2"}
}
]
instance_group [
{
kind: KIND_AUTO
count: 1
}
]

@ -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=768,
num_hidden_layers=12,
num_attention_heads=12,
intermediate_size=3072,
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)

48892
merges.txt

File diff suppressed because it is too large Load Diff

@ -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,
)

@ -1,20 +0,0 @@
[
{
"idx": 0,
"name": "0",
"path": "",
"type": "sentence_transformers.models.Transformer"
},
{
"idx": 1,
"name": "1",
"path": "1_Pooling",
"type": "sentence_transformers.models.Pooling"
},
{
"idx": 2,
"name": "2",
"path": "2_Normalize",
"type": "sentence_transformers.models.Normalize"
}
]

BIN
pytorch_model.bin (Stored with Git LFS)

Binary file not shown.

@ -1,4 +0,0 @@
{
"max_seq_length": 1024,
"do_lower_case": false
}

@ -1,28 +0,0 @@
{
"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|>",
"mask_token": "<mask>",
"pad_token": "<pad>",
"unk_token": "<|endoftext|>"
}

@ -1 +0,0 @@
25-11-18-Asia-Seoul

@ -1,277 +0,0 @@
import json
import os
from functools import lru_cache
from typing import List, Optional, Tuple
import regex as re
from transformers import AddedToken, PreTrainedTokenizer
import logging
logger = logging.getLogger(__name__)
VOCAB_FILES_NAMES = {
"vocab_file": "vocab.json",
"merges_file": "merges.txt",
}
# Taken from
# https://github.com/huggingface/transformers/blob/8aca43bdb3cb9a5020f6d57589d85679dc873b1c/src/transformers/models/gpt2/tokenization_gpt2.py#L62-L84
@lru_cache()
def bytes_to_unicode():
"""
Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
characters the bpe code barfs on.
The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
tables between utf-8 bytes and unicode strings.
"""
bs = (
list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
)
cs = bs[:]
n = 0
for b in range(2**8):
if b not in bs:
bs.append(b)
cs.append(2**8 + n)
n += 1
cs = [chr(n) for n in cs]
return dict(zip(bs, cs))
def get_pairs(word):
"""
Return set of symbol pairs in a word.
Word is represented as tuple of symbols (symbols being variable-length strings).
"""
pairs = set()
prev_char = word[0]
for char in word[1:]:
pairs.add((prev_char, char))
prev_char = char
return pairs
class CodeSageTokenizer(PreTrainedTokenizer):
"""A thin wrapper of the starcoder tokenizer.
See HuggingFace for further documentation on general tokenizer methods.
"""
vocab_files_names = VOCAB_FILES_NAMES
model_input_names = ["input_ids", "attention_mask"]
def __init__(
self,
vocab_file,
merges_file,
errors="replace",
unk_token="<|endoftext|>",
bos_token="<|endoftext|>",
eos_token="<|endoftext|>",
pad_token=None,
add_prefix_space=False,
add_bos_token=False,
add_eos_token=True,
**kwargs,
):
bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
unk_token = AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance(unk_token, str) else unk_token
pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
self.add_bos_token = add_bos_token
self.add_eos_token = add_eos_token
with open(vocab_file, encoding="utf-8") as vocab_handle:
self.encoder = json.load(vocab_handle)
self.decoder = {v: k for k, v in self.encoder.items()}
self.errors = errors # how to handle errors in decoding
self.byte_encoder = bytes_to_unicode()
self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
with open(merges_file, encoding="utf-8") as merges_handle:
bpe_merges = merges_handle.read().split("\n")[1:-1]
bpe_merges = [tuple(merge.split()) for merge in bpe_merges]
self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
self.cache = {}
self.add_prefix_space = add_prefix_space
# Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
super().__init__(
errors=errors,
unk_token=unk_token,
bos_token=bos_token,
eos_token=eos_token,
pad_token=pad_token,
add_prefix_space=add_prefix_space,
add_bos_token=add_bos_token,
add_eos_token=add_eos_token,
**kwargs,
)
@property
def vocab_size(self):
return len(self.encoder)
def get_vocab(self):
return dict(self.encoder, **self.added_tokens_encoder)
def bpe(self, token):
if token in self.cache:
return self.cache[token]
word = tuple(token)
pairs = get_pairs(word)
if not pairs:
return token
while True:
bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
if bigram not in self.bpe_ranks:
break
first, second = bigram
new_word = []
i = 0
while i < len(word):
try:
j = word.index(first, i)
except ValueError:
new_word.extend(word[i:])
break
else:
new_word.extend(word[i:j])
i = j
if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
new_word.append(first + second)
i += 2
else:
new_word.append(word[i])
i += 1
new_word = tuple(new_word)
word = new_word
if len(word) == 1:
break
else:
pairs = get_pairs(word)
word = " ".join(word)
self.cache[token] = word
return word
def build_inputs_with_special_tokens(
self,
token_ids_0: List[int],
token_ids_1: Optional[List[int]] = None) -> List[int]:
bos_token_id = [self.bos_token_id] if self.add_bos_token else []
eos_token_id = [self.eos_token_id] if self.add_eos_token else []
output = bos_token_id + token_ids_0 + eos_token_id
if token_ids_1 is not None:
output = output + bos_token_id + token_ids_1 + eos_token_id
return output
def get_special_tokens_mask(
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
) -> List[int]:
"""
Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.
Args:
token_ids_0 (`List[int]`):
List of IDs.
token_ids_1 (`List[int]`, *optional*):
Optional second list of IDs for sequence pairs.
already_has_special_tokens (`bool`, *optional*, defaults to `False`):
Whether or not the token list is already formatted with special tokens for the model.
Returns:
`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
"""
if already_has_special_tokens:
return super().get_special_tokens_mask(
token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
)
if not self.add_bos_token:
return super().get_special_tokens_mask(
token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=False
)
if token_ids_1 is None:
return [1] + ([0] * len(token_ids_0))
return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1))
def _tokenize(self, text):
"""Tokenize a string."""
bpe_tokens = []
for token in re.findall(self.pat, text):
token = "".join(
self.byte_encoder[b] for b in token.encode("utf-8")
) # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
return bpe_tokens
def _convert_token_to_id(self, token):
"""Converts a token (str) in an id using the vocab."""
return self.encoder.get(token, self.encoder.get(self.unk_token))
def _convert_id_to_token(self, index):
"""Converts an index (integer) in a token (str) using the vocab."""
return self.decoder.get(index)
def convert_tokens_to_string(self, tokens):
"""Converts a sequence of tokens (string) in a single string."""
text = "".join(tokens)
text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
return text
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
if not os.path.isdir(save_directory):
logger.error(f"Vocabulary path ({save_directory}) should be a directory")
return
vocab_file = os.path.join(
save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
)
merge_file = os.path.join(
save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
)
with open(vocab_file, "w", encoding="utf-8") as f:
f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
index = 0
with open(merge_file, "w", encoding="utf-8") as writer:
writer.write("#version: 0.2\n")
for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
if index != token_index:
logger.warning(
f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
" Please check that the tokenizer is not corrupted!"
)
index = token_index
writer.write(" ".join(bpe_tokens) + "\n")
index += 1
return vocab_file, merge_file
def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
add_prefix_space = kwargs.pop("add_prefix_space", self.add_prefix_space)
if is_split_into_words or add_prefix_space:
text = " " + text
return (text, kwargs)
@property
def default_chat_template(self):
"""
A simple chat template that ignores role information and just concatenates messages with EOS tokens.
"""
return "{% for message in messages %}" "{{ message.content }}{{ eos_token }}" "{% endfor %}"

File diff suppressed because it is too large Load Diff

@ -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]
}
}

File diff suppressed because one or more lines are too long