device conditionals

src/core/embedding.py

@@ -1,5 +1,6 @@
 from typing import List, Type
 
+import torch
 from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.embeddings.base import Embeddings
@@ -49,7 +50,13 @@ class FolderIndex:
 
 def embed_files(files: List[File]) -> FolderIndex:
     model_name = "adriancowham/letstalk-mythomax-embed-gte-small"
+
     model_kwargs = {'device': 'cpu'}
+    if torch.cuda.is_available():
+      model_kwargs['device'] = 'cuda'
+    if torch.backends.mps.is_available():
+      model_kwargs['device'] = 'mps'
+
     encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
     print("Loading model...")
     try: