简介

       公司做个大模型助手，需要提取用户query中的人名、公司名和产品名称来进行问答。目前我使用的是bert+crf模型 开源cluer数据+自造的数据，训练数据18w，测试数据1.3w。

目前这个方案有些瓶颈，主要表现如下：

1、产品名称识别错误 有的时候会把产品名称识别很长

2、产品名称简称识别不了，比如招白，这块数据训练集里面是没有的，训练集里面的产品名称是基金全称、基金简称、以及基金代码

加上目前都是用大模型方案来做，OneKE是阿里联合浙江大学开源的，我打算尝试一下。

ONEKE 魔搭社区

OnKE

我通过官方示例改成我的样本，尝试了几个效果不太行。我就打算进行微调。

首先我的数据集太多了18w，我打算进行过滤。

我使用Dmeta-embedding-zh-small向量Embedding模型，计算所有句子（训练集+测试集）的Embedding，然后使用以下代码进行过滤。

import pickle
import numpy as np
from numpy.linalg import norm

embedding_dict = pickle.load(open('embedding_dict.pkl', 'rb'))

texts = []
embedding = []
for k, v in embedding_dict.items():
    texts.append(k)
    embedding.append(v)

need_set = set()
not_need_set = set()
A = np.array(embedding)

for i in range(len(embedding)):
    if i in not_need_set:
        continue
    B = embedding[i]
    cosine = np.dot(A, B)
    need_set.add(i)
    for index in np.where(cosine > 0.80)[0]:
        not_need_set.add(index)
    if i %1000==0:
        print(i,len(not_need_set),len(need_set))

print(len(need_set))
print(len(not_need_set))

texts_set=set()
for i,text in enumerate(texts):
    if i not in need_set:
        continue
    texts_set.add(text)

print(len(texts_set))
pickle.dump(texts_set, open('texts_set.pkl', 'wb'))

只取 在texts_set里面的句子作为训练集

最后得到train 14498  dev 6214 ( train_size=0.7，来切割原始训练集的）

微调

数据格式按照 DeepKE/example/llm/OneKE.md at main · zjunlp/DeepKE · GitHub  数据准备 部分进行准备

{"text": "相比之下，青岛海牛队和广州松日队的雨中之战虽然也是0∶0，但乏善可陈。", "entity": [{"entity": "广州松日队", "entity_type": "组织机构"}, {"entity": "青岛海牛队", "entity_type": "组织机构"}]}

然后使用转换脚本进行转换 训练数据

python ie2instruction/convert_func.py \
    --src_path data/NER/sample.json \
    --tgt_path data/NER/train.json \
    --schema_path data/NER/schema.json \
    --language zh \
    --task NER \
    --split_num 6 \       
    --random_sort \
    --split train

测试数据

python ie2instruction/convert_func.py \
    --src_path data/NER/sample.json \
    --tgt_path data/NER/test.json \
    --schema_path data/NER/schema.json \
    --language zh \
    --task NER \
    --split_num 6 \
    --split test

我是在OneKE模型上微调的，我参照github上的

output_dir='lora/oneke-continue'
mkdir -p ${output_dir}
CUDA_VISIBLE_DEVICES="0,1,2,3" torchrun --nproc_per_node=4 --master_port=1287 src/finetune.py \
    --do_train --do_eval \
    --overwrite_output_dir \
    --model_name_or_path 'models/OneKE' \
    --stage 'sft' \
    --model_name 'llama' \
    --template 'llama2_zh' \
    --train_file 'data/train.json' \
    --valid_file 'data/dev.json' \
    --output_dir=${output_dir} \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 2 \
    --gradient_accumulation_steps 4 \
    --preprocessing_num_workers 16 \
    --num_train_epochs 10 \
    --learning_rate 5e-5 \
    --max_grad_norm 0.5 \
    --optim "adamw_torch" \
    --max_source_length 400 \
    --cutoff_len 700 \
    --max_target_length 300 \
    --evaluation_strategy "epoch" \
    --save_strategy "epoch" \
    --save_total_limit 10 \
    --lora_r 64 \
    --lora_alpha 64 \
    --lora_dropout 0.05 \
    --bf16 \
    --bits 4

将example/llm/InstructKGC/ft_scripts/fine_continue_full.bash改成了

output_dir='lora/oneke-continue'
mkdir -p ${output_dir}
CUDA_VISIBLE_DEVICES="0,1,2,3" torchrun --nproc_per_node=4 --master_port=1287 src/finetune.py \
    --do_train --do_eval \
    --overwrite_output_dir \
    --model_name_or_path '/data/lifengxin/OneKE' \
    --stage 'sft' \
    --model_name 'llama' \
    --template 'llama2_zh' \
    --train_file 'data/NER/train.json' \
    --valid_file 'data/NER/dev.json' \
    --output_dir=${output_dir} \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 2 \
    --gradient_accumulation_steps 4 \
    --preprocessing_num_workers 16 \
    --num_train_epochs 10 \
    --learning_rate 5e-5 \
    --max_grad_norm 0.5 \
    --optim "adamw_torch" \
    --max_source_length 400 \
    --cutoff_len 700 \
    --max_target_length 300 \
    --evaluation_strategy "epoch" \
    --save_strategy "epoch" \
    --save_total_limit 10 \
    --lora_r 16 \
    --lora_alpha 32 \
    --lora_dropout 0.05 \
    --bf16 \
    --bits 4

然后进行训练

训练的镜像是egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:24.07  从阿里云镜像仓库里面拉的。

训练完成

-rw-r--r-- 1 root root 5.0K Sep  6 17:16 README.md
-rw-r--r-- 1 root root  726 Sep  6 17:16 adapter_config.json
-rw-r--r-- 1 root root 239M Sep  6 17:16 adapter_model.safetensors
-rw-r--r-- 1 root root  371 Sep  6 17:18 all_results.json
drwxr-xr-x 2 root root 4.0K Sep  6 15:28 checkpoint-1359
drwxr-xr-x 2 root root 4.0K Sep  6 15:44 checkpoint-1813
drwxr-xr-x 2 root root 4.0K Sep  6 15:59 checkpoint-2266
drwxr-xr-x 2 root root 4.0K Sep  6 16:15 checkpoint-2719
drwxr-xr-x 2 root root 4.0K Sep  6 16:30 checkpoint-3172
drwxr-xr-x 2 root root 4.0K Sep  6 16:45 checkpoint-3626
drwxr-xr-x 2 root root 4.0K Sep  6 17:01 checkpoint-4079
drwxr-xr-x 2 root root 4.0K Sep  6 14:58 checkpoint-453
drwxr-xr-x 2 root root 4.0K Sep  6 17:14 checkpoint-4530
drwxr-xr-x 2 root root 4.0K Sep  6 15:13 checkpoint-906
-rw-r--r-- 1 root root  179 Sep  6 17:18 eval_results.json
-rw-r--r-- 1 root root  226 Sep  6 17:16 train_results.json
-rw-r--r-- 1 root root 388K Sep  6 17:16 trainer_state.json
-rw-r--r-- 1 root root 5.3K Sep  6 17:16 training_args.bin

进行合并

 python src/export_model.py     --model_name_or_path '/data/OneKE'     --checkpoint_dir '/data/lora/oneke-continue/checkpoint-4530'     --export_dir 'data/OneKE_v1'     --stage 'sft'     --model_name 'llama'     --template 'llama2_zh'     --output_dir 'data/OneKE_v1_test'

合并后进行效果验证

from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    GenerationConfig,
    BitsAndBytesConfig
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_path = '/data/OneKE_v1/'
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)


