class NestedNextTokenDataset(Dataset):
    def __init__(self, token_ids, target_positions, eval_context_length):
        self.token_ids = torch.tensor(token_ids)
        self.target_positions = target_positions
        self.eval_context_length = eval_context_length

    def __getitem__(self, index):
        target_position = self.target_positions[index]
        start = target_position - self.eval_context_length

        input_ids = self.token_ids[start:target_position]
        target_id = self.token_ids[target_position]

        return input_ids, target_id

        __
        __