gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Reformat

Browse source
This commit is contained in:
Tom Aarsen 2025-11-18 12:19:55 +01:00
parent 835edda6fc
commit 23e7ffbe1c
2 changed files with 10 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
lightrag/llm/sentence_transformers.py
View file

@ -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)

2
lightrag/rerank.py
View file

@ -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 ===")