Spaces:
Configuration error
Configuration error
File size: 6,201 Bytes
db10232 | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | import sys
import types
import datasets
# Mock sentence_transformers trainer, training_args, cross_encoder, and sparse_encoder to bypass Trainer imports
sys.modules['sentence_transformers.trainer'] = types.ModuleType('sentence_transformers.trainer')
sys.modules['sentence_transformers.trainer'].SentenceTransformerTrainer = None
sys.modules['sentence_transformers.training_args'] = types.ModuleType('sentence_transformers.training_args')
sys.modules['sentence_transformers.training_args'].SentenceTransformerTrainingArguments = None
sys.modules['sentence_transformers.training_args'].BatchSamplers = None
sys.modules['sentence_transformers.training_args'].MultiDatasetBatchSamplers = None
sys.modules['sentence_transformers.sparse_encoder'] = types.ModuleType('sentence_transformers.sparse_encoder')
sys.modules['sentence_transformers.sparse_encoder'].SparseEncoder = None
sys.modules['sentence_transformers.sparse_encoder'].SparseEncoderModelCardData = None
sys.modules['sentence_transformers.sparse_encoder'].SparseEncoderTrainer = None
sys.modules['sentence_transformers.sparse_encoder'].SparseEncoderTrainingArguments = None
sys.modules['sentence_transformers.cross_encoder'] = types.ModuleType('sentence_transformers.cross_encoder')
sys.modules['sentence_transformers.cross_encoder'].CrossEncoder = None
sys.modules['sentence_transformers.cross_encoder'].CrossEncoderModelCardData = None
sys.modules['sentence_transformers.cross_encoder'].CrossEncoderTrainer = None
sys.modules['sentence_transformers.cross_encoder'].CrossEncoderTrainingArguments = None
import logging
import os
import uuid
from pathlib import Path
from qdrant_client import QdrantClient, models
from sentence_transformers import SentenceTransformer
from dotenv import load_dotenv
# Add project root to path to resolve custom modules correctly
PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.append(str(PROJECT_ROOT))
from ingestion.pipeline import MultimodalIngestionPipeline
from ingestion.config import IngestionSettings
# Load environment variables
load_dotenv()
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
logger = logging.getLogger(__name__)
COLLECTION_NAME = "conversational_rag"
QDRANT_PATH = PROJECT_ROOT / "qdrant_db"
def _stable_chunk_id(chunk_content: str) -> str:
return str(uuid.uuid5(uuid.NAMESPACE_DNS, str(chunk_content)))
def main():
pdf_path = PROJECT_ROOT / "Data" / "Pdf" / "World Development Report 2025.pdf"
logger.info(f"Starting text-only PDF deployment to Qdrant for {pdf_path.name}...")
# 1. Initialize pipeline with figure extraction and vision disabled (text-only)
settings = IngestionSettings(
extract_figures=False,
use_vision=False
)
pipeline = MultimodalIngestionPipeline(settings)
# 2. Ingest the PDF to get text chunks
logger.info("Ingesting PDF text chunks...")
result = pipeline.ingest_sync(pdf_path)
chunks = result.chunks
logger.info(f"Successfully parsed {len(chunks)} PDF text chunks.")
# 3. Setup Qdrant Client (without recreating the collection)
client = QdrantClient(path=str(QDRANT_PATH))
if not client.collection_exists(COLLECTION_NAME):
logger.error(f"Collection '{COLLECTION_NAME}' does not exist! Please run CSV deployment first.")
sys.exit(1)
# 4. Load Local SentenceTransformer Model
logger.info("Loading local SentenceTransformer('all-MiniLM-L6-v2')...")
model = SentenceTransformer('all-MiniLM-L6-v2')
logger.info("Local SentenceTransformer model loaded successfully.")
# 5. Generate Embeddings for all text chunks
texts = [str(chunk.text).strip() for chunk in chunks]
logger.info(f"Generating embeddings for {len(texts)} chunks locally...")
embeddings = model.encode(texts, convert_to_numpy=True)
logger.info("Embedding generation completed.")
# 6. Construct Qdrant points
points = []
for i, chunk in enumerate(chunks):
text = str(chunk.text).strip()
metadata = dict(chunk.metadata or {})
chunk_id = str(metadata.get("chunk_id") or _stable_chunk_id(text))
payload = {
"text": text,
"page_content": text,
"source": str(metadata.get("source_file") or pdf_path.name),
"image_path": metadata.get("image_path"),
"contains_chart": bool(metadata.get("contains_chart")),
"contains_table": bool(metadata.get("contains_table")),
"contains_figure": bool(metadata.get("contains_figure")),
"contains_image": bool(metadata.get("contains_image")),
"contains_csv": bool(metadata.get("contains_csv")),
"metadata": metadata,
}
# Keep any asset specific metadata mappings same
for key in ["row_id", "columns"]:
if metadata.get(key) not in ("", None, [], {}):
payload[key] = metadata[key]
dense_vector = [float(val) for val in embeddings[i]]
points.append(
models.PointStruct(
id=str(uuid.uuid5(uuid.NAMESPACE_URL, chunk_id)),
vector={
"dense": dense_vector,
"sparse": models.SparseVector(indices=[], values=[]),
},
payload=payload,
)
)
# 7. Upsert points in batches (appending to the existing collection)
batch_size = 64
uploaded = 0
total_points = len(points)
for start in range(0, total_points, batch_size):
batch = points[start : start + batch_size]
client.upsert(collection_name=COLLECTION_NAME, points=batch, wait=True)
uploaded += len(batch)
logger.info(f"Upserted {uploaded}/{total_points} points.")
# Retrieve final exact count of all points (including CSV)
count = client.count(collection_name=COLLECTION_NAME, exact=True).count
client.close()
logger.info("PDF text-only deployment completed. Total points in collection '%s': %d", COLLECTION_NAME, count)
print(f"\nPDF Text-only Ingestion Completed.")
print(f"Upserted PDF Points: {uploaded}")
print(f"Qdrant Exact Count (Total): {count}")
if __name__ == "__main__":
main()
|