Farid Karimli commited on
Commit
40de40e
1 Parent(s): 71b8a7a

similarity score in response fix

Browse files
Files changed (1) hide show
  1. code/modules/vector_db.py +23 -3
code/modules/vector_db.py CHANGED
@@ -5,6 +5,7 @@ from langchain.vectorstores import FAISS, Chroma
5
  from langchain.schema.vectorstore import VectorStoreRetriever
6
  from langchain.callbacks.manager import CallbackManagerForRetrieverRun
7
  from langchain.schema.document import Document
 
8
 
9
  try:
10
  from modules.embedding_model_loader import EmbeddingModelLoader
@@ -17,12 +18,15 @@ except:
17
  from constants import *
18
  from helpers import *
19
 
 
 
20
 
21
  class VectorDBScore(VectorStoreRetriever):
 
22
  # See https://github.com/langchain-ai/langchain/blob/61dd92f8215daef3d9cf1734b0d1f8c70c1571c3/libs/langchain/langchain/vectorstores/base.py#L500
23
  def _get_relevant_documents(
24
- self, query: str, *, run_manager: CallbackManagerForRetrieverRun
25
- ):
26
  docs_and_similarities = (
27
  self.vectorstore.similarity_search_with_relevance_scores(
28
  query, **self.search_kwargs
@@ -35,6 +39,22 @@ class VectorDBScore(VectorStoreRetriever):
35
  docs = [doc for doc, _ in docs_and_similarities]
36
  return docs
37
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38
 
39
  class VectorDB:
40
  def __init__(self, config, logger=None):
@@ -160,7 +180,7 @@ class VectorDB:
160
  + self.config["embedding_options"]["model"],
161
  ),
162
  self.embedding_model,
163
- allow_dangerous_deserialization=True,
164
  )
165
  elif self.db_option == "Chroma":
166
  self.vector_db = Chroma(
 
5
  from langchain.schema.vectorstore import VectorStoreRetriever
6
  from langchain.callbacks.manager import CallbackManagerForRetrieverRun
7
  from langchain.schema.document import Document
8
+ from langchain_core.callbacks import AsyncCallbackManagerForRetrieverRun
9
 
10
  try:
11
  from modules.embedding_model_loader import EmbeddingModelLoader
 
18
  from constants import *
19
  from helpers import *
20
 
21
+ from typing import List
22
+
23
 
24
  class VectorDBScore(VectorStoreRetriever):
25
+
26
  # See https://github.com/langchain-ai/langchain/blob/61dd92f8215daef3d9cf1734b0d1f8c70c1571c3/libs/langchain/langchain/vectorstores/base.py#L500
27
  def _get_relevant_documents(
28
+ self, query: str, *, run_manager: CallbackManagerForRetrieverRun
29
+ ) -> List[Document]:
30
  docs_and_similarities = (
31
  self.vectorstore.similarity_search_with_relevance_scores(
32
  query, **self.search_kwargs
 
39
  docs = [doc for doc, _ in docs_and_similarities]
40
  return docs
41
 
42
+ async def _aget_relevant_documents(
43
+ self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
44
+ ) -> List[Document]:
45
+ docs_and_similarities = (
46
+ self.vectorstore.similarity_search_with_relevance_scores(
47
+ query, **self.search_kwargs
48
+ )
49
+ )
50
+ # Make the score part of the document metadata
51
+ for doc, similarity in docs_and_similarities:
52
+ doc.metadata["score"] = similarity
53
+
54
+ docs = [doc for doc, _ in docs_and_similarities]
55
+ return docs
56
+
57
+
58
 
59
  class VectorDB:
60
  def __init__(self, config, logger=None):
 
180
  + self.config["embedding_options"]["model"],
181
  ),
182
  self.embedding_model,
183
+ # allow_dangerous_deserialization=True, <- unexpected keyword argument to load_local
184
  )
185
  elif self.db_option == "Chroma":
186
  self.vector_db = Chroma(