"""Functions for downloading DaCy models."""
import os
import shutil
import zipfile
from pathlib import Path
from tqdm import tqdm
DACY_DEFAULT_PATH = Path.home() / ".cache" / "dacy"
DEFAULT_CACHE_DIR = Path(
os.getenv(
"DACY_CACHE_DIR",
DACY_DEFAULT_PATH,
),
)
# Models bundling a coref component, which needs the `dacy[coref]` extra.
COREF_MODELS = {
"da_dacy_small_trf-0.2.0",
"da_dacy_medium_trf-0.2.0",
"da_dacy_large_trf-0.2.0",
}
models_url = {
"da_dacy_small_trf-0.2.0": "https://huggingface.co/chcaa/da_dacy_small_trf/resolve/0eadea074d5f637e76357c46bbd56451471d0154/da_dacy_small_trf-any-py3-none-any.whl",
"da_dacy_medium_trf-0.2.0": "https://huggingface.co/chcaa/da_dacy_medium_trf/resolve/e7dba91f855a1d26679dc1ef3aa49f7874b50543/da_dacy_medium_trf-any-py3-none-any.whl",
"da_dacy_large_trf-0.2.0": "https://huggingface.co/chcaa/da_dacy_large_trf/resolve/963232f378190476503a1bfc35b520cb142e9e41/da_dacy_large_trf-any-py3-none-any.whl",
"small": None,
"medium": None,
"large": None,
"da_dacy_small_ner_fine_grained-0.1.0": "https://huggingface.co/chcaa/da_dacy_small_ner_fine_grained/resolve/43fedc5a1b1c1d193f461d13225f217f2ced507d/da_dacy_small_ner_fine_grained-any-py3-none-any.whl",
"da_dacy_medium_ner_fine_grained-0.1.0": "https://huggingface.co/chcaa/da_dacy_medium_ner_fine_grained/resolve/4bfc4397b720acdb6428d64f18e90bfd439c80fc/da_dacy_medium_ner_fine_grained-any-py3-none-any.whl",
"da_dacy_large_ner_fine_grained-0.1.0": "https://huggingface.co/chcaa/da_dacy_large_ner_fine_grained/resolve/08f973a1ff57120268bf30d3b7e7c4656ed25a58/da_dacy_large_ner_fine_grained-any-py3-none-any.whl",
}
[docs]def get_latest_version(model: str) -> str:
"""Returns the latest version of a DaCy model.
Args:
model: string indicating the model
Returns:
str: latest version of the model
"""
if model in {"small", "medium", "large"}:
model = f"da_dacy_{model}_trf"
versions = [mdl.split("-")[-1] for mdl in models_url if mdl.startswith(model)]
versions = sorted(
versions,
key=lambda s: [int(u) for u in s.split(".")],
reverse=True,
)
return versions[0]
class DownloadProgressBar(tqdm):
def update_to(self, b: int = 1, bsize: int = 1, tsize=None) -> None: # noqa
if tsize is not None:
self.total = tsize
self.update(b * bsize - self.n)
def download_url(url: str, output_path: str) -> None:
import urllib.request
with DownloadProgressBar(
unit="B",
unit_scale=True,
miniters=1,
desc=url.split("/")[-1],
) as t:
urllib.request.urlretrieve(url, filename=output_path, reporthook=t.update_to)
def _check_coref_dependencies(package: str) -> None:
"""Raises an informative ImportError if `package` requires the `dacy[coref]`
extra (`spacy-experimental`) and it is not installed.
"""
if package not in COREF_MODELS:
return
try:
import spacy_experimental # type: ignore # noqa: F401
except ImportError as e:
raise ImportError(
f"The DaCy model '{package}' includes a coreference resolution "
"component that requires the optional 'spacy-experimental' "
"dependency. Install it with `pip install dacy[coref]` (only "
"available for Python <3.12), or use a DaCy model that does not "
"include coreference resolution.",
) from e
[docs]def download_model(
model: str,
force: bool = False,
) -> str:
"""Downloads a specified DaCy pipeline to the DaCy cache
(`dacy.download.DEFAULT_CACHE_DIR`, configurable via the `DACY_CACHE_DIR`
environment variable) and returns the path to the pipeline. The returned
path can be loaded using `spacy.load`. Unlike installing the model as a
python package, this does not touch the environment's installed
dependencies.
Args:
model: string indicating DaCy model, use dacy.models() to get a list of
models.
force: Should it download the model regardless of it already
being present? Defaults to False.
Returns:
a string of the model location
Example:
>>> download_model(model="da_dacy_medium_trf-0.1.0")
"""
if model in {"small", "medium", "large"}:
latest_version = get_latest_version(model)
model = f"da_dacy_{model}_trf-{latest_version}"
if model not in models_url:
raise ValueError(
f"The model '{model}' is not available in DaCy. Please use dacy.models() to see a"
+ " list of all models",
)
_check_coref_dependencies(model)
package = model.split("-")[0]
model_dir = Path(DEFAULT_CACHE_DIR) / model
pipeline_dir = model_dir / package / model
if pipeline_dir.exists() and not force:
return str(pipeline_dir)
if model_dir.exists():
shutil.rmtree(model_dir)
model_dir.mkdir(parents=True, exist_ok=True)
whl_path = model_dir / f"{model}.whl"
download_url(models_url[model], str(whl_path)) # type: ignore
with zipfile.ZipFile(whl_path) as whl:
whl.extractall(model_dir)
whl_path.unlink()
if not pipeline_dir.exists():
raise ValueError(
f"Could not locate the pipeline data for '{model}' after extracting "
+ "the downloaded model.",
)
return str(pipeline_dir)