Kunal Pai
commited on
Commit
·
ab59793
1
Parent(s):
941e057
MPS support for memories for Mac users
Browse files- src/manager/manager.py +6 -1
src/manager/manager.py
CHANGED
@@ -164,7 +164,12 @@ class GeminiManager:
|
|
164 |
return []
|
165 |
top_k = min(k, len(memories))
|
166 |
# Semantic Retrieval with GPU
|
167 |
-
|
|
|
|
|
|
|
|
|
|
|
168 |
print(f"Using device: {device}")
|
169 |
model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
|
170 |
doc_embeddings = model.encode(memories, convert_to_tensor=True, device=device)
|
|
|
164 |
return []
|
165 |
top_k = min(k, len(memories))
|
166 |
# Semantic Retrieval with GPU
|
167 |
+
if torch.cuda.is_available():
|
168 |
+
device = 'cuda'
|
169 |
+
elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
|
170 |
+
device = 'mps'
|
171 |
+
else:
|
172 |
+
device = 'cpu'
|
173 |
print(f"Using device: {device}")
|
174 |
model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
|
175 |
doc_embeddings = model.encode(memories, convert_to_tensor=True, device=device)
|