Update app.py

app.py

@@ -1,5 +1,7 @@
 import os
 import re
+import nltk
+import copy
 import time
 import joblib
 import msgpack
@@ -196,6 +198,9 @@ def load_models_and_documents():
             providers=provider
         )
 
+        nltk.download('punkt')
+        nltk.download('averaged_perceptron_tagger')
+
     client = QdrantClient(':memory:')
     collection_name = 'collection_demo'
     
@@ -307,26 +312,31 @@ def chunk_documents(docs, dense_model, sparse_model):
         breakpoint_threshold_type='standard_deviation'
     )
 
-    texts = []
+    texts, metadatas = [], []
     for doc in docs:
         texts.append(doc.page_content)
+        metadatas.append(doc.metadata)
 
+    _metadatas = metadatas or [{}] * len(texts)
     documents = []
 
-    def create_document(text: str, i: int):
+    def create_document(text: str, i: int, _metadatas: list):
         index = -1
         for chunk in text_splitter.split_text(text):
+            metadata = copy.deepcopy(_metadatas[i])
             if text_splitter._add_start_index:
                 index = text.find(chunk, index + 1)
-        documents.append(chunk)
+                metadata['start_index'] = index
+        new_doc = Document(page_content=chunk, metadata=metadata)       
+        documents.append(new_doc)
         
     joblib.Parallel(n_jobs=joblib.cpu_count(), verbose=1, require='sharedmem')(
-        joblib.delayed(create_document)(text, i) for i, text in enumerate(texts))
+        joblib.delayed(create_document)(text, i, _metadatas) for i, text in enumerate(texts))
 
-    # documents = [doc.page_content for doc in text_splitter.transform_documents(list(docs))]
+    docs = [doc.page_content for doc in documents]
 
-    dense_embeddings = dense_model.embed_documents(documents,32)
-    sparse_embeddings = list(sparse_model.embed(documents, 32))
+    dense_embeddings = dense_model.embed_documents(docs,32)
+    sparse_embeddings = list(sparse_model.embed(docs, 32))
     
     return documents, dense_embeddings, sparse_embeddings
     
@@ -335,8 +345,6 @@ if __name__ == '__main__':
                        layout="wide"
                        )
     st.title("Video Game Assistant")
-    num_cores = joblib.cpu_count()
-    print(num_cores)
     
     client, collection_name, llm, dense_model, sparse_model = load_models_and_documents()