Upload index_miriad_to_qdrant.py

index_miriad_to_qdrant.py

@@ -0,0 +1,73 @@
+# index_miriad_to_qdrant.py
+
+from datasets import load_dataset
+from qdrant_client import QdrantClient, models
+from dotenv import load_dotenv
+import os
+
+load_dotenv() 
+
+# Connect to Qdrant Cloud
+client = QdrantClient(
+    url=os.environ.get("QDRANT_CLOUD_URL"),
+    api_key=os.environ.get("QDRANT_API_KEY"),
+    timeout=60.0,
+    prefer_grpc=True
+)
+
+# Load MIRIAD dataset (sample for demo)
+ds = load_dataset("miriad/miriad-5.8M", split="train").select(range(100000))
+
+dense_documents = [
+    models.Document(text=doc, model="BAAI/bge-small-en")
+    for doc in ds['passage_text']
+]
+
+colbert_documents = [
+    models.Document(text=doc, model="colbert-ir/colbertv2.0")
+    for doc in ds['passage_text']
+]
+
+collection_name = "medical_chat_bot"
+
+# Create collection
+if not client.collection_exists(collection_name):
+    client.recreate_collection(
+        collection_name=collection_name,
+        vectors_config={
+            "dense": models.VectorParams(size=384, distance=models.Distance.COSINE),
+            "colbert": models.VectorParams(
+                size=128,
+                distance=models.Distance.COSINE,
+                multivector_config=models.MultiVectorConfig(
+                    comparator=models.MultiVectorComparator.MAX_SIM
+                ),
+                hnsw_config=models.HnswConfigDiff(m=0)  # reranker: no indexing
+            )
+        }
+    )
+
+# Batch upload in chunks
+BATCH_SIZE = 3
+points_batch = []
+
+for i in range(len(ds['passage_text'])):
+    point = models.PointStruct(
+        id=i,
+        vector={
+            "dense": dense_documents[i],
+            "colbert": colbert_documents[i]
+        },
+        payload={"passage_text": ds['passage_text'][i], "paper_id": ds['paper_id'][i]}
+    )
+    points_batch.append(point)
+
+    if len(points_batch) == BATCH_SIZE:
+        client.upsert(collection_name=collection_name, points=points_batch)
+        print(f"Uploaded batch ending at index {i}")
+        points_batch = []
+
+# Final flush
+if points_batch:
+    client.upsert(collection_name=collection_name, points=points_batch)
+    print("Uploaded final batch.")