From 353d9fc72b6266c7d0407a882e8d772acec7308f Mon Sep 17 00:00:00 2001
From: imClumsyPanda <littlepanda0716@gmail.com>
Date: Sun, 13 Aug 2023 10:25:02 +0800
Subject: [PATCH] update init_database.py

---
 init_database.py       | 4 +++-
 requirements.txt       | 1 +
 requirements_api.txt   | 1 +
 requirements_webui.txt | 1 +
 webui_pages/utils.py   | 4 +++-
 5 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/init_database.py b/init_database.py
index 2d3f67a..61d00e1 100644
--- a/init_database.py
+++ b/init_database.py
@@ -1,5 +1,7 @@
 from server.knowledge_base.migrate import create_tables, folder2db, recreate_all_vs, list_kbs_from_folder
-
+from configs.model_config import NLTK_DATA_PATH
+import nltk
+nltk.data.path = [NLTK_DATA_PATH] + nltk.data.path
 
 if __name__ == "__main__":
     import argparse
diff --git a/requirements.txt b/requirements.txt
index fd1a0a1..f2e1d65 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,6 +14,7 @@ unstructured[all-docs]
 python-magic-bin; sys_platform == 'win32'
 SQLAlchemy==2.0.19
 faiss-cpu
+nltk
 
 # uncomment libs if you want to use corresponding vector store
 # pymilvus==2.1.3 # requires milvus==2.1.3
diff --git a/requirements_api.txt b/requirements_api.txt
index 9e60611..f077c94 100644
--- a/requirements_api.txt
+++ b/requirements_api.txt
@@ -14,6 +14,7 @@ unstructured[all-docs]
 python-magic-bin; sys_platform == 'win32'
 SQLAlchemy==2.0.19
 faiss-cpu
+nltk
 
 # uncomment libs if you want to use corresponding vector store
 # pymilvus==2.1.3 # requires milvus==2.1.3
diff --git a/requirements_webui.txt b/requirements_webui.txt
index 0f7e253..a832550 100644
--- a/requirements_webui.txt
+++ b/requirements_webui.txt
@@ -6,3 +6,4 @@ streamlit-antd-components>=0.1.11
 streamlit-chatbox>=1.1.6
 streamlit-aggrid>=0.3.4.post3
 httpx~=0.24.1
+nltk
\ No newline at end of file
diff --git a/webui_pages/utils.py b/webui_pages/utils.py
index d1e5013..e64403c 100644
--- a/webui_pages/utils.py
+++ b/webui_pages/utils.py
@@ -6,7 +6,6 @@ from configs.model_config import (
     DEFAULT_VS_TYPE,
     KB_ROOT_PATH,
     LLM_MODEL,
-    llm_model_dict,
     VECTOR_SEARCH_TOP_K,
     SEARCH_ENGINE_TOP_K,
 )
@@ -21,6 +20,9 @@ from server.db.repository.knowledge_base_repository import get_kb_detail
 from server.db.repository.knowledge_file_repository import get_file_detail
 from server.utils import run_async, iter_over_async
 
+from configs.model_config import NLTK_DATA_PATH
+import nltk
+nltk.data.path = [NLTK_DATA_PATH] + nltk.data.path
 
 def set_httpx_timeout(timeout=60.0):
     '''