Method _save_step_data

knn_utils/saveKNNMulti.py:180–246

Save data for current step using streaming Arrow format

(self, id_cnt, token_id, prob, label)

        return id_cnt_list, token_id_list, prob_list
    
    def _save_step_data(self, id_cnt, token_id, prob, label):
        """Save data for current step using streaming Arrow format"""
        # Convert id_cnt to tensor for gathering
        id_cnt_tensor = torch.tensor(id_cnt, device=self.device)
        
        # Step 1: Gather id_cnt and label first
        id_cnt_gathered = self.accelerator.gather_for_metrics(id_cnt_tensor)
        label_gathered = self.accelerator.gather_for_metrics(label)
        
        # Step 2: Get the largest value of id_cnt from gathered tensor
        max_tokens_global = torch.max(id_cnt_gathered).item()
        
        # Step 3: Pad the original token_id and prob lists to the largest id_cnt
        batch_size = id_cnt_tensor.shape[0]
        token_id_padded = torch.full((batch_size, max_tokens_global), -1, dtype=torch.long, device=self.device)
        prob_padded = torch.zeros((batch_size, max_tokens_global), dtype=torch.float16, device=self.device)
        
        for i in range(batch_size):
            valid_count = id_cnt[i]
            current_token_id = token_id[i]
            current_prob = prob[i]
            
            # For the last batch, there may be list longer than max_tokens_global since there are some remainder. Truncate them !
            if current_token_id.shape[0] > max_tokens_global:
                logger.debug(f"Truncating token_id and prob for batch {i} from {current_token_id.shape[0]} to {max_tokens_global}")
                valid_count = max_tokens_global
                current_token_id = current_token_id[:max_tokens_global]
                current_prob = current_prob[:max_tokens_global]
            
            token_id_padded[i, :valid_count] = current_token_id
            prob_padded[i, :valid_count] = current_prob
        
        # Step 4: Gather token_id_tensor and prob_tensor
        token_id_gathered_padded = self.accelerator.gather_for_metrics(token_id_padded)
        prob_gathered_padded = self.accelerator.gather_for_metrics(prob_padded)
        
        # Step 5: Unpad the result token_id_tensor and prob_tensor to original length list
        total_batch_size = id_cnt_gathered.shape[0]
        token_id_gathered = []
        prob_gathered = []
        
        for i in range(total_batch_size):
            valid_count = id_cnt_gathered[i].item()
            token_id_gathered.append(token_id_gathered_padded[i, :valid_count].cpu().numpy())
            prob_gathered.append(prob_gathered_padded[i, :valid_count].cpu().numpy())
        
        # Only main process writes to file
        if self.process_index == 0:
            id_cnt_np = id_cnt_gathered.cpu().numpy()
            label_np = label_gathered.cpu().numpy()
            
            logger.info(f"id_cnt_np shape: {id_cnt_np.shape} ; token_id_gathered length: {len(token_id_gathered)}; prob_gathered length: {len(prob_gathered)}; label_np shape: {label_np.shape}")
            
            # Create Arrow arrays
            id_cnt_array = pa.array(id_cnt_np, type=pa.int32())
            token_id_array = pa.array(token_id_gathered, type=pa.list_(pa.int32()))
            prob_array = pa.array(prob_gathered, type=pa.list_(pa.float16()))
            label_array = pa.array(label_np, type=pa.int32())