DataCollatorForCompletionOnlyLM is not able to find instruction prompt in utrachat dataset

[anandsarth]

from datasets import load_dataset
dataset_name = "stingning/ultrachat"
dataset = load_dataset(dataset_name, split="train[:1000]")
dataset = dataset.train_test_split(test_size=0.1,seed=42)

model_name = "mistralai/Mistral-7B-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
                             model_name, device_map=device_map,torch_dtype=torch_dtype,)

instruction_template = "<|user|>\n"
response_template = '<|assistant|>\n'
collator = DataCollatorForCompletionOnlyLM(instruction_template=instruction_template, response_template=response_template, tokenizer=tokenizer, mlm=False)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    max_seq_length=seq_length,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    dataset_text_field="text",
    peft_config=peft_config,
    data_collator=collator,
)

trainer.train()

But with collator i get 0 loss and the all the input sample are not able to compute loss

trl=='0.7.4'

[younesbelkada]

Hi @anandsarth

I am not an expert of ultra chat dataset (perhaps @lewtun @edbeeching ) but looking at the dataset it looks like you need to manually format the prompts

You can achieve that with SFTTrainer, please have a look at this section of the docs: https://huggingface.co/docs/trl/sft_trainer#format-your-input-prompts

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.