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( async def sentence_transformers_embed(
texts: list[str], model: SentenceTransformer texts: list[str], model: SentenceTransformer
) -> np.ndarray: ) -> 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( return model.encode(
texts, texts,
truncate_dim=embedding_dim, truncate_dim=embedding_dim,
@ -21,6 +23,10 @@ async def sentence_transformers_embed(
convert_to_tensor=False, convert_to_tensor=False,
show_progress_bar=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) 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']]}") print(f"Document: {docs[item['index']]}")
except Exception as e: except Exception as e:
print(f"Aliyun Error: {e}") print(f"Aliyun Error: {e}")
# Test Sentence Transformers rerank # Test Sentence Transformers rerank
try: try:
print("\n=== Sentence Transformers Rerank ===") print("\n=== Sentence Transformers Rerank ===")