yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2025-12-12 15:27:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

build: cache nltk models into the docker image (#4118)

Browse Source 
* separated nltk cache

* separated nltk caching

* fixed pylint lazy log error

* using model name as default value
This commit is contained in:
Mayank Jobanputra 2023-02-16 12:26:16 +01:00 committed by GitHub
parent ec72dd73fc
commit d27f372b67
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docker/Dockerfile.base
View File

@ -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
haystack/utils/docker.py
View File

@ -3,6 +3,13 @@ from typing import List, Union, Optional
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.
@ -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")
# Cache models
import transformers