feat(retrieval): prefer CUDA devices
This commit is contained in:
1 parent
5ec6d0b6ac
commit
e5b63965b8
2 files changed
+15
-2
No files matched your search
@@ -118,6 +118,11 @@ def document_text(path: str, symbol: str, content: str):
|
||||
return f"Path: {path}\nSymbol: {symbol}\n\n{content}"
|
||||
|
||||
|
||||
def preferred_device():
|
||||
import torch
|
||||
return "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
def reciprocal_rank_fusion(*rankings):
|
||||
scores = {}
|
||||
for ranking in rankings:
|
||||
@@ -165,13 +170,15 @@ class Index:
|
||||
if self.model is None:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
self.model = SentenceTransformer(MODEL, truncate_dim=EMBEDDING_DIMENSIONS,
|
||||
prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "})
|
||||
prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "},
|
||||
device=preferred_device())
|
||||
return self.model
|
||||
|
||||
def reranker(self):
|
||||
if self.reranker_model is None:
|
||||
from sentence_transformers import CrossEncoder
|
||||
self.reranker_model = CrossEncoder(RERANKER_MODEL, prompts={"query": RETRIEVAL_INSTRUCTION}, default_prompt_name="query")
|
||||
self.reranker_model = CrossEncoder(RERANKER_MODEL, prompts={"query": RETRIEVAL_INSTRUCTION},
|
||||
default_prompt_name="query", device=preferred_device())
|
||||
return self.reranker_model
|
||||
|
||||
def rerank(self, query: str, rows: list[dict], limit: int):
|
||||
|
||||
Reference in new issue
Block a user