Method _load_task_samples

train/src/train_multi_task.py:203–250

Load samples for a specific task type from the unified config.

(self, task_config)

                self.comparison_samples = self._load_task_samples(cfg["compare"])

    def _load_task_samples(self, task_config):
        """Load samples for a specific task type from the unified config."""
        samples = []
        
        if not isinstance(task_config, list):
            task_config = [task_config]
            
        for ds in task_config:
            path = ds.get("json_path")
            strategy = ds.get("sampling_strategy", "all")
            image_root = ds.get("image_root")
            sample = ds.get("sample")  # New field for sample count
            random_sample = ds.get("random", False)  # New field for random sampling

            if path.endswith(".jsonl"):
                data_list = [json.loads(line) for line in open(path, "r")]
            elif path.endswith(".json"):
                data_list = json.load(open(path, "r"))
            else:
                raise ValueError(f"Unsupported file type: {path}")

            # Handle sampling strategy
            count = None
            if ":" in strategy:
                strat, num_s = strategy.split(":")
                if "%" in num_s:
                    count = math.ceil(int(num_s.rstrip("%")) * len(data_list) / 100)
                else:
                    count = int(num_s)
                strategy = strat
            elif sample is not None:
                # Use the new 'sample' field if available
                count = sample
                strategy = "random" if random_sample else "first"

            if strategy == "first" and count is not None:
                data_list = data_list[:count]
            elif strategy == "end" and count is not None:
                data_list = data_list[-count:]
            elif strategy == "random" and count is not None:
                random.shuffle(data_list)
                data_list = data_list[:count]

            for sample in data_list:
                sample['image_root'] = image_root

            samples.extend(data_list)
        return samples

    def _load_samples_from_yaml(self, data_path: str):
        samples = []