Merge pull request #60 from weaviate/nltk-download-dir-fix

Fix downloading of nltk data into a local folder

.dockerignore

@@ -1 +1,2 @@
 models
+nltk_data

.gitignore

@@ -1,2 +1,3 @@
 __pycache__
 models
+nltk_data

cicd/build_custom_base.sh

@@ -41,7 +41,7 @@ function push_main() {
     # The ones that are always pushed
 
     tag="$remote_repo:custom-$git_hash"
-    docker buildx build -f custom.Dockerfile \
+    docker buildx build --platform=linux/arm64,linux/amd64 -f custom.Dockerfile \
       --tag "$tag" \
       --push \
       .
@@ -57,7 +57,7 @@ function push_tag() {
     echo "Tag & Push $tag, $tag_latest, $tag_git"
     docker tag "custom-base" "$tag" && docker push "$tag"
 
-    docker buildx build -f custom.Dockerfile \
+    docker buildx build --platform=linux/arm64,linux/amd64 -f custom.Dockerfile \
       --tag "$tag" \
       --tag "$tag_latest" \
       --tag "$tag_git" \

download.py

@@ -39,4 +39,4 @@
 model.save_pretrained('./models/model')
 tokenizer.save_pretrained('./models/model')
 
-nltk.download('punkt')
+nltk.download('punkt', download_dir='./nltk_data')

vectorizer.py

@@ -3,6 +3,7 @@
 import math
 from typing import Optional
 import torch
+import nltk
 from nltk.tokenize import sent_tokenize
 from pydantic import BaseModel
 from transformers import (
@@ -57,6 +58,8 @@ def __init__(self, model_path: str, cuda_support: bool, cuda_core: str, cuda_per
 
         self.executor = ThreadPoolExecutor()
 
+        nltk.data.path.append('./nltk_data')
+
     def tokenize(self, text:str):
         return self.tokenizer(text, padding=True, truncation=True, max_length=500,
                 add_special_tokens = True, return_tensors="pt")