-
Benachrichtigungen
You must be signed in to change notification settings - Fork 53
/
doc_builder.py
207 lines (171 loc) · 6.76 KB
/
doc_builder.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import importlib
import os
from abc import ABC, abstractmethod
from typing import List, Tuple
import dask.dataframe as dd
import pandas as pd
from dask import compute, delayed
from nemo_curator.datasets import DocumentDataset
from nemo_curator.utils.distributed_utils import (
read_single_partition,
single_partition_write_with_filename,
)
class DocumentDownloader(ABC):
"""Abstract class for downloading remote data to disk"""
def __init__(self):
super().__init__()
@abstractmethod
def download(self, url):
pass
class DocumentIterator(ABC):
"""
Abstract iterator class for reading in raw records that have been
downloaded to disk
"""
def __init__(self):
super().__init__()
@abstractmethod
def iterate(self, file_path):
pass
class DocumentExtractor(ABC):
"""Abstract class for extracting text from records read from disk"""
def __init__(self):
super().__init__()
@abstractmethod
def extract(self, content):
pass
def import_downloader(downloader_path):
module_path, downloader_name = downloader_path.rsplit(".", 1)
downloader_module = importlib.import_module(module_path)
downloader_class = getattr(downloader_module, downloader_name)
if not issubclass(downloader_class, DocumentDownloader):
raise ValueError(
f"Input downloader {downloader_class.__name__} "
"must be derived from DocumentDownloader defined in "
"nemo_curator.download.docbuilder"
)
return downloader_class
def import_iterator(iterator_path):
module_path, iterator_name = iterator_path.rsplit(".", 1)
iterator_module = importlib.import_module(module_path)
iterator_class = getattr(iterator_module, iterator_name)
if not issubclass(iterator_class, DocumentIterator):
raise ValueError(
f"Input iterator {iterator_class.__name__} "
"must be derived from DocumentIterator "
"defined in nemo_curator.download.docbuilder"
)
return iterator_class
def import_extractor(extractor_path):
module_path, extractor_name = extractor_path.rsplit(".", 1)
extractor_module = importlib.import_module(module_path)
extractor_class = getattr(extractor_module, extractor_name)
if not issubclass(extractor_class, DocumentExtractor):
raise ValueError(
f"Input extractor {extractor_class.__name__} "
"must be derived from DocumentExtractor defined "
"in nemo_curator.download.docbuilder"
)
return extractor_class
def _download_and_extract_single_partition(
paths: List[Tuple[str, str]],
downloader: DocumentDownloader,
iterator: DocumentIterator,
extractor: DocumentExtractor,
output_type: str,
keep_raw_download: bool,
force_download: bool,
) -> pd.DataFrame:
url, output_path = paths
if os.path.exists(output_path) and not force_download:
partition = read_single_partition(
[output_path], backend="pandas", filetype=output_type, add_filename=True
)
return partition
downloaded_file = downloader.download(url)
records = []
# Iterate over all records in file
for item in iterator.iterate(downloaded_file):
record_meta, content = item
# Extract the text from the record
extracted = extractor.extract(content)
if extracted is not None:
text_meta, text = extracted
if text is not None:
line = {
"text": text,
**text_meta,
**record_meta,
}
records.append(line)
partition = pd.DataFrame(records)
filename = os.path.basename(output_path)
output_dir = os.path.dirname(output_path)
partition["filename"] = filename
single_partition_write_with_filename(partition, output_dir, output_type=output_type)
if not keep_raw_download:
os.remove(downloaded_file)
return partition
def download_and_extract(
urls: List[str],
output_paths: List[str],
downloader: DocumentDownloader,
iterator: DocumentIterator,
extractor: DocumentExtractor,
output_format: dict,
output_type: str = "jsonl",
keep_raw_download=False,
force_download=False,
) -> DocumentDataset:
"""
Downloads and extracts a dataset into a format accepted by the NeMo Curator
Args:
urls: A list of urls to download the dataset from
output_paths: A list of paths to save the final extracted output to.
The raw output of the downloader will be saved using the path given by downloader.download(url).
downloader: A DocumentDownloader that handles retrieving each file from its url and saving it to storage
iterator: A DocumentIterator that handles iterating through the downloaded file's format
extractor: A DocumentExtractor that handles extracting the data from its raw format into text
output_format: A dictionary mappings columns to datatypes for the fields of each datapoint after extraction.
output_type: The file type to save the dataset as.
keep_raw_download: Whether to keep the pre-extracted download file.
force_download: If False, will skip processing all files in output_paths that already exist and
directly read from them instead.
Returns:
A DocumentDataset of the downloaded data
"""
if len(urls) != len(output_paths):
raise ValueError("Different number of urls and output_paths")
output_format = dict(sorted(output_format.items()))
df = dd.from_map(
_download_and_extract_single_partition,
zip(urls, output_paths),
downloader=downloader,
iterator=iterator,
extractor=extractor,
output_type=output_type,
keep_raw_download=keep_raw_download,
force_download=force_download,
enforce_metadata=False,
meta=output_format,
)
return DocumentDataset(df)
def batch_download(urls: List[str], downloader: DocumentDownloader) -> List[str]:
"""
Downloads all the urls using the downloader in parallel
"""
delayed_downloads = [delayed(downloader.download)(url) for url in urls]
return compute(*delayed_downloads)