klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-06 13:22:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix passing of device to setup method in /reload, /regenerate API

Browse Source 
- Use local variable to pass device to asymmetric.setup method via /reload, /regenerate API
- Set default argument to torch.device('cpu') instead of 'cpu' to be more formal
This commit is contained in:
Debanjum Singh Solanky
2022-06-30 01:32:56 +04:00
parent eda4b65ddb
commit 7677465f23
2 changed files with 10 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/search_type/asymmetric.py
View File

@@ -93,7 +93,7 @@ def compute_embeddings(entries, bi_encoder, embeddings_file, regenerate=False, d
return corpus_embeddings
def query(raw_query: str, model: TextSearchModel, device='cpu'):
def query(raw_query: str, model: TextSearchModel, device=torch.device('cpu')):
"Search all notes for entries that answer the query"
# Separate natural query from explicit required, blocked words filters
query = " ".join([word for word in raw_query.split() if not word.startswith("+") and not word.startswith("-")])
@@ -180,7 +180,7 @@ def collate_results(hits, entries, count=5):
in hits[0:count]]
def setup(config: TextContentConfig, search_config: AsymmetricSearchConfig, regenerate: bool, device='cpu', verbose: bool=False) -> TextSearchModel:
def setup(config: TextContentConfig, search_config: AsymmetricSearchConfig, regenerate: bool, device=torch.device('cpu'), verbose: bool=False) -> TextSearchModel:
# Initialize Model
bi_encoder, cross_encoder, top_k = initialize_model(search_config)