Function knn_collate_fn

train_memdec.py:428–474

Custom collate function that handles kNN data processing directly Args: batch: The batch of data knn_dstore: The KNN datastore vocab_size: Size of the vocabulary device: The device to move tensors to (accelerator.device)

(batch, knn_dstore, vocab_size)

    knn_dstore.set_format(type='torch', columns=['id_cnt', 'token_id', 'prob', 'label'])
    
    def knn_collate_fn(batch, knn_dstore, vocab_size):
        """
        Custom collate function that handles kNN data processing directly
        
        Args:
            batch: The batch of data
            knn_dstore: The KNN datastore
            vocab_size: Size of the vocabulary
            device: The device to move tensors to (accelerator.device)
        """
        # Apply default collation to the batch
        collated_batch = default_data_collator(batch)
        
        # Process kNN data for all items in the batch
        knn_labels_list = []
        knn_probs_list = []
        
        # Process each dstore_range in the batch
        for idx, cur_range in enumerate(collated_batch["dstore_range"]):
            # Get range boundaries
            start, end = int(cur_range[0]), int(cur_range[1])
            
            # Slice the knn_dstore
            knn_dstore_slice = knn_dstore.select(range(start, end))
            
            # Extract kNN label
            # Note that since datasets version 4.0.0, we can't use direct column selecting since the implementation of lazy columns, see pr https://github.com/huggingface/datasets/pull/7614
            cur_knn_label = knn_dstore_slice[:]["label"]
            knn_labels_list.append(cur_knn_label)
            
            # Extract token IDs and probabilities
            cur_token_id = knn_dstore_slice[:]["token_id"]
            cur_prob = knn_dstore_slice[:]["prob"]
            
            # Create sparse probability tensor
            cur_knn_prob = torch.zeros(size=(end - start, vocab_size))
            for i in range(end - start):
                assert cur_token_id[i].max() < vocab_size, f"token_id {cur_token_id[i]} is out of vocab size {vocab_size}"
                cur_knn_prob[i][cur_token_id[i]] = cur_prob[i]
            
            knn_probs_list.append(cur_knn_prob)
        
        # Concatenate and move to device
        collated_batch["knn_label"] = torch.cat(knn_labels_list, dim=0)
        collated_batch["knn_probs"] = torch.cat(knn_probs_list, dim=0)
        
        return collated_batch
    
    # --------------------------------------------------Evaluation-----------------------------------------------------------
    if args.do_test: