Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

build: cache nltk models into the docker image #4118

Merged
merged 6 commits into from
Feb 16, 2023
Merged

build: cache nltk models into the docker image #4118

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
2cb379a
separated nltk cache
mayankjobanputra Feb 8, 2023
9d15b5a
separated nltk caching
mayankjobanputra Feb 9, 2023
5e411cb
Merge branch 'main' of github.com:deepset-ai/haystack into mayank/cac…
mayankjobanputra Feb 9, 2023
116fdbd
fixed pylint lazy log error
mayankjobanputra Feb 9, 2023
f4be9e3
using model name as default value
mayankjobanputra Feb 9, 2023
fe888a8
Merge branch 'mayank/cache_nltk' of github.com:deepset-ai/haystack in…
mayankjobanputra Feb 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions docker/Dockerfile.base
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,3 +47,7 @@ ENV PATH="/opt/venv/bin:$PATH"
# - the schema will be already there when the container runs, saving the generation overhead when a container starts
# - derived images don't need to write the schema and can run with lower user privileges
RUN python3 -c "from haystack.utils.docker import cache_schema; cache_schema()"

# Haystack Preprocessor uses NLTK punkt model to divide text into a list of sentences.
# We cache these models for seemless user experience.
RUN python3 -c "from haystack.utils.docker import cache_nltk_model; cache_nltk_model()"
13 changes: 7 additions & 6 deletions haystack/utils/docker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,13 @@
from haystack.nodes._json_schema import load_schema


def cache_nltk_model(model: str = "punkt"):
logging.info("Caching %s model...", model)
import nltk

nltk.download(model)


def cache_models(models: Optional[List[str]] = None, use_auth_token: Optional[Union[str, bool]] = None):
"""
Small function that caches models and other data.
Expand All @@ -19,12 +26,6 @@ def cache_models(models: Optional[List[str]] = None, use_auth_token: Optional[Un
if models is None:
models = ["deepset/roberta-base-squad2"]

# download punkt tokenizer
logging.info("Caching punkt data")
import nltk

nltk.download("punkt")
mayankjobanputra marked this conversation as resolved.
Show resolved Hide resolved

# Cache models
import transformers

Expand Down