init commit: v0.1.0

query.py

@@ -0,0 +1,43 @@
+
+import pickle
+
+import numpy as np
+import faiss
+from dotenv import load_dotenv
+from openai import OpenAI
+
+load_dotenv()
+client = OpenAI()
+
+MODEL = "text-embedding-3-small"
+
+
+def embed(text):
+    res = client.embeddings.create(
+        model=MODEL,
+        input=text
+    )
+    return res.data[0].embedding
+
+
+index = faiss.read_index("index/index.faiss")
+
+with open("index/meta.pkl", "rb") as f:
+    chunks, meta = pickle.load(f)
+
+
+def search(query, k=5):
+    qvec = np.array([embed(query)]).astype("float32")
+
+    distances, ids = index.search(qvec, k)
+
+    for i in ids[0]:
+        print("\n----")
+        print(meta[i]["source"])
+        print(chunks[i][:500])
+
+
+if __name__ == "__main__":
+    import sys
+    query = " ".join(sys.argv[1:])
+    search(query)