diff --git a/lightrag/llm/sentence_transformers.py b/lightrag/llm/sentence_transformers.py index b5858a10..8a8ef5cd 100644 --- a/lightrag/llm/sentence_transformers.py +++ b/lightrag/llm/sentence_transformers.py @@ -13,7 +13,9 @@ from sentence_transformers import SentenceTransformer async def sentence_transformers_embed( texts: list[str], model: SentenceTransformer ) -> np.ndarray: - async def inner_encode(texts: list[str], model: SentenceTransformer, embedding_dim: int = 1024): + async def inner_encode( + texts: list[str], model: SentenceTransformer, embedding_dim: int = 1024 + ): return model.encode( texts, truncate_dim=embedding_dim, @@ -21,6 +23,10 @@ async def sentence_transformers_embed( convert_to_tensor=False, show_progress_bar=False, ) - - embedding_func = EmbeddingFunc(embedding_dim=model.get_sentence_embedding_dimension(), func=inner_encode, max_token_size=model.get_max_seq_length()) + + embedding_func = EmbeddingFunc( + embedding_dim=model.get_sentence_embedding_dimension(), + func=inner_encode, + max_token_size=model.get_max_seq_length(), + ) return await embedding_func(texts, model=model) diff --git a/lightrag/rerank.py b/lightrag/rerank.py index cf508a8a..4d5ba932 100644 --- a/lightrag/rerank.py +++ b/lightrag/rerank.py @@ -384,7 +384,7 @@ if __name__ == "__main__": print(f"Document: {docs[item['index']]}") except Exception as e: print(f"Aliyun Error: {e}") - + # Test Sentence Transformers rerank try: print("\n=== Sentence Transformers Rerank ===")