(self, dataset:BaseDataset, resume_path = None)
| 49 | return final_ans, all_messages |
| 50 | |
| 51 | def predict_dataset(self, dataset:BaseDataset, resume_path = None): |
| 52 | samples = dataset.load_data(use_retreival=True) |
| 53 | if resume_path: |
| 54 | assert os.path.exists(resume_path) |
| 55 | with open(resume_path, 'r') as f: |
| 56 | samples = json.load(f) |
| 57 | if self.config.truncate_len: |
| 58 | samples = samples[:self.config.truncate_len] |
| 59 | |
| 60 | sample_no = 0 |
| 61 | for sample in tqdm(samples): |
| 62 | if resume_path and self.config.ans_key in sample: |
| 63 | continue |
| 64 | question, texts, images = dataset.load_sample_retrieval_data(sample) |
| 65 | try: |
| 66 | final_ans, final_messages = self.predict(question, texts, images) |
| 67 | except RuntimeError as e: |
| 68 | print(e) |
| 69 | if "out of memory" in str(e): |
| 70 | torch.cuda.empty_cache() |
| 71 | final_ans, final_messages = None, None |
| 72 | sample[self.config.ans_key] = final_ans |
| 73 | if self.config.save_message: |
| 74 | sample[self.config.ans_key+"_message"] = final_messages |
| 75 | torch.cuda.empty_cache() |
| 76 | self.clean_messages() |
| 77 | |
| 78 | sample_no += 1 |
| 79 | if sample_no % self.config.save_freq == 0: |
| 80 | path = dataset.dump_reults(samples) |
| 81 | print(f"Save {sample_no} results to {path}.") |
| 82 | path = dataset.dump_reults(samples) |
| 83 | print(f"Save final results to {path}.") |
| 84 | |
| 85 | def clean_messages(self): |
| 86 | for agent in self.agents: |
no test coverage detected