Confissões de um PyMongoArrowholic: usando Atlas Vector Search e PyMongoArrow para pesquisar semanticamente itens de moda de luxo

1
def get_embedding(text):
2
       embedding = openai.embeddings.create(input=text, model=EMBEDDING_MODEL).data[0].embedding
3
       return embedding

1
!pip install openai
2
import pandas as pd
3
import openai

1
# we need our openai secret key
2
openai.api_key = 'OPENAI-APIKEY'
3
4
5
# we are going to use this embedding model for text
6
EMBEDDING_MODEL = "text-embedding-3-small"
7
8
9
# this is the embedding function we will use to create our embeddings for the descriptions
10
def get_embedding(text):
11
       embedding = openai.embeddings.create(input=text, model=EMBEDDING_MODEL).data[0].embedding
12
       return embedding

1
# we want to read in our file that we just uploaded
2
df = pd.read_csv('/content/net-a-porter.csv')

1
# drop the 'type' column
2
df.drop(columns=['type'], inplace=True)

1
# this is just saying to drop null values from our specific subsets, we want 'brand', price_usd', and 'description'
2
df.dropna(subset=['brand', 'price_usd', 'description'], inplace=True)

1
# cut our entire file down to just 100 rows so that we don't spend a million dollars embedding the file
2
df = df.head(100)

1
# this is creating a new column with the embeddings
2
df["description_embedding"] = df['description'].apply(get_embedding)
3
4
5
# this is going to give us the first twenty rows from our file
6
print(df.head(20))

1
!pip install pymongo pymongoarrow pyarrow

1
from pymongo import MongoClient
2
from pymongoarrow.api import write
3
import pyarrow as pa
4
5
6
# I want to sort by most expensive item to least expensive
7
df = df.sort_values(by=['price_usd'], ascending=False)
8
9
10
# this is your connection to your cluster
11
connection_string = "MONGODB-CONNECTION-STRING"
12
client = MongoClient(connection_string)
13
14
15
# you can name your database and collection anything you like
16
database = client['net-a-porter']
17
collection = database['average_prices_descending']
18
19
20
# in order to save our data, we need to first convert our Pandas DataFrame to an Arrow Table using pyarrow
21
arrow_table = pa.Table.from_pandas(df)
22
23
24
write(collection, arrow_table)
25
26
27
print("Successful")
28
print(arrow_table)

1
{
2
 "fields": [
3
   {
4
     "numDimensions": 1536,
5
     "path": "description_embedding",
6
     "similarity": "cosine",
7
     "type": "vector"
8
   }
9
 ]
10
}

1
# my query
2
query_description = "summer"
3
4
5
# we need to embed the query as well, since our documents are embedded
6
query_vector = get_embedding(query_description)

1
# write the aggregation pipeline
2
pipeline = [
3
   {
4
       '$vectorSearch': {
5
           'index': 'vector_index', 
6
           'path': 'description_embedding',
7
           'queryVector': query_vector,
8
           # I only had 100 rows saved so that it's easier to use OpenAI
9
           'numCandidates': 100,
10
           'limit': 5
11
       }
12
   },
13
   {
14
       '$project': {
15
           # i do not need to see the ID, but I do want to see my other columns.
16
           '_id': 0,
17
           'brand': 1,
18
           'description': 1,
19
           'price_usd': 1,
20
           'score': {
21
               '$meta': 'vectorSearchScore'
22
               }
23
       }
24
   },
25
   {
26
       '$sort': {
27
           'price_usd': -1  # sort by most expensive to least expensive
28
       }
29
   }
30
]

1
# the pipeline is run on this database and collection
2
database = client['net-a-porter']
3
collection = database['average_prices_descending']
4
result = collection.aggregate(pipeline)
5
6
7
8
9
for clothing in result:
10
   print(clothing)