423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-08-23 22:32:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix trainer predict

Browse Source 
Former-commit-id: 99fd9637bdc25f41fd1abc8a162f1069cb9060d4
This commit is contained in:
hoshi-hiyouga 2024-09-02 10:15:29 +08:00 committed by GitHub
parent 5c9972a2d5
commit 5af92971bc
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/llamafactory/train/sft/trainer.py
View File

@ -87,9 +87,10 @@ class CustomSeq2SeqTrainer(Seq2SeqTrainer):
Subclass and override to inject custom behavior.
"""
labels = inputs["labels"].detach().clone().cpu() if "labels" in inputs else None # backup labels (d2h)
labels = inputs["labels"] if "labels" in inputs else None
if self.args.predict_with_generate:
assert self.tokenizer.padding_side == "left", "This method only accepts left-padded tensor."
labels = labels.detach().clone() if labels is not None else None # backup labels
prompt_len, label_len = inputs["input_ids"].size(-1), inputs["labels"].size(-1)
if prompt_len > label_len:
inputs["labels"] = self._pad_tensors_to_target_len(inputs["labels"], inputs["input_ids"])
@ -101,7 +102,7 @@ class CustomSeq2SeqTrainer(Seq2SeqTrainer):
)
if generated_tokens is not None and self.args.predict_with_generate:
generated_tokens[:, :prompt_len] = self.tokenizer.pad_token_id
generated_tokens = generated_tokens.contiguous().cpu() # d2h
generated_tokens = generated_tokens.contiguous()
return loss, generated_tokens, labels