Save the data in a dictionary format cache, and write to a CSV file finally. Typically, the data can be classification predictions, call `save` for single data or call `save_batch` to save a batch of data together, and call `finalize` to write the cached data into CSV file. If no me
| 24 | |
| 25 | |
| 26 | class CSVSaver: |
| 27 | """ |
| 28 | Save the data in a dictionary format cache, and write to a CSV file finally. |
| 29 | Typically, the data can be classification predictions, call `save` for single data |
| 30 | or call `save_batch` to save a batch of data together, and call `finalize` to write |
| 31 | the cached data into CSV file. If no metadata provided, use index from 0 to save data. |
| 32 | Note that this saver can't support multi-processing because it reads / writes single |
| 33 | CSV file and can't guarantee the data order in multi-processing situation. |
| 34 | |
| 35 | """ |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | output_dir: PathLike = "./", |
| 40 | filename: str = "predictions.csv", |
| 41 | overwrite: bool = True, |
| 42 | flush: bool = False, |
| 43 | delimiter: str = ",", |
| 44 | ) -> None: |
| 45 | """ |
| 46 | Args: |
| 47 | output_dir: output CSV file directory. |
| 48 | filename: name of the saved CSV file name. |
| 49 | overwrite: whether to overwriting existing CSV file content, if True, will clear the file before saving. |
| 50 | otherwise, will append new content to the CSV file. |
| 51 | flush: whether to write the cache data to CSV file immediately when `save_batch` and clear the cache. |
| 52 | default to False. |
| 53 | delimiter: the delimiter character in the saved file, default to "," as the default output type is `csv`. |
| 54 | to be consistent with: https://docs.python.org/3/library/csv.html#csv.Dialect.delimiter. |
| 55 | |
| 56 | """ |
| 57 | self.output_dir = Path(output_dir) |
| 58 | self._cache_dict: OrderedDict = OrderedDict() |
| 59 | if not (isinstance(filename, str) and filename[-4:] == ".csv"): |
| 60 | warnings.warn("CSV filename is not a string ends with '.csv'.") |
| 61 | self._filepath = self.output_dir / filename |
| 62 | if self._filepath.exists() and overwrite: |
| 63 | os.remove(self._filepath) |
| 64 | |
| 65 | self.flush = flush |
| 66 | self.delimiter = delimiter |
| 67 | self._data_index = 0 |
| 68 | |
| 69 | def finalize(self) -> None: |
| 70 | """ |
| 71 | Writes the cached dict to a csv |
| 72 | |
| 73 | """ |
| 74 | if not self.output_dir.exists(): |
| 75 | self.output_dir.mkdir(parents=True, exist_ok=True) |
| 76 | with open(self._filepath, "a") as f: |
| 77 | for k, v in self._cache_dict.items(): |
| 78 | f.write(k) |
| 79 | for result in v.flatten(): |
| 80 | f.write(self.delimiter + str(result)) |
| 81 | f.write("\n") |
| 82 | # clear cache content after writing |
| 83 | self.reset_cache() |
no outgoing calls
searching dependent graphs…