fix(sentence-transformers): use cuda if available

Author: rti

gswikichat/vector_store_interface.py

@@ -11,6 +11,7 @@
 from haystack.components.preprocessors import DocumentSplitter
 from haystack.components.preprocessors import DocumentCleaner
 
+import torch
 
 from .logger import get_logger
 
@@ -25,6 +26,13 @@
 top_k = 5
 input_documents = []
 
+device = "cpu"
+
+if torch.cuda.is_available():
+    logger.info('GPU is available.')
+    device = "cuda"
+
+
 # TODO: Add the json strings as env variables
 json_dir = 'json_input'
 json_fname = 'excellent-articles_10.json'
@@ -95,6 +103,7 @@
 
 embedder = SentenceTransformersDocumentEmbedder(
     model=sentence_transformer_model,
+    device=device
 )
 embedder.warm_up()