update collection selector

app.py

@@ -11,13 +11,13 @@ from langchain.chains import ConversationalRetrievalChain
 
 
 def load_embeddings():
-    print("Loading embeddings...")
+    print("loading embeddings...")
     model_name = os.environ['HUGGINGFACEHUB_EMBEDDINGS_MODEL_NAME']
     return HuggingFaceInstructEmbeddings(model_name=model_name)
 
 
 def split_file(file, chunk_size, chunk_overlap):
-    print('spliting file', file.name)
+    print('spliting file...', file.name, chunk_size, chunk_overlap)
     loader = TextLoader(file.name)
     documents = loader.load()
     text_splitter = CharacterTextSplitter(
@@ -30,17 +30,17 @@ def get_persist_directory(file_name):
 
 
 def process_file(file, chunk_size, chunk_overlap):
-    docs = split_file(file, chunk_size, chunk_overlap)
-    embeddings = load_embeddings()
-
     file_name, _ = os.path.splitext(os.path.basename(file.name))
     persist_directory = get_persist_directory(file_name)
     print("persist directory", persist_directory)
+
+    docs = split_file(file, chunk_size, chunk_overlap)
+    embeddings = load_embeddings()
     vectordb = Chroma.from_documents(documents=docs, embedding=embeddings,
                                      collection_name=file_name, persist_directory=persist_directory)
     print(vectordb._client.list_collections())
     vectordb.persist()
-    return 'Done!'
+    return 'Done!', gr.Dropdown.update(choices=get_vector_dbs(), value=file_name)
 
 
 def is_dir(root, name):
@@ -53,8 +53,9 @@ def get_vector_dbs():
     if not os.path.exists(root):
         return []
 
+    print('get vector dbs...', root)
     files = os.listdir(root)
-    dirs = filter(lambda x: is_dir(root, x), files)
+    dirs = list(filter(lambda x: is_dir(root, x), files))
     print(dirs)
     return dirs
 
@@ -71,7 +72,7 @@ def load_vectordb(file_name):
 
 
 def create_qa_chain(collection_name, temperature, max_length):
-    print('creating qa chain...')
+    print('creating qa chain...', collection_name, temperature, max_length)
     memory = ConversationBufferMemory(
         memory_key="chat_history", return_messages=True)
     llm = HuggingFaceHub(
@@ -116,8 +117,9 @@ with gr.Blocks() as demo:
     with gr.Tab("Bot"):
         with gr.Row():
             with gr.Column(scale=0.5):
+                choices = get_vector_dbs()
                 collection = gr.Dropdown(
-                    choices=get_vector_dbs(), label="Document")
+                    choices, value=choices[0] if choices else None, label="Document")
                 temperature = gr.Slider(
                     0.0, 1.0, value=0.5, step=0.05, label="Temperature")
                 max_length = gr.Slider(20, 1000, value=64, label="Max Length")
@@ -128,7 +130,8 @@ with gr.Blocks() as demo:
                     show_label=False, placeholder="Ask me anything!")
                 clear = gr.Button("Clear")
 
-    process.click(process_file, [upload, chunk_size, chunk_overlap], result)
+    process.click(process_file, [upload, chunk_size,
+                  chunk_overlap], [result, collection])
 
     message.submit(submit_message, [chatbot, message], [chatbot, message]).then(
         bot, [chatbot, collection, temperature, max_length], chatbot