process new dataset (#9)

cfg.py

@@ -21,18 +21,21 @@ PASSWORD = os.getenv("BUSTER_PASSWORD")
 
 HUB_TOKEN = os.getenv("HUB_TOKEN")
 REPO_ID = os.getenv("HF_DATASET")
-HUB_DB_FILE = "deeplake_store.zip"
+# HUB_DB_FILE = "deeplake_store.zip"
 
-logger.info(f"Downloading {HUB_DB_FILE} from hub...")
+DEEPLAKE_DATASET = os.getenv("DEEPLAKE_DATASET", "wiki_tai_langchain")
+ZIP_FILE = DEEPLAKE_DATASET + ".zip"
+
+logger.info(f"Downloading {ZIP_FILE} from hub...")
 hf_hub_download(
     repo_id=REPO_ID,
     repo_type="dataset",
-    filename=HUB_DB_FILE,
+    filename=ZIP_FILE,
     token=HUB_TOKEN,
     local_dir=".",
 )
 
-extract_zip(zip_file_path=HUB_DB_FILE, output_path="deeplake_store")
+extract_zip(zip_file_path=ZIP_FILE, output_path=DEEPLAKE_DATASET)
 
 example_questions = [
     "What is the LLama model?",
@@ -71,7 +74,7 @@ A user will now submit a question. Respond 'true' if it is valid, respond 'false
         },
     },
     retriever_cfg={
-        "path": "./deeplake_store",
+        "path": f"./{DEEPLAKE_DATASET}",
         "top_k": 3,
         "thresh": 0.7,
         "max_tokens": 2000,

embed_documents.py

@@ -2,13 +2,17 @@ import pandas as pd
 from buster.documents_manager import DeepLakeDocumentsManager
 
 if __name__ == "__main__":
-    vector_store_path = "deeplake_store"
-    chunk_file = "data/wiki_and_tai.csv"
+    vector_store_path = "wiki_tai_langchain"
+    chunk_file = "./data/wiki_tai_langchain.csv"
     overwrite = True
 
     df = pd.read_csv(chunk_file)
 
-    dm = DeepLakeDocumentsManager(vector_store_path, overwrite=overwrite)
+    print(f"before drop: {len(df)}")
+    df = df.dropna()
+    print(f"after drop: {len(df)}")
+
+    dm = DeepLakeDocumentsManager(vector_store_path, overwrite=overwrite, required_columns=["url", "source", "content", "title"])
     dm.batch_add(df)
     zipped_file_path = dm.to_zip()
     print(f"Contents zipped to: {zipped_file_path}")