mirror of https://github.com/aliasrobotics/cai.git
83 lines
1.8 KiB
Python
83 lines
1.8 KiB
Python
import json
|
|
import os
|
|
|
|
import pandas as pd
|
|
import qdrant_client
|
|
from openai import OpenAI
|
|
from qdrant_client.http import models as rest
|
|
|
|
client = OpenAI()
|
|
GPT_MODEL = "gpt-4o"
|
|
EMBEDDING_MODEL = "text-embedding-3-large"
|
|
|
|
article_list = os.listdir("data")
|
|
|
|
articles = []
|
|
|
|
for x in article_list:
|
|
article_path = "data/" + x
|
|
|
|
# Opening JSON file
|
|
f = open(article_path)
|
|
|
|
# returns JSON object as
|
|
# a dictionary
|
|
data = json.load(f)
|
|
|
|
articles.append(data)
|
|
|
|
# Closing file
|
|
f.close()
|
|
|
|
for i, x in enumerate(articles):
|
|
try:
|
|
embedding = client.embeddings.create(
|
|
model=EMBEDDING_MODEL, input=x["text"])
|
|
articles[i].update({"embedding": embedding.data[0].embedding})
|
|
except Exception as e:
|
|
print(x["title"])
|
|
print(e)
|
|
|
|
qdrant = qdrant_client.QdrantClient(host="localhost")
|
|
qdrant.get_collections()
|
|
|
|
collection_name = "help_center"
|
|
|
|
vector_size = len(articles[0]["embedding"])
|
|
vector_size
|
|
|
|
article_df = pd.DataFrame(articles)
|
|
article_df.head()
|
|
|
|
# Delete the collection if it exists, so we can rewrite it changes to
|
|
# articles were made
|
|
if qdrant.get_collection(collection_name=collection_name):
|
|
qdrant.delete_collection(collection_name=collection_name)
|
|
|
|
# Create Vector DB collection
|
|
qdrant.create_collection(
|
|
collection_name=collection_name,
|
|
vectors_config={
|
|
"article": rest.VectorParams(
|
|
distance=rest.Distance.COSINE,
|
|
size=vector_size,
|
|
)
|
|
},
|
|
)
|
|
|
|
# Populate collection with vectors
|
|
|
|
qdrant.upsert(
|
|
collection_name=collection_name,
|
|
points=[
|
|
rest.PointStruct(
|
|
id=k,
|
|
vector={
|
|
"article": v["embedding"],
|
|
},
|
|
payload=v.to_dict(),
|
|
)
|
|
for k, v in article_df.iterrows()
|
|
],
|
|
)
|