Skip to content

Pytorch

PyTorch RAG Finetuning Dataset

PyTorchRAGFinetuningDataset

Bases: Dataset

PyTorch RAG Fine-Tuning Dataset Class.

Parameters:

Name Type Description Default
Dataset _type_

description

required
Source code in src/fed_rag/utils/data/finetuning_datasets/pytorch.py
class PyTorchRAGFinetuningDataset(Dataset):
    """PyTorch RAG Fine-Tuning Dataset Class.

    Args:
        Dataset (_type_): _description_
    """

    def __init__(
        self, input_ids: list[torch.Tensor], target_ids: list[torch.Tensor]
    ):
        self.input_ids = input_ids
        self.target_ids = target_ids

    def __len__(self) -> int:
        return len(self.input_ids)

    def __getitem__(self, idx: int) -> Any:
        return self.input_ids[idx], self.target_ids[idx]