(self,
dstore_path,
val_path,
index_path,
output_path,
model_path,
k=1024,
knn_temp=1.0,
probe=32,
batch_size=32,
knn_gpu=True,
ignore_first=False,
threshold=1e-10)
| 19 | |
| 20 | class KNNSearchMulti: |
| 21 | def __init__(self, |
| 22 | dstore_path, |
| 23 | val_path, |
| 24 | index_path, |
| 25 | output_path, |
| 26 | model_path, |
| 27 | k=1024, |
| 28 | knn_temp=1.0, |
| 29 | probe=32, |
| 30 | batch_size=32, |
| 31 | knn_gpu=True, |
| 32 | ignore_first=False, |
| 33 | threshold=1e-10): |
| 34 | |
| 35 | self.dstore_path = dstore_path |
| 36 | self.val_path = val_path |
| 37 | self.index_path = index_path |
| 38 | self.output_path = output_path |
| 39 | self.model_path = model_path |
| 40 | self.k = k |
| 41 | self.knn_temp = knn_temp |
| 42 | self.probe = probe |
| 43 | self.batch_size = batch_size |
| 44 | self.knn_gpu = knn_gpu |
| 45 | self.ignore_first = ignore_first |
| 46 | |
| 47 | self.threshold = threshold |
| 48 | |
| 49 | # Initialize accelerator |
| 50 | self.accelerator = Accelerator() |
| 51 | self.device = self.accelerator.device |
| 52 | self.world_size = self.accelerator.num_processes |
| 53 | self.process_index = self.accelerator.local_process_index |
| 54 | |
| 55 | # Get vocab size from tokenizer |
| 56 | tokenizer = AutoTokenizer.from_pretrained(self.model_path) |
| 57 | self.vocab_size = len(tokenizer) |
| 58 | logger.info(f"Vocab size: {self.vocab_size}") |
| 59 | |
| 60 | # Load FAISS index (each process loads it) |
| 61 | self.reconstruct_index, self.index = self._load_faiss_index() |
| 62 | |
| 63 | # Create dataset and dataloader |
| 64 | dataset = Dataset.from_file(self.dstore_path) |
| 65 | # Set format to torch for proper tensor conversion |
| 66 | dataset.set_format(type='torch', columns=['keys', 'vals']) |
| 67 | |
| 68 | if self.val_path is not None: |
| 69 | # Load val_file using pickle |
| 70 | with open(self.val_path, 'rb') as f: |
| 71 | self.vals = pickle.load(f).to(self.device) |
| 72 | else: |
| 73 | self.vals = dataset['vals'].to(self.device) |
| 74 | |
| 75 | self.dataloader = DataLoader( |
| 76 | dataset, |
| 77 | batch_size=self.batch_size, |
| 78 | shuffle=False, |
nothing calls this directly
no test coverage detected