从零接触代码模型：第一次让 UniXcoder 读取 C 代码

概览#

该系列文章用于跟踪信安作品的模型训练入门相关

特别鸣谢：G老师、英俊潇洒于学长

环境初步搭建#

使用的是Ubuntu系统，并以此为基础创建虚拟环境

1
mkdir -p ~/projects/vuln-ai
2
cd ~/projects/vuln-ai
3

4
python3 -m venv .venv
5
source .venv/bin/activate

接着升级基础工具

1
pip install -U pip setuptools wheel

然后先安装最基础的一组

1
pip install torch transformers datasets scikit-learn pandas matplotlib

这一组的作用如下

1
torch          → 深度学习框架
2
transformers   → Hugging Face 模型库
3
datasets       → 数据集处理
4
scikit-learn   → 指标计算
5
pandas         → 表格数据处理
6
matplotlib     → 画图

第一次加载UniXcoder tokenizer#

在项目目录新建

1
mkdir -p src
2
nano src/test_tokenizer.py

写入代码

1
from transformers import AutoTokenizer
2

3
model_name = "microsoft/unixcoder-base"
4

5
tokenizer = AutoTokenizer.from_pretrained(model_name)
6

7
code = """
8
void f(char *input)
9
{
10
    char buf[8];
11
    strcpy(buf, input);
12
}
13
"""
14

15
tokens = tokenizer.tokenize(code)
16
token_ids = tokenizer.convert_tokens_to_ids(tokens)
17

18
print("原始代码：")
19
print(code)
20

21
print("\nTokens：")
22
print(tokens)
23

24
print("\nToken IDs：")
25
print(token_ids)
26

27
print("\nToken数量：")
28
print(len(tokens))

保存后运行并查看结果

1
原始代码：
2
void f(char *input)
3
{
4
char buf[8];
5
strcpy(buf, input);
6
}
7

8
Tokens： ['Ċ', 'void', 'Ġf', '(', 'char', 'Ġ*', 'input', ')', 'Ċ', '{', 'Ċ', 'ĠĠĠ', 'Ġchar', 'Ġbuf', '[', '8', '];', 'Ċ', 'ĠĠĠ', 'Ġstrcpy', '(', 'buf', ',', 'Ġinput', ');', 'Ċ', '}', 'Ċ']
9

10
Token IDs： [317, 895, 412, 126, 1285, 426, 1376, 127, 317, 209, 317, 377, 1108, 2200, 177, 142, 823, 317, 377, 16542, 126, 1091, 130, 1586, 388, 317, 211, 317]
11

12
Token数量： 28

不难看出，Unicoder在读这段代码时，是先将原代码按模型自己的子词规则切为若干部分，转化为token，再将token映射为对应的数字来输入并处理

接下来新建py文件并写入

1
nano src/test_encode.py
2

3
from transformers import AutoTokenizer
4

5
model_name = "microsoft/unixcoder-base"
6
tokenizer = AutoTokenizer.from_pretrained(model_name)
7

8
code = """
9
void f(char *input)
10
{
11
    char buf[8];
12
    strcpy(buf, input);
13
}
14
"""
15

16
encoded = tokenizer(
17
    code,
18
    truncation=True,
19
    max_length=512,
20
    padding="max_length",
21
    return_tensors="pt"
22
)
23

24
print("input_ids shape:")
25
print(encoded["input_ids"].shape)
26

27
print("\nattention_mask shape:")
28
print(encoded["attention_mask"].shape)
29

30
print("\n前40个 input_ids:")
31
print(encoded["input_ids"][0][:40])
32

33
print("\n前40个 attention_mask:")
34
print(encoded["attention_mask"][0][:40])

运行后看到

1
input_ids shape:
2
torch.Size([1, 512])
3

4
attention_mask shape:
5
torch.Size([1, 512])
6

7
前40个 input_ids:
8
tensor([ 0, 317, 895, 412, 126, 1285, 426, 1376, 127, 317, 209, 317, 377, 1108, 2200, 177, 142, 823, 317, 377, 16542, 126, 1091, 130, 1586, 388, 317, 211, 317, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
9

10
前40个 attention_mask:
11
tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])

torch.Size 的1 代表一个样本，512 代表每条样本统一整理成 512 个 token 位置，而 attention_mask 后面这么多 0 也是因为原样本只有 28 个 token，所以后面部分基本上都是填充的 padding，0 代表的是无关紧要的 padding，1 则是模型需要认真看的内容，这就是“padding + mask”的配合工作