# 4bit量化OneKE
quantization_config=BitsAndBytesConfig(
    load_in_4bit=True,
    llm_int8_threshold=6.0,
    llm_int8_has_fp16_weight=False,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
)

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    config=config,
    device_map="auto",
    quantization_config=quantization_config,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)
model.eval()


system_prompt = '<<SYS>>\nYou are a helpful assistant. 你是一个乐于助人的助手。\n<</SYS>>\n\n'

sintruct="{\"instruction\": \"你是专门进行实体抽取的专家。请从input中抽取出符合schema定义的实体，不存在的实体类型返回空列表。请按照JSON字符串的格式回答。\", \"schema\": [\"人名\",
\"公司名称\", \"产品名称\"], \"input\": \"164105的投资风格是什么？\"}"


sintruct = '[INST] ' + system_prompt + sintruct + '[/INST]'

input_ids = tokenizer.encode(sintruct, return_tensors="pt").to(device)
input_length = input_ids.size(1)
generation_output = model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_length=1024, max_new_tokens=512, return_dict_in_generate=True), pad_token_id=tokenize
r.eos_token_id)
generation_output = generation_output.sequences[0]
generation_output = generation_output[input_length:]
output = tokenizer.decode(generation_output, skip_special_tokens=True)

print(output)

量化

由于模型占用内存较大，因此产生了量化的想法。

使用llama.cpp进行量化

python llama.cpp/convert.py /data/OneKE_v1   --outfile OneKE_v1.gguf   --outtype q8_0

使用llama-server启动服务进行测试

  ./llama-server  -m OneKE_v1.gguf --port 80 --host 0.0.0.0

效果还可以，需要8s左右推断完

想快速推断完 尝试 tq1_0 tq2_0进行压缩

python3 convert_hf_to_gguf.py /data/OneKE_v1   --outfile OneKE_v1_tq1_0.gguf   --outtype tq1_0
python3 convert_hf_to_gguf.py /data/OneKE_v1   --outfile OneKE_v1_tq2_0.gguf   --outtype tq2_0

-rw-r--r--  1 root root  14G Sep  7 03:24 OneKE_v1.gguf
-rw-r--r--  1 root root 3.6G Sep  7 03:40 OneKE_v1_tq1_0.gguf
-rw-r--r--  1 root root 4.2G Sep  7 03:47 OneKE_v1_tq2_0.gguf

./llama-server  -m OneKE_v1_tq1_0.gguf --port 80  --host 0.0.0.0
./llama-server  -m OneKE_v1_tq2_0.gguf --port 80  --host 0.0.0.0

tq1_0 tq2_0 出不来结果，放弃了。

用autoawq进行量化 然后用GPU推断

pip install autoawq

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = '/data/OneKE_v1'
quant_path = '/data/OneKE_v1-awq'
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }

# Load model
model = AutoAWQForCausalLM.from_pretrained(model_path, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize
model.quantize(tokenizer, quant_config=quant_config)

# Save quantized model
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)
~

OneKE_v1-awq 7.2G 

OneKE_v1 25G

发布

使用vllm进行发布

镜像使用 egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/llm-inference:vllm0.5.4-deepgpu-llm24.7-pytorch2.4.0-cuda12.4-ubuntu22.04  阿里云的

OneKE_v1-awq  占用显卡15G 推断500ms

python3 -m vllm.entrypoints.openai.api_server --model /data/OneKE_v1-awq --host 0.0.0.0 --port 80 --max-model-len 1024 --gpu-memory-utilization 0.3

OneKE_v1 模型 占用显卡 30G，推断1317ms

python3 -m vllm.entrypoints.openai.api_server --model /data/OneKE_v1-awq --host 0.0.0.0 --port 80 --max-model-len 1024 --gpu-memory-utilization 0.5