wehub-resource-sync
80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
import pandas as pd
|
|
from typing import List, Dict
|
|
from autoagent.memory.rag_memory import Memory, Reranker
|
|
import json
|
|
import math
|
|
import os
|
|
from litellm import completion
|
|
from autoagent.memory.utils import chunking_by_token_size
|
|
|
|
class TextMemory(Memory):
|
|
def __init__(
|
|
self,
|
|
project_path: str,
|
|
db_name: str = '.text_table',
|
|
platform: str = 'OpenAI',
|
|
api_key: str = None,
|
|
embedding_model: str = "text-embedding-3-small",
|
|
):
|
|
super().__init__(
|
|
project_path=project_path,
|
|
db_name=db_name,
|
|
platform=platform,
|
|
api_key=api_key,
|
|
embedding_model=embedding_model
|
|
)
|
|
self.collection_name = 'text_memory'
|
|
|
|
def add_text_content(self, paper_content: str, batch_size: int = 100, collection = None):
|
|
assert collection is not None, "Collection is required. Should be the path of the paper."
|
|
queries = []
|
|
content_chunks = chunking_by_token_size(paper_content, max_token_size=4096)
|
|
|
|
idx_list = ["chunk_" + str(chunk['chunk_order_index']) for chunk in content_chunks]
|
|
for chunk in content_chunks:
|
|
query = {
|
|
'query': chunk['content'],
|
|
'response': chunk['content']
|
|
}
|
|
queries.append(query)
|
|
|
|
# self.add_query(queries, collection=collection)
|
|
print(f'Adding {len(queries)} queries to {collection} with batch size {batch_size}')
|
|
num_batches = math.ceil(len(queries) / batch_size)
|
|
|
|
for i in range(num_batches):
|
|
start_idx = i * batch_size
|
|
end_idx = min((i + 1) * batch_size, len(queries))
|
|
batch_queries = queries[start_idx:end_idx]
|
|
batch_idx = idx_list[start_idx:end_idx]
|
|
|
|
# Add the current batch of queries
|
|
self.add_query(batch_queries, collection=collection, idx=batch_idx)
|
|
|
|
print(f"Batch {i+1}/{num_batches} added")
|
|
|
|
def query_text_content(
|
|
self,
|
|
query_text: str,
|
|
collection: str = None,
|
|
n_results: int = 5
|
|
) -> List[str]:
|
|
"""
|
|
Query the table and return the results
|
|
"""
|
|
assert collection is not None, "Collection is required. Should be the path of the paper."
|
|
results = self.query([query_text], collection=collection, n_results=n_results)
|
|
|
|
metadata_results = results['metadatas'][0]
|
|
|
|
results = [item['response'] for item in metadata_results]
|
|
return results
|
|
|
|
def peek_table(self, collection: str = None, n_results: int = 20) -> pd.DataFrame:
|
|
"""
|
|
Peek at the data in the table
|
|
"""
|
|
assert collection is not None, "Collection is required. Should be the path of the paper."
|
|
raw_results = self.peek(collection=collection, n_results=n_results)
|
|
results = [item['response'] for item in raw_results['metadatas']]
|
|
return results |