Cleans up the working directory and leaves models if available. Should not assume any functions from the framework class has been called. Return: None
(
self,
workspace: NNFolderWorkspace,
keep_onnx_model: bool = True,
keep_pytorch_model: bool = True,
)
| 144 | return NetworkModels(torch=torch_models, onnx=onnx_models, trt=None) |
| 145 | |
| 146 | def cleanup( |
| 147 | self, |
| 148 | workspace: NNFolderWorkspace, |
| 149 | keep_onnx_model: bool = True, |
| 150 | keep_pytorch_model: bool = True, |
| 151 | ) -> None: |
| 152 | """ |
| 153 | Cleans up the working directory and leaves models if available. |
| 154 | Should not assume any functions from the framework class has been called. |
| 155 | Return: |
| 156 | None |
| 157 | """ |
| 158 | # Clean-up generated files |
| 159 | if not keep_onnx_model: |
| 160 | if self.onnx_BART_decoder is not None: |
| 161 | self.onnx_BART_decoder.cleanup() |
| 162 | if self.onnx_BART_encoder is not None: |
| 163 | self.onnx_BART_encoder.cleanup() |
| 164 | |
| 165 | if not keep_pytorch_model: |
| 166 | # Using rmtree can be dangerous, have user confirm before deleting. |
| 167 | confirm_folder_delete( |
| 168 | self.torch_BART_dir, |
| 169 | prompt="Confirm you want to delete downloaded pytorch model folder?", |
| 170 | ) |
| 171 | |
| 172 | if not keep_pytorch_model and not keep_onnx_model: |
| 173 | workspace.cleanup(force_remove=False) |
| 174 | |
| 175 | def setup_tokenizer_and_model( |
| 176 | self, |
no test coverage detected