readd remove cols after adding samples + no grad for base chat output

This commit is contained in:
51616 2025-01-09 20:19:14 +00:00
parent a5a64581ca
commit 4dbeac1340
2 changed files with 17 additions and 5 deletions

View file

@ -126,9 +126,19 @@ def get_tokenized_dataset(
ds = ds.map(get_preprocessing_fn(ds_name))
ds = ds.filter(filter_long_samples, batched=True)
if add_negative_prompt:
ds = ds.map(add_negative_prompt_fn, batched=True)
cols_to_remove = [
col
for col in ds.column_names
if col not in ["context", "prompt", "response"]
]
ds = ds.map(add_negative_prompt_fn, batched=True, remove_columns=cols_to_remove)
if add_repeat_prompt and "context_numbers" not in ds_name:
ds = ds.map(add_repeat_prompt_fn, batched=True)
cols_to_remove = [
col
for col in ds.column_names
if col not in ["context", "prompt", "response"]
]
ds = ds.map(add_repeat_prompt_fn, batched=True, remove_columns=cols_to_remove)
# for sft + chat_model, we need to convert the dataset to chat format
# add "messages" field

View file

@ -688,9 +688,11 @@ class ModulatedPretrainedModel(nn.Module):
first_label_pos = torch.argmax((chat_labels != -100).float(), dim=-1)
chat_labels[torch.arange(bs), first_label_pos - 1] = 1
chat_outputs = self.base_model(
input_ids=chat_ids, attention_mask=chat_attn_mask
)
with torch.no_grad():
chat_outputs = self.base_model(
input_ids=chat_ids, attention_mask=chat_attn_mask
)
with apply_generated_loras(
self.base_model,
generated_loras,