Skip to content

Commit

Permalink
fix: reformatting...
Browse files Browse the repository at this point in the history
  • Loading branch information
Um Changyong committed Dec 16, 2024
1 parent 3188e44 commit 8fb2be1
Show file tree
Hide file tree
Showing 6 changed files with 23 additions and 10 deletions.
5 changes: 4 additions & 1 deletion autorag/nodes/retrieval/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,10 @@ def cast_queries(queries: Union[str, List[str]]) -> List[str]:


def evenly_distribute_passages(
ids: List[List[str]], scores: List[List[float]], contents: [List[List[str]]], top_k: int
ids: List[List[str]],
scores: List[List[float]],
contents: [List[List[str]]],
top_k: int,
) -> Tuple[List[str], List[float], List[str]]:
assert len(ids) == len(scores), "ids and scores must have same length."
query_cnt = len(ids)
Expand Down
2 changes: 1 addition & 1 deletion autorag/nodes/retrieval/vectordb.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
flatten_apply,
result_to_dataframe,
pop_params,
fetch_contents,
# fetch_contents,
empty_cuda_cache,
convert_inputs_to_list,
make_batch,
Expand Down
4 changes: 3 additions & 1 deletion autorag/vectordb/chroma.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,9 @@ async def add(self, ids: List[str], texts: List[str]):
texts = self.truncated_inputs(texts)
text_embeddings = await self.embedding.aget_text_embedding_batch(texts)
if isinstance(self.collection, AsyncCollection):
await self.collection.add(ids=ids, embeddings=text_embeddings, documents=texts)
await self.collection.add(
ids=ids, embeddings=text_embeddings, documents=texts
)
else:
self.collection.add(ids=ids, embeddings=text_embeddings, documents=texts)

Expand Down
2 changes: 1 addition & 1 deletion autorag/vectordb/milvus.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ def __init__(
field = FieldSchema(
name="vector", dtype=DataType.FLOAT_VECTOR, dim=dimension
)
content = FieldSchema(
content = FieldSchema(
name="content", dtype=DataType.VARCHAR, max_length=65535
)
schema = CollectionSchema(fields=[pk, field, content])
Expand Down
4 changes: 2 additions & 2 deletions autorag/vectordb/pinecone.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,8 +62,8 @@ async def add(self, ids: List[str], texts: List[str]):

metadatas = [{} for _ in texts]
for metadata, text in zip(metadatas, texts):
metadata[self.text_key] = text
metadata[self.text_key] = text

vector_tuples = list(zip(ids, text_embeddings, metadatas))
batch_vectors = make_batch(vector_tuples, self.ingest_batch)

Expand Down
16 changes: 12 additions & 4 deletions autorag/vectordb/qdrant.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,10 +86,13 @@ async def add(self, ids: List[str], texts: List[str]):

metadatas = [{} for _ in texts]
for metadata, text in zip(metadatas, texts):
metadata[self.text_key] = text
metadata[self.text_key] = text

points = list(
map(lambda x: PointStruct(id=x[0], vector=x[1], payload=x[2]), zip(ids, text_embeddings, metadatas))
map(
lambda x: PointStruct(id=x[0], vector=x[1], payload=x[2]),
zip(ids, text_embeddings, metadatas),
)
)

self.client.upload_points(
Expand Down Expand Up @@ -133,7 +136,9 @@ async def query(

search_queries = list(
map(
lambda x: SearchRequest(vector=x, limit=top_k, with_vector=True, with_payload=True),
lambda x: SearchRequest(
vector=x, limit=top_k, with_vector=True, with_payload=True
),
query_embeddings,
)
)
Expand All @@ -145,7 +150,10 @@ async def query(
# Extract IDs and distances
ids = [[str(hit.id) for hit in result] for result in search_result]
scores = [[hit.score for hit in result] for result in search_result]
contents = [[hit.payload.get(self.text_key) for hit in result] for result in search_result]
contents = [
[hit.payload.get(self.text_key) for hit in result]
for result in search_result
]

return ids, scores, contents

Expand Down

0 comments on commit 8fb2be1

Please sign in to comment.