lmops

1
'''
2
for knn search
3
'''
4
from datasets import load_dataset
5
from typing import Any, Dict, Iterable
6
import torch
7
import pandas as pd
8
import tqdm
9

10
class IndexerDatasetReader(torch.utils.data.Dataset):
11
    def __init__(self, tokenizer, data) -> None:
12
        self.tokenizer = tokenizer
13
        self.dataset=data
14
                 
15
    def __getitem__(self, index):
16
        return self.text_to_instance(self.dataset[index],index=index)
17

18
    def __len__(self):
19
        return len(self.dataset)
20

21
    def text_to_instance(self, entry: Dict[str, Any],index=-1):
22
        enc_text = entry['instruction']
23
        tokenized_inputs = self.tokenizer.encode_plus(enc_text,truncation=True,return_tensors='pt')
24
        return {
25
                        'input_ids': tokenized_inputs.input_ids.squeeze(),
26
                        'attention_mask': tokenized_inputs.attention_mask.squeeze(),
27
                        "metadata":{"id":index}
28
                        
29
                    }
30

31
        
32