github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-17 18:23:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adds embeddings filter option to return scores in state (#12489)

Browse Source 
CC @baskaryan @assafelovic
This commit is contained in:
Jacob Lee 2023-11-08 17:50:06 -08:00 committed by GitHub
parent 18601bd4c8
commit 76283e9625
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
libs/langchain/langchain/retrievers/document_compressors/embeddings_filter.py
View File

@ -67,4 +67,6 @@ class EmbeddingsFilter(BaseDocumentCompressor):
similarity[included_idxs] > self.similarity_threshold
)
included_idxs = included_idxs[similar_enough]
for i in included_idxs:
stateful_documents[i].state["query_similarity_score"] = similarity[i]
return [stateful_documents[i] for i in included_idxs]

4
libs/langchain/tests/integration_tests/retrievers/document_compressors/test_embeddings_filter.py
View File

@ -35,7 +35,9 @@ def test_embeddings_filter_with_state() -> None:
state = {"embedded_doc": np.zeros(len(embedded_query))}
docs = [_DocumentWithState(page_content=t, state=state) for t in texts]
docs[-1].state = {"embedded_doc": embedded_query}
relevant_filter = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.75)
relevant_filter = EmbeddingsFilter(
embeddings=embeddings, similarity_threshold=0.75, return_similarity_scores=True
)
actual = relevant_filter.compress_documents(docs, query)
assert len(actual) == 1
assert texts[-1] == actual[0].page_content