Main processing loop
(self)
| 246 | self.accelerator.wait_for_everyone() |
| 247 | |
| 248 | def process(self): |
| 249 | """Main processing loop""" |
| 250 | logger.info(f"Process {self.process_index}: Starting kNN search and processing") |
| 251 | |
| 252 | for batch_idx, batch in enumerate(tqdm(self.dataloader, desc=f"Process {self.process_index}")): |
| 253 | keys = batch['keys'].to(torch.float16).to(self.device) |
| 254 | vals = batch['vals'].to(torch.int32).to(self.device) |
| 255 | |
| 256 | # Perform kNN search |
| 257 | dists, knns = self.get_knns(keys, self.ignore_first) |
| 258 | neg_dists = -dists |
| 259 | |
| 260 | # Compute probability distribution |
| 261 | knn_probs = self.knns_to_probs(knns, neg_dists) |
| 262 | |
| 263 | # Sparsify distribution |
| 264 | id_cnt, token_id, prob = self.sparsify_distribution(knn_probs) |
| 265 | |
| 266 | # Save step data |
| 267 | self._save_step_data(id_cnt, token_id, prob, vals) |
| 268 | |
| 269 | # Close Arrow writer on main process |
| 270 | if self.process_index == 0: |
| 271 | self.arrow_writer.close() |
| 272 | self.arrow_file.close() |
| 273 | logger.info(f"Finished writing to {self.output_path}") |
| 274 | |
| 275 | def parse_args(): |
| 276 | import argparse |
no test coverage detected