Get API Key and host URL from env vars

.gitignore

@@ -1,3 +1,3 @@
 __pycache__
 *.npy
-*.jsonl
+*.json

create_index.py

@@ -1,18 +1,46 @@
+import argparse
+import os
 from encoder import encode
 
 import numpy as np
 import tqdm
 from datasets import load_dataset
 
-batch_size = 64
 
-dataset = load_dataset('wiki_qa', split='train+test+validation')
+def get_embeddings_and_payload(dataset, batch_size):
+    answer_embeddings = np.zeros(shape=(len(dataset), 384), dtype=np.float32)
+    answer_texts = []
 
-answer_embeddings = np.zeros(shape=(len(dataset), 384), dtype=np.float32)
+    for i in tqdm.tqdm(range(0, len(dataset), batch_size), desc='Extracting vectors'):
+        batch_answers = dataset['answer'][i:i+batch_size]
+        answer_texts.extend([{"text": answer in batch_answers}])
 
-for i in tqdm.tqdm(range(0, len(dataset), batch_size), desc='Extracting vectors'):
-    batch_answers = dataset['answer'][i:i+batch_size]
-    batch_embeddings = encode(batch_answers)
-    answer_embeddings[i:i+batch_size] = batch_embeddings
+        batch_embeddings = encode(batch_answers)
+        answer_embeddings[i:i+batch_size] = batch_embeddings
 
-np.save('answer_embeddings.npy', answer_embeddings)
+    return answer_embeddings, answer_texts
+
+
+def main(args):
+    if args.host_url is None:
+        print("You must specify a host URL either by passing it as a command line argument with `--host-url` or as an environment variable named `QDRANT_HOST_URL`")
+        exit()
+    if args.api_key is None:
+        print("You must specify an API key  either by passing it as a command line argument with `--api-key` or as an environment variable named `QDRANT_API_KEY`")
+        exit()
+
+    dataset = load_dataset('wiki_qa', split='train+test+validation')
+    embeddings, payload_data = get_embeddings_and_payload(
+        dataset, args.batch_size)
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser(prog="Create index on on Qdrant Cloud")
+    ap.add_argument("-b", "--batch-size", type=int, default=64,
+                    help="Batch size when extracting vectors")
+    ap.add_argument("-u", "--host-url", type=str,
+                    default=os.environ['QDRANT_HOST_URL', None], help="Host URL from Qdrant Cloud")
+    ap.add_argument("-k", "--api-key", type=str,
+                    default=os.environ['QDRANT_API_KEY', None], help="API key from Qdrant Cloud")
+    args = ap.parse_args()
+    main(args)