Skip to content

Commit 570b2b7

Browse files
sarahyuricksarahyurick
authored
Update Wikipedia URL generation and download (#2183)
* Update Wikipedia URL generation and download Signed-off-by: Sarah Yurick <sarahyurick@gmail.com> * ruff Signed-off-by: Sarah Yurick <sarahyurick@gmail.com> * add greptile feedback Signed-off-by: Sarah Yurick <sarahyurick@gmail.com> * greptile and ruff Signed-off-by: Sarah Yurick <sarahyurick@gmail.com> * more greptile Signed-off-by: Sarah Yurick <sarahyurick@gmail.com> * greptile again Signed-off-by: Sarah Yurick <sarahyurick@gmail.com> * greptile and slimmer tests Signed-off-by: Sarah Yurick <sarahyurick@gmail.com> --------- Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
1 parent 44d37e0 commit 570b2b7

5 files changed

Lines changed: 193 additions & 35 deletions

File tree

‎nemo_curator/stages/text/download/wikipedia/constants.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from nemo_curator.package_info import __repository_url__, __shortversion__
16+
17+
WIKIMEDIA_USER_AGENT = f"NeMo-Curator-WikipediaDownloader/{__shortversion__} ({__repository_url__})"
18+
WIKIMEDIA_REQUEST_HEADERS = {"User-Agent": WIKIMEDIA_USER_AGENT}

‎nemo_curator/stages/text/download/wikipedia/download.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.
@@ -19,6 +19,8 @@
1919

2020
from nemo_curator.stages.text.download import DocumentDownloader
2121

22+
from .constants import WIKIMEDIA_USER_AGENT
23+
2224

2325
class WikipediaDownloader(DocumentDownloader):
2426
"""Downloads Wikipedia dump files (.bz2) from wikimedia.org."""
@@ -53,7 +55,7 @@ def _download_to_path(self, url: str, path: str) -> tuple[bool, str | None]:
5355
logger.info(f"Downloading {url} to {path}")
5456

5557
# Download with wget
56-
cmd = ["wget", url, "-O", path]
58+
cmd = ["wget", f"--user-agent={WIKIMEDIA_USER_AGENT}", url, "-O", path]
5759

5860
# Set up stdout/stderr handling
5961
if self._verbose:

‎nemo_curator/stages/text/download/wikipedia/url_generation.py‎

Lines changed: 64 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.
@@ -22,6 +22,8 @@
2222

2323
from nemo_curator.stages.text.download import URLGenerator
2424

25+
from .constants import WIKIMEDIA_REQUEST_HEADERS
26+
2527
# Request timeout in seconds
2628
REQUEST_TIMEOUT = 30
2729

@@ -47,15 +49,23 @@ def _get_data_for_dump(self, dump_date: str, wiki_index_url: str) -> dict | None
4749
wiki_latest_dump = urljoin(wiki_index_url + "/", dump_date)
4850
wiki_latest_dump_status = urljoin(wiki_latest_dump, "dumpstatus.json")
4951

50-
raw_dump_data = requests.get(wiki_latest_dump_status, timeout=REQUEST_TIMEOUT)
52+
raw_dump_data = requests.get(
53+
wiki_latest_dump_status,
54+
headers=WIKIMEDIA_REQUEST_HEADERS,
55+
timeout=REQUEST_TIMEOUT,
56+
)
57+
if raw_dump_data.status_code == requests.codes.not_found:
58+
logger.warning(f"Dump status is not available at {wiki_latest_dump_status}")
59+
return None
60+
raw_dump_data.raise_for_status()
5161
try:
5262
dump_data = json.loads(raw_dump_data.content)
53-
except json.JSONDecodeError as e:
63+
except (json.JSONDecodeError, UnicodeDecodeError) as e:
5464
logger.warning(f"Unable to load dump data for {wiki_latest_dump_status}: {e}")
5565
return None
5666
return dump_data
5767

58-
def _get_wikipedia_urls(self) -> list[str]:
68+
def _get_wikipedia_urls(self) -> list[str]: # noqa: C901, PLR0912, PLR0915
5969
"""
6070
Retrieves all URLs pointing to Wikipedia dumps for the specified language and date.
6171
@@ -68,31 +78,64 @@ def _get_wikipedia_urls(self) -> list[str]:
6878
if not dump_date:
6979
# Get the latest dump date from the index
7080
logger.info(f"Fetching latest dump date from {wiki_index_url}")
71-
raw_wiki_index = requests.get(wiki_index_url, timeout=REQUEST_TIMEOUT)
81+
raw_wiki_index = requests.get(
82+
wiki_index_url,
83+
headers=WIKIMEDIA_REQUEST_HEADERS,
84+
timeout=REQUEST_TIMEOUT,
85+
)
86+
raw_wiki_index.raise_for_status()
7287
wiki_index = raw_wiki_index.content.decode("utf-8")
7388
wiki_index_parsed = BeautifulSoup(wiki_index, "lxml")
7489

75-
# Get all dumps available in the index
76-
dumps = wiki_index_parsed.find_all("a")
77-
for dump in reversed(dumps[:-1]):
78-
if dump.text.strip("/").isdigit():
79-
candidate_dump_date = dump.text
80-
dump_data = self._get_data_for_dump(candidate_dump_date, wiki_index_url)
81-
if dump_data is None:
82-
logger.warning(f"Cannot load dump data for {candidate_dump_date[:-1]}")
83-
continue
84-
85-
if dump_data["jobs"].get("articlesmultistreamdump", {}).get("status") == "done":
86-
dump_date = candidate_dump_date
87-
break
88-
else:
89-
logger.warning(f"Dump {candidate_dump_date[:-1]} is not finished, trying next dump")
90+
# Extract and sort date directories instead of relying on link order or
91+
# assuming that "latest/" is the final link in the index.
92+
dump_dates = {
93+
link.text.strip("/")
94+
for link in wiki_index_parsed.find_all("a")
95+
if len(link.text.strip("/")) == 8 and link.text.strip("/").isdigit() # noqa: PLR2004
96+
}
97+
dump_data = None
98+
for candidate_date in sorted(dump_dates, reverse=True):
99+
candidate_dump_date = f"{candidate_date}/"
100+
try:
101+
candidate_dump_data = self._get_data_for_dump(candidate_dump_date, wiki_index_url)
102+
except requests.HTTPError as e:
103+
status_code = e.response.status_code if e.response is not None else None
104+
if status_code is not None and (
105+
status_code in (requests.codes.request_timeout, requests.codes.too_many_requests)
106+
or status_code >= requests.codes.internal_server_error
107+
):
108+
logger.warning(
109+
f"Unable to load dump data for {candidate_date} due to HTTP {status_code}; "
110+
"trying next dump"
111+
)
90112
continue
113+
raise
114+
except requests.RequestException as e:
115+
logger.warning(
116+
f"Unable to load dump data for {candidate_date} due to {type(e).__name__}: {e}; "
117+
"trying next dump"
118+
)
119+
continue
120+
if candidate_dump_data is None:
121+
logger.warning(f"Cannot load dump data for {candidate_date}")
122+
continue
123+
124+
if candidate_dump_data.get("jobs", {}).get("articlesmultistreamdump", {}).get("status") == "done":
125+
dump_date = candidate_dump_date
126+
dump_data = candidate_dump_data
127+
break
128+
129+
logger.warning(f"Dump {candidate_date} is not finished, trying next dump")
130+
131+
if dump_date is None or dump_data is None:
132+
error_msg = f"Unable to find a completed Wikipedia dump at {wiki_index_url}"
133+
raise ValueError(error_msg)
91134

92135
logger.info(f"Found latest dump date: {dump_date[:-1]}")
93136
else:
94137
# A trailing / is needed for the URL
95-
dump_date = dump_date + "/"
138+
dump_date = dump_date.rstrip("/") + "/"
96139
dump_data = self._get_data_for_dump(dump_date, wiki_index_url)
97140
if dump_data is None:
98141
error_msg = f"Unable to load dump data for {dump_date[:-1]}"

‎tests/stages/text/download/wikipedia/test_download.py‎

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.
@@ -16,9 +16,14 @@
1616
from pathlib import Path
1717
from unittest import mock
1818

19+
from nemo_curator.stages.text.download.wikipedia.constants import WIKIMEDIA_USER_AGENT
1920
from nemo_curator.stages.text.download.wikipedia.download import WikipediaDownloader
2021

2122

23+
def _wget_command(url: str, path: str) -> list[str]:
24+
return ["wget", f"--user-agent={WIKIMEDIA_USER_AGENT}", url, "-O", path]
25+
26+
2227
class TestWikipediaDownloader:
2328
"""Test suite for WikipediaDownloader."""
2429

@@ -77,7 +82,7 @@ def test_download_to_path_success(self, mock_run: mock.Mock, tmp_path: Path):
7782
assert success is True
7883
assert error_message is None
7984
mock_run.assert_called_once_with(
80-
["wget", url, "-O", temp_path],
85+
_wget_command(url, temp_path),
8186
stdout=subprocess.DEVNULL,
8287
stderr=subprocess.PIPE,
8388
)
@@ -95,7 +100,7 @@ def test_download_to_path_verbose_logging(self, mock_run: mock.Mock, tmp_path: P
95100
assert success is True
96101
assert error_message is None
97102
mock_run.assert_called_once_with(
98-
["wget", url, "-O", temp_path],
103+
_wget_command(url, temp_path),
99104
stdout=None,
100105
stderr=None,
101106
)
@@ -113,7 +118,7 @@ def test_download_to_path_quiet_mode(self, mock_run: mock.Mock, tmp_path: Path):
113118
assert success is True
114119
assert error_message is None
115120
mock_run.assert_called_once_with(
116-
["wget", url, "-O", temp_path],
121+
_wget_command(url, temp_path),
117122
stdout=subprocess.DEVNULL,
118123
stderr=subprocess.PIPE,
119124
)
@@ -131,7 +136,7 @@ def test_download_to_path_failed(self, mock_run: mock.Mock, tmp_path: Path):
131136
assert success is False
132137
assert error_message == "File not found"
133138
mock_run.assert_called_once_with(
134-
["wget", url, "-O", temp_path],
139+
_wget_command(url, temp_path),
135140
stdout=subprocess.DEVNULL,
136141
stderr=subprocess.PIPE,
137142
)
@@ -149,7 +154,7 @@ def test_download_to_path_network_error(self, mock_run: mock.Mock, tmp_path: Pat
149154
assert success is False
150155
assert error_message == "Network error"
151156
mock_run.assert_called_once_with(
152-
["wget", url, "-O", temp_path],
157+
_wget_command(url, temp_path),
153158
stdout=subprocess.DEVNULL,
154159
stderr=subprocess.PIPE,
155160
)
@@ -167,7 +172,7 @@ def test_download_to_path_failed_no_stderr(self, mock_run: mock.Mock, tmp_path:
167172
assert success is False
168173
assert error_message == "Unknown error"
169174
mock_run.assert_called_once_with(
170-
["wget", url, "-O", temp_path],
175+
_wget_command(url, temp_path),
171176
stdout=subprocess.DEVNULL,
172177
stderr=subprocess.PIPE,
173178
)
@@ -183,7 +188,7 @@ def test_download_command_construction(self, mock_run: mock.Mock, tmp_path: Path
183188

184189
downloader._download_to_path(url, temp_path)
185190

186-
expected_cmd = ["wget", url, "-O", temp_path]
191+
expected_cmd = _wget_command(url, temp_path)
187192
mock_run.assert_called_once_with(
188193
expected_cmd,
189194
stdout=subprocess.DEVNULL,
@@ -234,8 +239,9 @@ def test_realistic_download_scenario(self, mock_run: mock.Mock, tmp_path: Path):
234239
for i, call in enumerate(mock_run.call_args_list):
235240
args, kwargs = call
236241
assert args[0][0] == "wget"
237-
assert args[0][1] == test_urls[i]
238-
assert args[0][2] == "-O"
242+
assert args[0][1] == f"--user-agent={WIKIMEDIA_USER_AGENT}"
243+
assert args[0][2] == test_urls[i]
244+
assert args[0][3] == "-O"
239245
assert kwargs["stdout"] is None # verbose mode
240246
assert kwargs["stderr"] is None # verbose mode
241247

‎tests/stages/text/download/wikipedia/test_url_generation.py‎

Lines changed: 91 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.
@@ -19,6 +19,7 @@
1919
import requests
2020
from bs4 import BeautifulSoup
2121

22+
from nemo_curator.stages.text.download.wikipedia.constants import WIKIMEDIA_REQUEST_HEADERS
2223
from nemo_curator.stages.text.download.wikipedia.url_generation import WikipediaUrlGenerator
2324

2425

@@ -97,6 +98,90 @@ def mock_get_side_effect(url: str, **kwargs) -> Mock: # noqa: ARG001
9798
"https://dumps.wikimedia.org/enwiki/20230501/enwiki-20230501-pages-articles-multistream2.xml.bz2",
9899
]
99100
assert sorted(urls) == sorted(expected_urls)
101+
for call in mock_get.call_args_list:
102+
assert call.kwargs["headers"] == WIKIMEDIA_REQUEST_HEADERS
103+
104+
@pytest.mark.parametrize(
105+
"candidate_failure",
106+
[
107+
pytest.param(404, id="missing-status"),
108+
pytest.param(408, id="request-timeout-status"),
109+
pytest.param(429, id="rate-limited"),
110+
pytest.param(500, id="server-error"),
111+
pytest.param(503, id="service-unavailable"),
112+
pytest.param(requests.Timeout, id="request-timeout"),
113+
pytest.param(requests.ConnectionError, id="connection-error"),
114+
pytest.param(requests.exceptions.ChunkedEncodingError, id="truncated-response"),
115+
pytest.param(b"not JSON", id="invalid-json"),
116+
pytest.param(b"\xff", id="invalid-utf8"),
117+
],
118+
)
119+
@patch("requests.get")
120+
def test_get_latest_dump_date_falls_back_from_unavailable_candidate(
121+
self,
122+
mock_get: Mock,
123+
candidate_failure: int | bytes | type[requests.RequestException],
124+
):
125+
"""An unavailable newest candidate falls back to an older completed dump."""
126+
completed_dump = {
127+
"jobs": {
128+
"articlesmultistreamdump": {
129+
"status": "done",
130+
"files": {"enwiki-20230501-pages-articles-multistream1.xml.bz2": {}},
131+
}
132+
}
133+
}
134+
135+
def mock_get_side_effect(url: str, **kwargs) -> Mock: # noqa: ARG001
136+
response = Mock(status_code=200)
137+
if url == "https://dumps.wikimedia.org/enwiki":
138+
response.content = b'<a href="20230501/">20230501/</a><a href="20230601/">20230601/</a>'
139+
elif url.endswith("20230601/dumpstatus.json"):
140+
if isinstance(candidate_failure, type):
141+
error_msg = "Temporary request failure"
142+
raise candidate_failure(error_msg)
143+
if isinstance(candidate_failure, int):
144+
response.status_code = candidate_failure
145+
if candidate_failure != requests.codes.not_found:
146+
response.raise_for_status.side_effect = requests.HTTPError(
147+
f"HTTP {candidate_failure}", response=response
148+
)
149+
else:
150+
response.content = candidate_failure
151+
elif url.endswith("20230501/dumpstatus.json"):
152+
response.content = json.dumps(completed_dump).encode("utf-8")
153+
else:
154+
error_msg = f"Unexpected URL: {url}"
155+
raise ValueError(error_msg)
156+
return response
157+
158+
mock_get.side_effect = mock_get_side_effect
159+
160+
urls = WikipediaUrlGenerator(language="en").generate_urls()
161+
162+
assert urls == [
163+
"https://dumps.wikimedia.org/enwiki/20230501/enwiki-20230501-pages-articles-multistream1.xml.bz2"
164+
]
165+
166+
@patch("requests.get")
167+
def test_get_latest_dump_date_http_error(self, mock_get: Mock):
168+
"""HTTP errors from the dump index are surfaced instead of parsed as HTML."""
169+
mock_response = Mock()
170+
mock_response.raise_for_status.side_effect = requests.HTTPError("403 Client Error")
171+
mock_get.return_value = mock_response
172+
173+
with pytest.raises(requests.HTTPError, match="403 Client Error"):
174+
WikipediaUrlGenerator(language="en").generate_urls()
175+
176+
@patch("requests.get")
177+
def test_get_latest_dump_date_no_completed_dump(self, mock_get: Mock):
178+
"""A clear error is raised when the index has no completed dumps."""
179+
mock_response = Mock(status_code=200)
180+
mock_response.content = b'<a href="latest/">latest/</a>'
181+
mock_get.return_value = mock_response
182+
183+
with pytest.raises(ValueError, match="Unable to find a completed Wikipedia dump"):
184+
WikipediaUrlGenerator(language="en").generate_urls()
100185

101186
@patch("requests.get")
102187
def test_get_wikipedia_urls_with_specified_date(self, mock_get: Mock):
@@ -368,4 +453,8 @@ def test_realistic_scenario_specific_dump(self, mock_get: Mock):
368453
"https://dumps.wikimedia.org/eswiki/20230301/eswiki-20230301-pages-articles-multistream2.xml.bz2",
369454
]
370455
assert sorted(urls) == sorted(expected_urls)
371-
mock_get.assert_called_once_with("https://dumps.wikimedia.org/eswiki/20230301/dumpstatus.json", timeout=30)
456+
mock_get.assert_called_once_with(
457+
"https://dumps.wikimedia.org/eswiki/20230301/dumpstatus.json",
458+
headers=WIKIMEDIA_REQUEST_HEADERS,
459+
timeout=30,
460+
)

0 commit comments

Comments
 (0)