Iterate over tar file, yielding filename, content pairs for the given tar stream. :param fileobj: byte stream suitable for tarfile :param meta_names: key of different items in meta file :param skip_meta: regexp for keys that are skipped entirely (Default value = r"__[^/]*__($|/)")
(
fileobj, meta_names, skip_meta=r"__[^/]*__($|/)", suffix=None, handler=reraise_exception, meta_stream=None
)
| 93 | |
| 94 | |
| 95 | def tar_file_iterator_with_meta( |
| 96 | fileobj, meta_names, skip_meta=r"__[^/]*__($|/)", suffix=None, handler=reraise_exception, meta_stream=None |
| 97 | ): |
| 98 | """Iterate over tar file, yielding filename, content pairs for the given tar stream. |
| 99 | |
| 100 | :param fileobj: byte stream suitable for tarfile |
| 101 | :param meta_names: key of different items in meta file |
| 102 | :param skip_meta: regexp for keys that are skipped entirely (Default value = r"__[^/]*__($|/)") |
| 103 | |
| 104 | """ |
| 105 | stream = tarfile.open(fileobj=fileobj, mode="r|*") |
| 106 | data_dir, filename = fileobj.name.rsplit("/", 1) |
| 107 | meta_data = {} # {id: {meta_name: meta_value, meta_name2: meta_value2, ...}} |
| 108 | |
| 109 | if meta_stream is None: |
| 110 | meta_file_name = filename.split(".")[0] + ".meta.jsonl" |
| 111 | meta_path = os.path.join(data_dir, meta_file_name) |
| 112 | if os.path.exists(meta_path): |
| 113 | meta_stream = open(meta_path, "r") |
| 114 | else: |
| 115 | meta_file_name = meta_stream.name |
| 116 | |
| 117 | if meta_stream is not None: |
| 118 | for lineno, line in enumerate(meta_stream): |
| 119 | meta_list = [] |
| 120 | try: |
| 121 | meta_list.append(json.loads(line)) |
| 122 | except Exception as exn: |
| 123 | from sat.helpers import print_rank0 |
| 124 | |
| 125 | print_rank0(f"Error in loading jsonl {meta_file_name}, lineno {lineno}: {line}", level="DEBUG") |
| 126 | continue |
| 127 | for item in meta_list: |
| 128 | if not item["key"] in meta_data: |
| 129 | meta_data[item["key"]] = {} |
| 130 | for meta_name in meta_names: |
| 131 | if meta_name in item: |
| 132 | meta_data[item["key"]][meta_name] = item[meta_name] |
| 133 | meta_stream.close() |
| 134 | |
| 135 | try: |
| 136 | for tarinfo in stream: |
| 137 | fname = tarinfo.name |
| 138 | try: |
| 139 | if not tarinfo.isreg(): |
| 140 | continue |
| 141 | if fname is None: |
| 142 | continue |
| 143 | if "/" not in fname and fname.startswith("__") and fname.endswith("__"): |
| 144 | # skipping metadata for now |
| 145 | continue |
| 146 | if skip_meta is not None and re.match(skip_meta, fname): |
| 147 | continue |
| 148 | if fname.endswith(".txt") and suffix is not None: |
| 149 | data = (stream.extractfile(tarinfo).read().decode() + suffix).encode() |
| 150 | else: |
| 151 | data = stream.extractfile(tarinfo).read() |
| 152 | result = dict(fname=fname, data=data) |
no test coverage detected