简介
公司做个大模型助手，需要提取用户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.0tq1_0 tq2_0 出不来结果，放弃了。
用autoawq进行量化 然后用GPU推断
pip install autoawqfrom 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.3OneKE_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
