43 lines
759 B
Python
43 lines
759 B
Python
|
|
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) |