mirror of
https://github.com/SirBlobby/Hoya26.git
synced 2026-02-03 19:24:34 -05:00
63 lines
2.3 KiB
Python
63 lines
2.3 KiB
Python
import os
|
|
import sys
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
|
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
|
|
from src.rag.ingest import process_file
|
|
from src.rag.store import ingest_documents
|
|
from src.mongo.metadata import is_file_processed, log_processed_file
|
|
|
|
def populate_from_dataset(dataset_dir, category=None):
|
|
dataset_path = Path(dataset_dir)
|
|
if not dataset_path.exists():
|
|
print(f"Dataset directory not found: {dataset_dir}")
|
|
return
|
|
|
|
print(f"Scanning {dataset_dir}...")
|
|
if category:
|
|
print(f"Category: {category}")
|
|
|
|
total_chunks = 0
|
|
files_processed = 0
|
|
|
|
for file_path in dataset_path.glob('*'):
|
|
if file_path.is_file() and file_path.suffix.lower() in ['.csv', '.pdf', '.txt', '.xlsx']:
|
|
if is_file_processed(file_path.name):
|
|
print(f"Skipping {file_path.name} (already processed)")
|
|
continue
|
|
|
|
print(f"Processing {file_path.name}...")
|
|
try:
|
|
chunks = process_file(str(file_path))
|
|
if chunks:
|
|
count = ingest_documents(chunks, source_file=file_path.name, category=category)
|
|
print(f" Ingested {count} chunks.")
|
|
if count > 0:
|
|
log_processed_file(file_path.name, category=category, chunk_count=count)
|
|
total_chunks += count
|
|
files_processed += 1
|
|
else:
|
|
print(" No text found/extracted.")
|
|
except Exception as e:
|
|
print(f" Error processing file: {e}")
|
|
|
|
print(f"\nFinished! Processed {files_processed} files. Total chunks ingested: {total_chunks}")
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Populate vector database from dataset files")
|
|
parser.add_argument("--category", "-c", type=str, help="Category to assign to ingested documents")
|
|
parser.add_argument("--dir", "-d", type=str, default=None, help="Dataset directory path")
|
|
args = parser.parse_args()
|
|
|
|
if args.dir:
|
|
dataset_dir = args.dir
|
|
else:
|
|
dataset_dir = os.path.join(os.path.dirname(__file__), '../dataset')
|
|
|
|
populate_from_dataset(dataset_dir, category=args.category)
|