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,
)
| 124 | return NetworkModels(torch=torch_models, onnx=onnx_models, trt=None) |
| 125 | |
| 126 | def cleanup( |
| 127 | self, |
| 128 | workspace: NNFolderWorkspace, |
| 129 | keep_onnx_model: bool = True, |
| 130 | keep_pytorch_model: bool = True, |
| 131 | ) -> None: |
| 132 | """ |
| 133 | Cleans up the working directory and leaves models if available. |
| 134 | Should not assume any functions from the framework class has been called. |
| 135 | Return: |
| 136 | None |
| 137 | """ |
| 138 | # Clean-up generated files |
| 139 | if not keep_onnx_model: |
| 140 | if self.onnx_t5_decoder is not None: |
| 141 | self.onnx_t5_decoder.cleanup() |
| 142 | if self.onnx_t5_encoder is not None: |
| 143 | self.onnx_t5_encoder.cleanup() |
| 144 | |
| 145 | if not keep_pytorch_model: |
| 146 | # Using rmtree can be dangerous, have user confirm before deleting. |
| 147 | confirm_folder_delete( |
| 148 | self.torch_t5_dir, |
| 149 | prompt="Confirm you want to delete downloaded pytorch model folder?", |
| 150 | ) |
| 151 | |
| 152 | if not keep_pytorch_model and not keep_onnx_model: |
| 153 | workspace.cleanup(force_remove=False) |
| 154 | |
| 155 | def setup_tokenizer_and_model( |
| 156 | self, |
no test coverage detected