refac: lazily load sentence_transformers to reduce start up memory usage

This commit is contained in:
Jun Siang Cheah
2024-07-01 08:13:56 +08:00
parent 17c684369e
commit a48ac6a209
2 changed files with 6 additions and 4 deletions
+2 -2
View File
@@ -442,8 +442,6 @@ from langchain_core.documents import BaseDocumentCompressor, Document
from langchain_core.callbacks import Callbacks
from langchain_core.pydantic_v1 import Extra
from sentence_transformers import util
class RerankCompressor(BaseDocumentCompressor):
embedding_function: Any
@@ -468,6 +466,8 @@ class RerankCompressor(BaseDocumentCompressor):
[(query, doc.page_content) for doc in documents]
)
else:
from sentence_transformers import util
query_embedding = self.embedding_function(query)
document_embedding = self.embedding_function(
[doc.page_content for doc in documents]