klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-05 21:29:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

Test memory leak on MPS device when generating vector embeddings

Browse Source 
Slope threshold of 2.0 determined qualitatively on local Mac device
Minor unused import and clean-up
This commit is contained in:
Debanjum Singh Solanky
2023-11-05 03:32:29 -08:00
parent ef24485ada
commit a4f407f595
4 changed files with 41 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/database/adapters/__init__.py
View File

@@ -1,4 +1,3 @@
import secrets
from typing import Type, TypeVar, List
from datetime import date
import secrets
@@ -36,9 +35,6 @@ from database.models import (
OfflineChatProcessorConversationConfig,
)
from khoj.utils.helpers import generate_random_name
from khoj.utils.rawconfig import (
ConversationProcessorConfig as UserConversationProcessorConfig,
)
from khoj.search_filter.word_filter import WordFilter
from khoj.search_filter.file_filter import FileFilter
from khoj.search_filter.date_filter import DateFilter

6
src/khoj/processor/embeddings.py
View File

@@ -8,10 +8,10 @@ from khoj.utils.rawconfig import SearchResponse
class EmbeddingsModel:
def __init__(self):
self.model_name = "thenlper/gte-small"
self.encode_kwargs = {"normalize_embeddings": True}
model_kwargs = {"device": get_device()}
self.embeddings_model = SentenceTransformer(self.model_name, **model_kwargs)
self.model_kwargs = {"device": get_device()}
self.model_name = "thenlper/gte-small"
self.embeddings_model = SentenceTransformer(self.model_name, **self.model_kwargs)
def embed_query(self, query):
return self.embeddings_model.encode([query], show_progress_bar=False, **self.encode_kwargs)[0]