Train code
#Train
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model, PeftModel
from datasets import load_dataset, Dataset
import re
import transformers
model_name = "/mnt/workspace/.cache/modelscope/hub/qwen/Qwen2.5-7B-Instruct"
# model = AutoModelForCausalLM.from_pretrained(
# model_name,
# torch_dtype="auto",
# device_map="auto"
# )
# model_name = "Qwen/Qwen2.5-7B-Instruct"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config, device_map={"":0})
def prompt_query(q):
# q = re.sub(r'请?(理解|阅读).*，?.*(回答|答复|选项|options)(。|.)', "", q).strip()
# if '问题：'in q or 'question:'in q.lower() or '问：'in q:
# # q = q.replace('问：', 'Question:').replace('问题：', 'Question:').replace('question:', 'Question:')
# q = re.sub(r'问：|问题：|question:', 'Question:', q)
# q = re.sub(r'.*?Question:', 'Question:', q, flags=re.DOTALL)
# else:
# q = f'Question: {q}'
q = f'User: {q}\n请回答以上问题。Answer:'
return q
def prompt_answer(a):
if '答：'in a:
a = a.replace('答：', 'Answer:')
else:
a = f'Answer: {a}'
return a
def preprocess_function(examples):
conversations = examples["conversations"]
# input_texts = [one[0]['value'] for one in conversations]
# target_texts = [one[1]['value'] for one in conversations]
input_texts = [prompt_query(one[0]['value']) for one in conversations]
target_texts = [prompt_answer(one[1]['value']) for one in conversations]
# Format inputs for the model
# inputs = [f"User: {inp}\nAssistant:" for inp in input_texts]
model_inputs = tokenizer(input_texts, max_length=512, truncation=True, padding="max_length", return_tensors="pt")
# Tokenize target texts with the 'text_target' argument instead of using 'as_target_tokenizer'
labels = tokenizer(target_texts, max_length=512, truncation=True, padding="max_length", return_tensors="pt")
# Replace padding tokens in labels with -100 to ignore them during training
labels["input_ids"][labels["input_ids"] == tokenizer.pad_token_id] = -100
# Add labels to model inputs
model_inputs["labels"] = labels["input_ids"]
return model_inputs
def filter_objective(examples):
return examples
# keep_indices = []
# conversations = examples["conversations"]
# for idx, one in enumerate(conversations):
# if re.search(r'^(问：|问题：|question:)', one[0]['value']):
# keep_indices.append(idx)
# filtered_examples = {
# "instruction": [examples["instruction"][i] for i in keep_indices],
# "id": [examples["id"][i] for i in keep_indices],
# "conversations": [examples["conversations"][i] for i in keep_indices],
# "lang": [examples["lang"][i] for i in keep_indices]
# }
# return filtered_examples
train = './data/train.jsonl'
val = './data/eval.jsonl'
dataset = load_dataset('json', data_files={'train': train, 'validation': val})
filter_dataset = dataset.map(filter_objective, batched=True, num_proc=1, desc="filter objective")
filter_dataset['tokenized_train'] = filter_dataset['train'].map(preprocess_function, batched=True, num_proc=1, desc="Tokenizing Dataset")
filter_dataset['tokenized_val'] = filter_dataset['validation'].map(preprocess_function, batched=True, num_proc=1, desc="Tokenizing Dataset")
filter_dataset
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)
def print_trainable_parameters(model):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
all_param += param.numel()
if param.requires_grad:
trainable_params += param.numel()
print(
f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
)
config = LoraConfig(
r=8,
lora_alpha=8,
target_modules=["q_proj", "k_proj"],
lora_dropout=0.08,
bias="none",
task_type="CAUSAL_LM"
)
model = get_peft_model(model, config)
# # 微调 lm_head 层
# for param in model.parameters():
# param.requires_grad = False
# # 解冻 `lm_head` 层
# model.lm_head.requires_grad = True
# for name, param in model.named_parameters():
# if 'lora' in name: # LoRA 相关的参数
# param.requires_grad = True
print_trainable_parameters(model)
out = './runs/lora_2epoches'
tokenizer.pad_token = tokenizer.eos_token
trainer = transformers.Trainer(
model=model,
train_dataset=filter_dataset["tokenized_train"],
eval_dataset=filter_dataset['tokenized_val'],
args=transformers.TrainingArguments(
per_device_train_batch_size=8,
gradient_accumulation_steps=2,
warmup_steps=2,
eval_steps=100,
# eval_strategy="epoch",
save_steps=100,
# max_steps=1000,
num_train_epochs=1,
learning_rate=5e-5,
fp16=True,
logging_steps=100,
output_dir=out,
optim="paged_adamw_8bit",
),
data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False # silence the warnings. Please re-enable for inference!
trainer.train()
model.save_pretrained(out)
Inference code
import json
from tqdm import tqdm
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
import re
import transformers
model_name = "qwen/Qwen2.5-7B-Instruct"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config, device_map={"":0})
# lora_model = PeftModel.from_pretrained(model, './runs/lora_10w/checkpoint-8500')
lora_model = PeftModel.from_pretrained(model, './runs/lora_10w/checkpoint-8500/')
lora_model
log_file = './data/eval_only_query.jsonl'
# log_file = './data/eval.jsonl'
query = []
with open(log_file, 'r')as f:
for line in f:
# 每一行是一个 JSON 字符串，可以用 json.loads() 解析
data = json.loads(line.strip()) # strip() 去除换行符
query.append(data)
def prompt_query(q):
# q = re.sub(r'请?(理解|阅读).*，?.*(回答|答复|选项|options)(。|.)', "", q).strip()
# if '问题：'in q or 'question:'in q.lower() or '问：'in q:
# # q = q.replace('问：', 'Question:').replace('问题：', 'Question:').replace('question:', 'Question:')
# # q = re.sub(r'问：|问题：|question:', 'Question:', q)
# # q = re.sub(r'.*?Question:', 'Question:', q, flags=re.DOTALL)
# pass
# else:
# q = f'question: {q}'
# q = f'question: {q}\nAnswer:'
promp = "You are an AI assistant. User will provide a question related to health or medical field. Your task is to answer the user's question as accurately and clearly as possible, using scientific facts, research results, or medical knowledge. Do not provide any medical advice."
q = f'{promp}\nUser: {q} \nAnswer:'
return q
max_new_tokens=512
# prompt = f'User: {q}\nAssistant: '
params = {
'max_new_tokens': max_new_tokens,
'top_p': 0.95,
'temperature': 0.8,
'do_sample': True}
for idx in tqdm(range(len(query))):
# if query[idx]['query_type'] == 'objective':
# # prompt = prompt_query('请阅读并理解问题，然后给出准确的回答。\n' + query[idx]['query'])
# q = query[idx]['query']
# prompt = f'User: {q}\n请回答以上问题。Answer:'
# else:
# prompt = prompt_query(query[idx]['query'])
prompt = prompt_query(query[idx]['query'])
# Tokenize the batch (with padding and truncation)
inputs = tokenizer(prompt, return_tensors='pt').to(lora_model.device)
with torch.no_grad():
outputs = lora_model.generate(**inputs, **params)
prompt_len = len(inputs[0])
output = outputs[0][prompt_len:]
response = tokenizer.decode(output, skip_special_tokens=True)
query[idx]['answer'] = response.strip()
if query[idx]['query_type'] == 'objective':
query[idx]['answer'] = re.sub(r'请阅读并理解问题，然后给出准确的回答.*?$', '', query[idx]['answer'], flags=re.DOTALL).strip()
# if idx == 1:
# break
ans_file = './data/res_lora_10w_8500_ai.jsonl'
with open(ans_file, 'w') as f:
for item in query:
f.write(json.dumps(item, ensure_ascii=False) + '\n')
print(f"Generated responses saved to {ans_file}")