Update cal_lr.py

This commit is contained in:
hiyouga 2023-11-14 21:14:42 +08:00
parent 5619e76dc5
commit 829e879e04
1 changed files with 2 additions and 7 deletions

View File

@ -32,19 +32,14 @@ def calculate_lr(
dataset=dataset, dataset=dataset,
template="default", template="default",
cutoff_len=cutoff_len, cutoff_len=cutoff_len,
output_dir="dummy_dir", output_dir="dummy_dir"
fp16=True
)) ))
trainset = get_dataset(model_args, data_args) trainset = get_dataset(model_args, data_args)
_, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, is_trainable=False, stage="sft") _, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, is_trainable=False, stage="sft")
trainset = preprocess_dataset(trainset, tokenizer, data_args, training_args, stage="sft") trainset = preprocess_dataset(trainset, tokenizer, data_args, training_args, stage="sft")
data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX) data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX)
dataloader = DataLoader( dataloader = DataLoader(
dataset=trainset, dataset=trainset, batch_size=batch_size, shuffle=True, collate_fn=data_collator, pin_memory=True
batch_size=batch_size,
shuffle=True,
collate_fn=data_collator,
pin_memory=True
) )
valid_tokens, total_tokens = 0, 0 valid_tokens, total_tokens = 0, 0
for batch in tqdm(dataloader): for batch in tqdm(dataloader):