MCPcopy Create free account
hub / github.com/BAAI-DCAI/SpatialBot / DataCollatorForSupervisedDataset

Class DataCollatorForSupervisedDataset

bunny/util/data_utils.py:965–1025  ·  view source on GitHub ↗

Collate examples for supervised fine-tuning.

Source from the content-addressed store, hash-verified

963
964@dataclass
965class DataCollatorForSupervisedDataset(object):
966 """Collate examples for supervised fine-tuning."""
967
968 tokenizer: transformers.PreTrainedTokenizer
969
970 def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
971 input_ids, labels = tuple([instance[key] for instance in instances]
972 for key in ("input_ids", "labels"))
973
974 if self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:
975 for input_id in input_ids:
976 input_id[input_id == self.tokenizer.eos_token_id] = -300
977
978 if conversation_lib.default_conversation.version == "minicpm":
979 for input_id in input_ids:
980 input_id[input_id == self.tokenizer.eos_token_id] = -300
981 input_ids = torch.nn.utils.rnn.pad_sequence(
982 input_ids,
983 batch_first=True,
984 padding_value=self.tokenizer.eos_token_id)
985 input_ids = input_ids[:, :self.tokenizer.model_max_length]
986 attention_mask = input_ids.ne(self.tokenizer.eos_token_id)
987 for input_id in input_ids:
988 input_id[input_id == -300] = self.tokenizer.eos_token_id
989 else:
990 input_ids = torch.nn.utils.rnn.pad_sequence(
991 input_ids,
992 batch_first=True,
993 padding_value=self.tokenizer.pad_token_id)
994 input_ids = input_ids[:, :self.tokenizer.model_max_length]
995 attention_mask = input_ids.ne(self.tokenizer.pad_token_id)
996
997 labels = torch.nn.utils.rnn.pad_sequence(
998 labels,
999 batch_first=True,
1000 padding_value=IGNORE_INDEX)
1001 labels = labels[:, :self.tokenizer.model_max_length]
1002
1003 if self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:
1004 for input_id in input_ids:
1005 input_id[input_id == -300] = self.tokenizer.eos_token_id
1006
1007 batch = dict(
1008 input_ids=input_ids,
1009 labels=labels,
1010 attention_mask=attention_mask,
1011 )
1012
1013 if 'image' in instances[0]:
1014 images = [instance['image'] for instance in instances]
1015 new_images = []
1016 for image in images:
1017 if type(image) is list:
1018 for i in image:
1019 new_images.append(i)
1020 else:
1021 new_images.append(image)
1022 images = new_images

Callers 1

Calls

no outgoing calls

Tested by

no test coverage detected