Update scripts/load_documents.py

scripts/load_documents.py

@@ -7,31 +7,35 @@ from langchain_community.document_loaders import (
 )
 import pickle
 
-BASE_DIR = Path(__file__).resolve().parent.parent
-DATA_DIR = BASE_DIR / "data"
-OUTPUT_DIR = BASE_DIR / "output"
-OUTPUT_PATH = OUTPUT_DIR / "all_docs.pkl"
+def main():
+    BASE_DIR = Path(__file__).resolve().parent.parent
+    DATA_DIR = BASE_DIR / "data"
+    OUTPUT_DIR = BASE_DIR / "output"
+    OUTPUT_PATH = OUTPUT_DIR / "all_docs.pkl"
+    
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    
+    loaders = {
+        ".pdf": PyPDFLoader,
+        ".txt": lambda path: TextLoader(path, encoding="utf-8"),
+        ".py": PythonLoader,
+        ".ipynb": NotebookLoader,
+    }
+    
+    documents = []
+    for file in DATA_DIR.rglob("*"):
+        loader_class = loaders.get(file.suffix.lower())
+        if loader_class:
+            try:
+                docs = loader_class(str(file)).load()
+                documents.extend(docs)
+                print(f"[✓] Loaded: {file.name}")
+            except Exception as e:
+                print(f"[!] Failed to load {file.name}: {e}")
+    
+    with open(OUTPUT_PATH, "wb") as f:
+        pickle.dump(documents, f)
+    print(f"📦 Saved {len(documents)} documents to {OUTPUT_PATH}")
 
-OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-
-loaders = {
-    ".pdf": PyPDFLoader,
-    ".txt": lambda path: TextLoader(path, encoding="utf-8"),
-    ".py": PythonLoader,
-    ".ipynb": NotebookLoader,
-}
-
-documents = []
-for file in DATA_DIR.rglob("*"):
-    loader_class = loaders.get(file.suffix.lower())
-    if loader_class:
-        try:
-            docs = loader_class(str(file)).load()
-            documents.extend(docs)
-            print(f"[✓] Loaded: {file.name}")
-        except Exception as e:
-            print(f"[!] Failed to load {file.name}: {e}")
-
-with open(OUTPUT_PATH, "wb") as f:
-    pickle.dump(documents, f)
-print(f"📦 Saved {len(documents)} documents to {OUTPUT_PATH}")
+if __name__ == "__main__":
+    main()