Spaces:
Runtime error
Runtime error
File size: 2,542 Bytes
8b6eec6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import os
import weaviate
from llama_index import download_loader
from llama_index.vector_stores import WeaviateVectorStore
from llama_index import VectorStoreIndex, StorageContext
from pathlib import Path
import argparse
def get_pdf_files(base_path, loader):
"""
Get paths to all PDF files in a directory and its subdirectories.
Parameters:
- base_path (str): The path to the starting directory.
Returns:
- list of str: A list of paths to all PDF files found.
"""
pdf_paths = []
# Check if the base path exists and is a directory
if not os.path.exists(base_path):
raise FileNotFoundError(f"The specified base path does not exist: {base_path}")
if not os.path.isdir(base_path):
raise NotADirectoryError(f"The specified base_path is not a directory: {base_path}")
# Loop through all directories and files starting from the base path
for root, dirs, files in os.walk(base_path):
for filename in files:
# If a file has a .pdf extension, add its path to the list
if filename.endswith('.pdf'):
pdf_file = loader.load_data(file=Path(root, filename))
pdf_paths.extend(pdf_file)
return pdf_paths
def main(args):
PDFReader = download_loader("PDFReader")
loader = PDFReader()
documents = get_pdf_files(args.pdf_dir, loader)
client = weaviate.Client(
url=os.environ["WEAVIATE_URL"],
auth_client_secret=weaviate.AuthApiKey(api_key=os.environ["WEAVIATE_API_KEY"]),
additional_headers={
"X-OpenAI-Api-Key": os.environ["OPENAI_API_KEY"]
}
)
# construct vector store
vector_store = WeaviateVectorStore(weaviate_client=client, index_name=args.customer, text_key="content")
# setting up the storage for the embeddings
storage_context = StorageContext.from_defaults(vector_store=vector_store)
# set up the index
index = VectorStoreIndex(documents, storage_context=storage_context)
query_engine = index.as_query_engine()
response = query_engine.query(args.query)
print(response)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Process and query PDF files.')
parser.add_argument('--customer', default='Ausy', help='Customer name')
parser.add_argument('--pdf_dir', default='./data', help='Directory containing PDFs')
parser.add_argument('--query', default='What is CX0 customer exprience office?', help='Query to execute')
args = parser.parse_args()
main(args)
|