Update README.md
Browse files
README.md
CHANGED
@@ -1,302 +0,0 @@
|
|
1 |
-
---
|
2 |
-
language: code
|
3 |
-
thumbnail: https://cdn-media.huggingface.co/CodeBERTa/CodeBERTa.png
|
4 |
-
datasets:
|
5 |
-
- code_search_net
|
6 |
-
license: apache-2.0
|
7 |
-
base_model: huggingface/CodeBERTa-small-v1
|
8 |
-
---
|
9 |
-
|
10 |
-
# CodeBERTa-language-id: The World’s fanciest programming language identification algo 🤯
|
11 |
-
|
12 |
-
|
13 |
-
To demonstrate the usefulness of our CodeBERTa pretrained model on downstream tasks beyond language modeling, we fine-tune the [`CodeBERTa-small-v1`](https://huggingface.co/huggingface/CodeBERTa-small-v1) checkpoint on the task of classifying a sample of code into the programming language it's written in (*programming language identification*).
|
14 |
-
|
15 |
-
We add a sequence classification head on top of the model.
|
16 |
-
|
17 |
-
On the evaluation dataset, we attain an eval accuracy and F1 > 0.999 which is not surprising given that the task of language identification is relatively easy (see an intuition why, below).
|
18 |
-
|
19 |
-
## Quick start: using the raw model
|
20 |
-
|
21 |
-
```python
|
22 |
-
CODEBERTA_LANGUAGE_ID = "huggingface/CodeBERTa-language-id"
|
23 |
-
|
24 |
-
tokenizer = RobertaTokenizer.from_pretrained(CODEBERTA_LANGUAGE_ID)
|
25 |
-
model = RobertaForSequenceClassification.from_pretrained(CODEBERTA_LANGUAGE_ID)
|
26 |
-
|
27 |
-
input_ids = tokenizer.encode(CODE_TO_IDENTIFY)
|
28 |
-
logits = model(input_ids)[0]
|
29 |
-
|
30 |
-
language_idx = logits.argmax() # index for the resulting label
|
31 |
-
```
|
32 |
-
|
33 |
-
|
34 |
-
## Quick start: using Pipelines 💪
|
35 |
-
|
36 |
-
```python
|
37 |
-
from transformers import TextClassificationPipeline
|
38 |
-
|
39 |
-
pipeline = TextClassificationPipeline(
|
40 |
-
model=RobertaForSequenceClassification.from_pretrained(CODEBERTA_LANGUAGE_ID),
|
41 |
-
tokenizer=RobertaTokenizer.from_pretrained(CODEBERTA_LANGUAGE_ID)
|
42 |
-
)
|
43 |
-
|
44 |
-
pipeline(CODE_TO_IDENTIFY)
|
45 |
-
```
|
46 |
-
|
47 |
-
Let's start with something very easy:
|
48 |
-
|
49 |
-
```python
|
50 |
-
pipeline("""
|
51 |
-
def f(x):
|
52 |
-
return x**2
|
53 |
-
""")
|
54 |
-
# [{'label': 'python', 'score': 0.9999965}]
|
55 |
-
```
|
56 |
-
|
57 |
-
Now let's probe shorter code samples:
|
58 |
-
|
59 |
-
```python
|
60 |
-
pipeline("const foo = 'bar'")
|
61 |
-
# [{'label': 'javascript', 'score': 0.9977546}]
|
62 |
-
```
|
63 |
-
|
64 |
-
What if I remove the `const` token from the assignment?
|
65 |
-
```python
|
66 |
-
pipeline("foo = 'bar'")
|
67 |
-
# [{'label': 'javascript', 'score': 0.7176245}]
|
68 |
-
```
|
69 |
-
|
70 |
-
For some reason, this is still statistically detected as JS code, even though it's also valid Python code. However, if we slightly tweak it:
|
71 |
-
|
72 |
-
```python
|
73 |
-
pipeline("foo = u'bar'")
|
74 |
-
# [{'label': 'python', 'score': 0.7638422}]
|
75 |
-
```
|
76 |
-
This is now detected as Python (Notice the `u` string modifier).
|
77 |
-
|
78 |
-
Okay, enough with the JS and Python domination already! Let's try fancier languages:
|
79 |
-
|
80 |
-
```python
|
81 |
-
pipeline("echo $FOO")
|
82 |
-
# [{'label': 'php', 'score': 0.9995257}]
|
83 |
-
```
|
84 |
-
|
85 |
-
(Yes, I used the word "fancy" to describe PHP 😅)
|
86 |
-
|
87 |
-
```python
|
88 |
-
pipeline("outcome := rand.Intn(6) + 1")
|
89 |
-
# [{'label': 'go', 'score': 0.9936151}]
|
90 |
-
```
|
91 |
-
|
92 |
-
Why is the problem of language identification so easy (with the correct toolkit)? Because code's syntax is rigid, and simple tokens such as `:=` (the assignment operator in Go) are perfect predictors of the underlying language:
|
93 |
-
|
94 |
-
```python
|
95 |
-
pipeline(":=")
|
96 |
-
# [{'label': 'go', 'score': 0.9998052}]
|
97 |
-
```
|
98 |
-
|
99 |
-
By the way, because we trained our own custom tokenizer on the [CodeSearchNet](https://github.blog/2019-09-26-introducing-the-codesearchnet-challenge/) dataset, and it handles streams of bytes in a very generic way, syntactic constructs such `:=` are represented by a single token:
|
100 |
-
|
101 |
-
```python
|
102 |
-
self.tokenizer.encode(" :=", add_special_tokens=False)
|
103 |
-
# [521]
|
104 |
-
```
|
105 |
-
|
106 |
-
<br>
|
107 |
-
|
108 |
-
## Fine-tuning code
|
109 |
-
|
110 |
-
<details>
|
111 |
-
|
112 |
-
```python
|
113 |
-
import gzip
|
114 |
-
import json
|
115 |
-
import logging
|
116 |
-
import os
|
117 |
-
from pathlib import Path
|
118 |
-
from typing import Dict, List, Tuple
|
119 |
-
|
120 |
-
import numpy as np
|
121 |
-
import torch
|
122 |
-
from sklearn.metrics import f1_score
|
123 |
-
from tokenizers.implementations.byte_level_bpe import ByteLevelBPETokenizer
|
124 |
-
from tokenizers.processors import BertProcessing
|
125 |
-
from torch.nn.utils.rnn import pad_sequence
|
126 |
-
from torch.utils.data import DataLoader, Dataset
|
127 |
-
from torch.utils.data.dataset import Dataset
|
128 |
-
from torch.utils.tensorboard.writer import SummaryWriter
|
129 |
-
from tqdm import tqdm, trange
|
130 |
-
|
131 |
-
from transformers import RobertaForSequenceClassification
|
132 |
-
from transformers.data.metrics import acc_and_f1, simple_accuracy
|
133 |
-
|
134 |
-
|
135 |
-
logging.basicConfig(level=logging.INFO)
|
136 |
-
|
137 |
-
|
138 |
-
CODEBERTA_PRETRAINED = "huggingface/CodeBERTa-small-v1"
|
139 |
-
|
140 |
-
LANGUAGES = [
|
141 |
-
"go",
|
142 |
-
"java",
|
143 |
-
"javascript",
|
144 |
-
"php",
|
145 |
-
"python",
|
146 |
-
"ruby",
|
147 |
-
]
|
148 |
-
FILES_PER_LANGUAGE = 1
|
149 |
-
EVALUATE = True
|
150 |
-
|
151 |
-
# Set up tokenizer
|
152 |
-
tokenizer = ByteLevelBPETokenizer("./pretrained/vocab.json", "./pretrained/merges.txt",)
|
153 |
-
tokenizer._tokenizer.post_processor = BertProcessing(
|
154 |
-
("</s>", tokenizer.token_to_id("</s>")), ("<s>", tokenizer.token_to_id("<s>")),
|
155 |
-
)
|
156 |
-
tokenizer.enable_truncation(max_length=512)
|
157 |
-
|
158 |
-
# Set up Tensorboard
|
159 |
-
tb_writer = SummaryWriter()
|
160 |
-
|
161 |
-
|
162 |
-
class CodeSearchNetDataset(Dataset):
|
163 |
-
examples: List[Tuple[List[int], int]]
|
164 |
-
|
165 |
-
def __init__(self, split: str = "train"):
|
166 |
-
"""
|
167 |
-
train | valid | test
|
168 |
-
"""
|
169 |
-
|
170 |
-
self.examples = []
|
171 |
-
|
172 |
-
src_files = []
|
173 |
-
for language in LANGUAGES:
|
174 |
-
src_files += list(
|
175 |
-
Path("../CodeSearchNet/resources/data/").glob(f"{language}/final/jsonl/{split}/*.jsonl.gz")
|
176 |
-
)[:FILES_PER_LANGUAGE]
|
177 |
-
for src_file in src_files:
|
178 |
-
label = src_file.parents[3].name
|
179 |
-
label_idx = LANGUAGES.index(label)
|
180 |
-
print("🔥", src_file, label)
|
181 |
-
lines = []
|
182 |
-
fh = gzip.open(src_file, mode="rt", encoding="utf-8")
|
183 |
-
for line in fh:
|
184 |
-
o = json.loads(line)
|
185 |
-
lines.append(o["code"])
|
186 |
-
examples = [(x.ids, label_idx) for x in tokenizer.encode_batch(lines)]
|
187 |
-
self.examples += examples
|
188 |
-
print("🔥🔥")
|
189 |
-
|
190 |
-
def __len__(self):
|
191 |
-
return len(self.examples)
|
192 |
-
|
193 |
-
def __getitem__(self, i):
|
194 |
-
# We’ll pad at the batch level.
|
195 |
-
return self.examples[i]
|
196 |
-
|
197 |
-
|
198 |
-
model = RobertaForSequenceClassification.from_pretrained(CODEBERTA_PRETRAINED, num_labels=len(LANGUAGES))
|
199 |
-
|
200 |
-
train_dataset = CodeSearchNetDataset(split="train")
|
201 |
-
eval_dataset = CodeSearchNetDataset(split="test")
|
202 |
-
|
203 |
-
|
204 |
-
def collate(examples):
|
205 |
-
input_ids = pad_sequence([torch.tensor(x[0]) for x in examples], batch_first=True, padding_value=1)
|
206 |
-
labels = torch.tensor([x[1] for x in examples])
|
207 |
-
# ^^ uncessary .unsqueeze(-1)
|
208 |
-
return input_ids, labels
|
209 |
-
|
210 |
-
|
211 |
-
train_dataloader = DataLoader(train_dataset, batch_size=256, shuffle=True, collate_fn=collate)
|
212 |
-
|
213 |
-
batch = next(iter(train_dataloader))
|
214 |
-
|
215 |
-
|
216 |
-
model.to("cuda")
|
217 |
-
model.train()
|
218 |
-
for param in model.roberta.parameters():
|
219 |
-
param.requires_grad = False
|
220 |
-
## ^^ Only train final layer.
|
221 |
-
|
222 |
-
print(f"num params:", model.num_parameters())
|
223 |
-
print(f"num trainable params:", model.num_parameters(only_trainable=True))
|
224 |
-
|
225 |
-
|
226 |
-
def evaluate():
|
227 |
-
eval_loss = 0.0
|
228 |
-
nb_eval_steps = 0
|
229 |
-
preds = np.empty((0), dtype=np.int64)
|
230 |
-
out_label_ids = np.empty((0), dtype=np.int64)
|
231 |
-
|
232 |
-
model.eval()
|
233 |
-
|
234 |
-
eval_dataloader = DataLoader(eval_dataset, batch_size=512, collate_fn=collate)
|
235 |
-
for step, (input_ids, labels) in enumerate(tqdm(eval_dataloader, desc="Eval")):
|
236 |
-
with torch.no_grad():
|
237 |
-
outputs = model(input_ids=input_ids.to("cuda"), labels=labels.to("cuda"))
|
238 |
-
loss = outputs[0]
|
239 |
-
logits = outputs[1]
|
240 |
-
eval_loss += loss.mean().item()
|
241 |
-
nb_eval_steps += 1
|
242 |
-
preds = np.append(preds, logits.argmax(dim=1).detach().cpu().numpy(), axis=0)
|
243 |
-
out_label_ids = np.append(out_label_ids, labels.detach().cpu().numpy(), axis=0)
|
244 |
-
eval_loss = eval_loss / nb_eval_steps
|
245 |
-
acc = simple_accuracy(preds, out_label_ids)
|
246 |
-
f1 = f1_score(y_true=out_label_ids, y_pred=preds, average="macro")
|
247 |
-
print("=== Eval: loss ===", eval_loss)
|
248 |
-
print("=== Eval: acc. ===", acc)
|
249 |
-
print("=== Eval: f1 ===", f1)
|
250 |
-
# print(acc_and_f1(preds, out_label_ids))
|
251 |
-
tb_writer.add_scalars("eval", {"loss": eval_loss, "acc": acc, "f1": f1}, global_step)
|
252 |
-
|
253 |
-
|
254 |
-
### Training loop
|
255 |
-
|
256 |
-
global_step = 0
|
257 |
-
train_iterator = trange(0, 4, desc="Epoch")
|
258 |
-
optimizer = torch.optim.AdamW(model.parameters())
|
259 |
-
for _ in train_iterator:
|
260 |
-
epoch_iterator = tqdm(train_dataloader, desc="Iteration")
|
261 |
-
for step, (input_ids, labels) in enumerate(epoch_iterator):
|
262 |
-
optimizer.zero_grad()
|
263 |
-
outputs = model(input_ids=input_ids.to("cuda"), labels=labels.to("cuda"))
|
264 |
-
loss = outputs[0]
|
265 |
-
loss.backward()
|
266 |
-
tb_writer.add_scalar("training_loss", loss.item(), global_step)
|
267 |
-
optimizer.step()
|
268 |
-
global_step += 1
|
269 |
-
if EVALUATE and global_step % 50 == 0:
|
270 |
-
evaluate()
|
271 |
-
model.train()
|
272 |
-
|
273 |
-
|
274 |
-
evaluate()
|
275 |
-
|
276 |
-
os.makedirs("./models/CodeBERT-language-id", exist_ok=True)
|
277 |
-
model.save_pretrained("./models/CodeBERT-language-id")
|
278 |
-
```
|
279 |
-
|
280 |
-
</details>
|
281 |
-
|
282 |
-
<br>
|
283 |
-
|
284 |
-
## CodeSearchNet citation
|
285 |
-
|
286 |
-
<details>
|
287 |
-
|
288 |
-
```bibtex
|
289 |
-
@article{husain_codesearchnet_2019,
|
290 |
-
title = {{CodeSearchNet} {Challenge}: {Evaluating} the {State} of {Semantic} {Code} {Search}},
|
291 |
-
shorttitle = {{CodeSearchNet} {Challenge}},
|
292 |
-
url = {http://arxiv.org/abs/1909.09436},
|
293 |
-
urldate = {2020-03-12},
|
294 |
-
journal = {arXiv:1909.09436 [cs, stat]},
|
295 |
-
author = {Husain, Hamel and Wu, Ho-Hsiang and Gazit, Tiferet and Allamanis, Miltiadis and Brockschmidt, Marc},
|
296 |
-
month = sep,
|
297 |
-
year = {2019},
|
298 |
-
note = {arXiv: 1909.09436},
|
299 |
-
}
|
300 |
-
```
|
301 |
-
|
302 |
-
</details>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|