Compare commits
No commits in common. "refs/deployment/triton" and "main" have entirely different histories.
refs/deplo
...
main
35
.gitattributes
vendored
Normal file
35
.gitattributes
vendored
Normal file
@ -0,0 +1,35 @@
|
||||
*.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
170
1/model.py
@ -1,170 +0,0 @@
|
||||
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
|
||||
7
1_Pooling/config.json
Normal file
7
1_Pooling/config.json
Normal file
@ -0,0 +1,7 @@
|
||||
{
|
||||
"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
Normal file
100
README.md
Normal file
@ -0,0 +1,100 @@
|
||||
---
|
||||
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}
|
||||
}
|
||||
```
|
||||
4
added_tokens.json
Normal file
4
added_tokens.json
Normal file
@ -0,0 +1,4 @@
|
||||
{
|
||||
"<mask>": 49152,
|
||||
"<pad>": 49153
|
||||
}
|
||||
25
config.json
Normal file
25
config.json
Normal file
@ -0,0 +1,25 @@
|
||||
{
|
||||
"_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
37
config.pbtxt
@ -1,37 +0,0 @@
|
||||
# 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
|
||||
}
|
||||
]
|
||||
52
config_codesage.py
Normal file
52
config_codesage.py
Normal file
@ -0,0 +1,52 @@
|
||||
#!/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
Normal file
48892
merges.txt
Normal file
File diff suppressed because it is too large
Load Diff
426
modeling_codesage.py
Normal file
426
modeling_codesage.py
Normal file
@ -0,0 +1,426 @@
|
||||
#!/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,
|
||||
)
|
||||
20
modules.json
Normal file
20
modules.json
Normal file
@ -0,0 +1,20 @@
|
||||
[
|
||||
{
|
||||
"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)
Normal file
BIN
pytorch_model.bin
(Stored with Git LFS)
Normal file
Binary file not shown.
4
sentence_bert_config.json
Normal file
4
sentence_bert_config.json
Normal file
@ -0,0 +1,4 @@
|
||||
{
|
||||
"max_seq_length": 1024,
|
||||
"do_lower_case": false
|
||||
}
|
||||
28
special_tokens_map.json
Normal file
28
special_tokens_map.json
Normal file
@ -0,0 +1,28 @@
|
||||
{
|
||||
"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
sub-dir/date-w-locale.txt
Normal file
1
sub-dir/date-w-locale.txt
Normal file
@ -0,0 +1 @@
|
||||
25-11-18-Asia-Seoul
|
||||
277
tokenization_codesage.py
Normal file
277
tokenization_codesage.py
Normal file
@ -0,0 +1,277 @@
|
||||
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 %}"
|
||||
98274
tokenizer.json
Normal file
98274
tokenizer.json
Normal file
File diff suppressed because it is too large
Load Diff
34
tokenizer_config.json
Normal file
34
tokenizer_config.json
Normal file
@ -0,0 +1,34 @@
|
||||
{
|
||||
"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]
|
||||
}
|
||||
}
|
||||
1
vocab.json
Normal file
1
vocab.json
Normal file
File diff suppressed because one or more lines are too long
Loading…
Reference in New Issue
Block a user