Skip to content

Commit

Permalink
..
Browse files Browse the repository at this point in the history
  • Loading branch information
ShashankMosaicML committed Jan 17, 2024
1 parent d36f5f7 commit ed667db
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions llmfoundry/models/mpt/modeling_mpt.py
Original file line number Diff line number Diff line change
Expand Up @@ -1003,6 +1003,12 @@ def __init__(
def get_targets(self, batch: Mapping) -> torch.Tensor:
targets = torch.roll(batch['labels'], shifts=-1)
targets[:, -1] = -100
# The model should not be trained to predict the word after the eos_token, because it comes from a different sequence.
if self.tokenizer is not None and hasattr(self.tokenizer,
'eos_token_id'):
targets = torch.where(
batch['input_ids'] == self.tokenizer.eos_token_id, -100,
targets)
return targets

def forward(self, batch: MutableMapping) -> CausalLMOutputWithPast:
Expand Down

0 comments on commit ed667db

Please sign in to comment.