def obtain_first_n_samples(self, seed=0):
        """Get first nsample data as the real calibration dataset."""
        self.dataloader.clear()
        random.seed(seed)
        for batch in self.dataloader_original:
            # process data, depends on its data type.
            if len(self.dataloader) == self.nsamples:
                logger.info(
                    f"Successfully collect {self.nsamples} calibration samples."
                )
                break
            # list, tuple
            if isinstance(batch, list) or isinstance(batch, tuple):
                if batch[0].shape[-1] > self.pad_max_length:
                    i = random.randint(0, batch[0].shape[-1] - self.pad_max_length - 1)
                    j = i + self.pad_max_length
                    batch_final = []
                    for item in batch:
                        if isinstance(item, torch.Tensor) and item.shape.__len__() == 2:
                            batch_final.append(item[:, i:j])
                        else:
                            batch_final.append(item)
                else:
                    batch_final = batch[:]
            # dict
            elif isinstance(batch, dict):
                try:
                    length = batch["input_ids"].shape[-1]
                except Exception:
                    logger.warning(
                        "Please make sure your dict'like data contains key of 'input_ids'."
                    )
                    continue
                batch_final = {}
                if length > self.pad_max_length:
                    i = random.randint(0, length - self.pad_max_length - 1)
                    j = i + self.pad_max_length
                    # may have to slice every sequence related data
                    for key in batch.keys():
                        if isinstance(batch[key], torch.Tensor):
                            batch_final[key] = batch[key][
                                :, i:j
                            ]  # slice on sequence length dim
                        else:
                            batch_final[key] = batch[key]
                else:
                    batch_final = batch
            # tensor
            else:
                if batch.shape[-1] > self.pad_max_length:
                    i = random.randint(0, batch.shape[-1] - self.pad_max_length - 1)
                    j = i + self.pad_max_length
                    batch_final = batch[:, i:j]
                else:
                    batch_final = batch
            self.dataloader.append(batch_final)

        if len(self.dataloader) < self.nsamples:
            logger.warning(
                f"Try to use {self.nsamples} data, but entire dataset size is {len(self.dataloader)}."
            )
Get first nsample data as the real calibration dataset.
intel_extension_for_pytorch/quantization/_GPTQ/gptq/gptq.py