Fix doc indexing and use accelerator (#30)

Signed-off-by: Aivin V. Solatorio <[email protected]>

llm4data/embeddings/base.py

@@ -5,12 +5,22 @@
 from qdrant_client.http import models
 from pydantic.main import ModelMetaclass
 from dataclasses import dataclass, asdict
-
+import torch
 
 # Make the model atomically available
 LOADED_MODELS: dict = {}
 
 
+def get_device():
+    device = "cpu"
+    if torch.cuda.device_count() > 0:
+        device = "cuda:0"
+    elif torch.backends.mps.is_available():
+        device = "mps:0"
+
+    return device
+
+
 @dataclass
 class BaseEmbeddingModel:
     model_size = {
@@ -101,7 +111,7 @@ def _create_embeddings(self):
             raise ValueError("`config.kwargs` must be a dict")
 
         self.embeddings = getattr(langchain_embeddings, self.embedding_cls)(
-            **self.kwargs
+            **{"model_kwargs": {"device": get_device()}, **self.kwargs}
         )
 
         if self.max_tokens is None and self.embeddings:

llm4data/scripts/indexing/docs/docs.py

@@ -1,7 +1,7 @@
 from typing import Optional, Union
 from pathlib import Path
 from langchain.docstore.document import Document
-from langchain.document_loaders import PyMuPDFLoader
+from langchain_community.document_loaders import PyMuPDFLoader
 from langchain.text_splitter import (
     NLTKTextSplitter,
     CharacterTextSplitter,

llm4data/scripts/indexing/docs/load_docs.py

@@ -106,5 +106,5 @@ def main(path: Union[str, Path], strict: bool = False):
 
 
 if __name__ == "__main__":
-    # python -m llm4data.scripts.indexing.docs.load_docs --path=data/knowledge/docs/prwp/pdf --strict
+    # python -m llm4data.scripts.indexing.docs.load_docs --path=data/sources/docs/prwp/pdf --strict
     fire.Fire(main)