ensure nltk tokenizer works

salt/logic/embeddings.py

@@ -1,3 +1,4 @@
+import nltk
 import pickle
 import numpy as np
 import pandas as pd
@@ -7,9 +8,9 @@
 from itertools import chain
 from typing import Dict, List
 from salt.constants import NA
-from nltk import sent_tokenize
 from sentence_transformers import SentenceTransformer
 
+nltk.download('punkt')
 
 TEXTS_KEY = "texts"
 VECTORS_KEY = "vectors"
@@ -42,7 +43,7 @@ def get_relevant_texts(texts: list[str]) -> list[str]:
 
 def embed_texts(texts: List[str]) -> List[List[float]]:
     model = get_model(MODEL_NAME)
-    texts_sentences = [sent_tokenize(text) for text in texts]
+    texts_sentences = [nltk.sent_tokenize(text) for text in texts]
     texts_lengths = [len(sentences) for sentences in texts_sentences]
     all_sentences = list(chain.from_iterable(texts_sentences))
     if len(all_sentences) <= BATCH_SIZE:

salt/resources/thin_classifier.py

@@ -2,21 +2,22 @@
 # scikit-learn==1.3.2
 # sentence-transformers==2.2.2
 
+import nltk
 import pickle
 import numpy as np
 from tqdm.auto import tqdm
-from nltk import sent_tokenize
 from sklearn.linear_model import LogisticRegression
 from sentence_transformers import SentenceTransformer
 
+nltk.download('punkt')
 
 def load_model(path: str) -> LogisticRegression:
     with open(path, "rb") as f:
         return pickle.load(f)
 
 
 def vectorize(text: str, model: SentenceTransformer) -> np.array:
-    sentences = sent_tokenize(text)
+    sentences = nltk.sent_tokenize(text)
     sentences_embeddings = model.encode(sentences)
     return sentences_embeddings.mean(axis=0)