Update app.py

app.py

@@ -15,6 +15,8 @@ from pathlib import Path
 import chromadb
 from unidecode import unidecode
 
+from babi_app import BabyAGI
+
 from transformers import AutoTokenizer
 import transformers
 import torch
@@ -22,16 +24,16 @@ import tqdm
 import accelerate
 import re
 
-
-
-# default_persist_directory = './chroma_HF/'
-list_llm = ["mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.1", \
-    "google/gemma-7b-it","google/gemma-2b-it", \
-    "HuggingFaceH4/zephyr-7b-beta", "HuggingFaceH4/zephyr-7b-gemma-v0.1", \
-    "meta-llama/Llama-2-7b-chat-hf", "microsoft/phi-2", \
-    "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "mosaicml/mpt-7b-instruct", "tiiuae/falcon-7b-instruct", \
-    "google/flan-t5-xxl"
-]
+default_persist_directory = './chroma_HF/'
+list_llm = ["mistralai/Mistral-7B-Instruct-v0.2", 
+            "mistralai/Mixtral-8x7B-Instruct-v0.1", 
+            "mistralai/Mistral-7B-Instruct-v0.1", 
+            "google/gemma-7b-it", "google/gemma-2b-it", 
+            "HuggingFaceH4/zephyr-7b-beta", "HuggingFaceH4/zephyr-7b-gemma-v0.1", 
+            "meta-llama/Llama-2-7b-chat-hf", "microsoft/phi-2", 
+            "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "mosaicml/mpt-7b-instruct", 
+            "tiiuae/falcon-7b-instruct", "google/flan-t5-xxl"
+        ]
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
 # Load PDF document and create doc splits
@@ -43,7 +45,6 @@ def load_doc(list_file_path, chunk_size, chunk_overlap):
     pages = []
     for loader in loaders:
         pages.extend(loader.load())
-    # text_splitter = RecursiveCharacterTextSplitter(chunk_size = 600, chunk_overlap = 50)
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size = chunk_size, 
         chunk_overlap = chunk_overlap)
@@ -69,7 +70,7 @@ def create_db(splits, collection_name):
 def load_db():
     embedding = HuggingFaceEmbeddings()
     vectordb = Chroma(
-        # persist_directory=default_persist_directory, 
+        persist_directory=default_persist_directory, 
         embedding_function=embedding)
     return vectordb