pricer-agent/vector_db.ipynb
2025-06-10 12:03:09 +05:30

349 lines
15 KiB
Plaintext

{
"cells": [
{
"cell_type": "code",
"execution_count": 15,
"id": "4263f6ec",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import re\n",
"import math\n",
"import json\n",
"from tqdm import tqdm\n",
"import random\n",
"from dotenv import load_dotenv\n",
"from huggingface_hub import login\n",
"import numpy as np\n",
"import pickle\n",
"from sentence_transformers import SentenceTransformer\n",
"from datasets import load_dataset\n",
"import chromadb\n",
"from items import Item\n",
"from sklearn.manifold import TSNE\n",
"import plotly.graph_objects as go"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "b91d1bf7",
"metadata": {},
"outputs": [],
"source": [
"with open('train.pkl', 'rb') as file:\n",
" train = pickle.load(file)"
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "2b7a04b6",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'How much does this cost to the nearest dollar?\\n\\nDelphi FG0166 Fuel Pump Module\\nDelphi brings 80 years of OE Heritage into each Delphi pump, ensuring quality and fitment for each Delphi part. Part is validated, tested and matched to the right vehicle application Delphi brings 80 years of OE Heritage into each Delphi assembly, ensuring quality and fitment for each Delphi part Always be sure to check and clean fuel tank to avoid unnecessary returns Rigorous OE-testing ensures the pump can withstand extreme temperatures Brand Delphi, Fit Type Vehicle Specific Fit, Dimensions LxWxH 19.7 x 7.7 x 5.1 inches, Weight 2.2 Pounds, Auto Part Position Unknown, Operation Mode Mechanical, Manufacturer Delphi, Model FUEL PUMP, Dimensions 19.7\\n\\nPrice is $227.00'"
]
},
"execution_count": 17,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"train[0].prompt"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "e9d7cb72",
"metadata": {},
"outputs": [],
"source": [
"DB = \"products_vectorstore\""
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "40bc2796",
"metadata": {},
"outputs": [],
"source": [
"client = chromadb.PersistentClient(path=DB)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e9303c73",
"metadata": {},
"outputs": [],
"source": [
"# Check if the collection exists and delete it if it does\n",
"collection_name = \"products\"\n",
"\n",
"# For old versions of Chroma, use this line instead of the subsequent one\n",
"# existing_collection_names = [collection.name for collection in client.list_collections()]\n",
"existing_collection_names = client.list_collections()\n",
"\n",
"if collection_name in existing_collection_names:\n",
" client.delete_collection(collection_name)\n",
" print(f\"Deleted existing collection: {collection_name}\")\n",
"\n",
"collection = client.create_collection(collection_name)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "586dd619",
"metadata": {},
"outputs": [],
"source": [
"model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')"
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "337adc93",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array([-9.46715921e-02, 4.27619852e-02, 5.51620908e-02, -5.10982354e-04,\n",
" 1.16203353e-02, -6.80130422e-02, 2.76406445e-02, 6.06974810e-02,\n",
" 2.88530681e-02, -1.74128301e-02, -4.94346023e-02, 2.30993200e-02,\n",
" -1.28614539e-02, -4.31402288e-02, 2.17510201e-02, 4.26548794e-02,\n",
" 5.10500222e-02, -7.79727176e-02, -1.23247258e-01, 3.67455557e-02,\n",
" 4.54121176e-03, 9.47937518e-02, -5.53098992e-02, 1.70641821e-02,\n",
" -2.92872526e-02, -4.47125323e-02, 2.06785351e-02, 6.39319792e-02,\n",
" 2.27427743e-02, 4.87789363e-02, -2.33508484e-03, 4.72859591e-02,\n",
" -2.86258645e-02, 2.30625179e-02, 2.45130546e-02, 3.95680927e-02,\n",
" -4.33176570e-02, -1.02316678e-01, 2.79876147e-03, 2.39304081e-02,\n",
" 1.61556378e-02, -8.99088103e-03, 2.07255036e-02, 6.40122816e-02,\n",
" 6.89178854e-02, -6.98361173e-02, 2.89762835e-03, -8.10989514e-02,\n",
" 1.71122309e-02, 2.50655250e-03, -1.06529057e-01, -4.87733409e-02,\n",
" -1.67761575e-02, -2.28662733e-02, 1.14816569e-01, 4.87412885e-02,\n",
" -1.64961927e-02, -6.90832511e-02, 1.13612078e-01, -7.18486458e-02,\n",
" -9.01571363e-02, 3.93730397e-06, -8.66768882e-02, -4.05915864e-02,\n",
" 3.71316932e-02, -1.77618898e-02, -5.57463802e-02, -4.57097068e-02,\n",
" -5.43141440e-02, -4.00611497e-02, -4.60227989e-02, 2.82194242e-02,\n",
" -2.33606212e-02, 1.44406464e-02, -1.52777331e-02, -4.34328392e-02,\n",
" 6.81274012e-02, 8.21894854e-02, 7.83894584e-03, -2.85973307e-02,\n",
" 6.14309609e-02, -4.92684729e-02, 3.27059231e-03, -2.72683222e-02,\n",
" -4.59346585e-02, -2.60167439e-02, 5.70650063e-02, -5.86422451e-04,\n",
" -4.37076800e-02, -2.47870252e-04, -8.73550028e-02, 3.48776542e-02,\n",
" 1.39328036e-02, -1.60043947e-02, 2.86958888e-02, -9.48595777e-02,\n",
" 9.40611809e-02, 5.92685640e-02, -8.65013823e-02, 1.45011336e-01,\n",
" 2.35388372e-02, 3.43324915e-02, 2.51512363e-04, 7.15541244e-02,\n",
" -3.12182605e-02, 3.86666618e-02, -2.47745272e-02, 6.52673766e-02,\n",
" -8.28818083e-02, -2.80246828e-02, 9.34542995e-03, -7.85537530e-03,\n",
" 5.30727953e-02, 2.96895131e-02, 3.27329263e-02, 4.94736545e-02,\n",
" 2.52207797e-02, 4.67067957e-02, 5.03403731e-02, -7.23745152e-02,\n",
" 2.54435204e-02, -3.67217027e-02, 1.27570834e-02, 1.46063920e-02,\n",
" 2.11491529e-02, -5.56909665e-02, -9.18515865e-03, -5.63395203e-34,\n",
" 9.58770216e-02, 5.94577529e-02, 5.11444733e-02, 3.36359628e-02,\n",
" -1.34414285e-02, -2.77203768e-02, -3.48436199e-02, 1.80352516e-02,\n",
" -2.53210664e-02, 6.73579378e-03, 5.47839850e-02, -3.60573716e-02,\n",
" -5.20869531e-02, -2.90346015e-02, 4.38963762e-03, 6.50022700e-02,\n",
" 3.07485145e-02, 2.00219769e-02, 1.73045334e-03, 2.96729826e-03,\n",
" 3.40948463e-03, -6.78144172e-02, 3.41305025e-02, 8.37866776e-03,\n",
" 5.39903976e-02, 2.70389207e-02, 7.84119666e-02, -1.30136281e-01,\n",
" 4.84650210e-02, 5.14179468e-02, -7.94673618e-03, 5.57883270e-03,\n",
" -5.31024933e-02, 3.81298959e-02, -3.05512622e-02, -7.69778490e-02,\n",
" 1.20531330e-02, -4.08992879e-02, -8.69358629e-02, 6.38055578e-02,\n",
" 1.68675203e-02, 1.68741110e-03, 6.28894791e-02, -1.67711601e-02,\n",
" 2.15586051e-02, 7.10061751e-04, 2.81031057e-03, -8.89795925e-03,\n",
" -1.80887170e-02, -2.16217209e-02, -5.59149943e-02, 1.78774614e-02,\n",
" -9.27094072e-02, 7.27912458e-03, -1.27753600e-01, -4.86938767e-02,\n",
" 1.45872040e-02, -1.62750706e-02, 6.75623193e-02, 3.87702435e-02,\n",
" 7.23295733e-02, 9.14992169e-02, -9.65292305e-02, 4.84791510e-02,\n",
" -1.06274128e-01, -1.05042122e-02, 8.90350118e-02, -8.07525739e-02,\n",
" 7.87081793e-02, -2.04917695e-02, -5.55080362e-02, -3.31532545e-02,\n",
" -2.14428790e-02, 4.94698957e-02, -7.05119073e-02, 6.63998201e-02,\n",
" 7.39670321e-02, -2.70114886e-03, 1.62262768e-02, -3.98229063e-02,\n",
" 5.94091974e-02, -7.14365626e-03, -3.33479904e-02, 2.30419226e-02,\n",
" 1.87185612e-02, -6.15725555e-02, -8.55416001e-04, -1.05786592e-01,\n",
" -8.43618810e-02, -3.92993018e-02, -3.16446945e-02, 6.60644248e-02,\n",
" 9.41816568e-02, -8.35982710e-02, 9.50875413e-03, 1.25501758e-34,\n",
" 6.38197511e-02, 2.11371575e-02, -1.65899675e-02, 1.88641343e-02,\n",
" -5.57019338e-02, 1.82812882e-03, -1.37587301e-02, 8.16278681e-02,\n",
" -9.13296789e-02, 7.06855804e-02, 6.79991096e-02, -5.44536524e-02,\n",
" 3.80394794e-02, 3.80513002e-03, 1.03689805e-01, 7.32862041e-04,\n",
" 2.95660980e-02, 4.19424325e-02, -1.20444320e-01, 1.24932900e-02,\n",
" -5.53505979e-02, 1.75228491e-02, -2.28164997e-02, -5.79300001e-02,\n",
" 9.42930207e-03, -5.42277237e-03, -3.94948432e-03, 2.82348841e-02,\n",
" -1.28066897e-01, -1.31304078e-02, 7.42957667e-02, -1.74529944e-02,\n",
" -9.72759053e-02, 8.25622305e-03, 2.06900500e-02, -5.29765012e-03,\n",
" -1.37695735e-02, -3.50973941e-02, 1.74978618e-02, -1.76233463e-02,\n",
" -6.50825202e-02, -3.84675600e-02, -8.76396373e-02, 3.21291536e-02,\n",
" 2.55025318e-03, -2.09378786e-02, 5.55309318e-02, 2.57095862e-02,\n",
" -2.94735916e-02, 1.25047630e-02, -6.83466420e-02, -8.00623894e-02,\n",
" -1.46906180e-02, 1.03744986e-02, -8.51863623e-02, -1.10538909e-02,\n",
" 2.14596596e-02, 4.08609137e-02, 3.31647359e-02, -2.76757330e-02,\n",
" -2.01877356e-02, 8.98887776e-03, 3.92048731e-02, 1.15103371e-01,\n",
" 5.50441183e-02, 2.72755288e-02, -1.09526664e-01, -1.72623191e-02,\n",
" 1.33438576e-02, -1.73701961e-02, -5.04375855e-03, -2.00292692e-02,\n",
" 1.16672337e-01, -1.84322931e-02, 3.70627306e-02, 1.60885453e-02,\n",
" 3.48830558e-02, 5.50573654e-02, -6.60796463e-03, 7.06828311e-02,\n",
" 4.07849178e-02, -1.43314470e-02, -2.85443966e-03, 2.74251904e-02,\n",
" -4.26768474e-02, 1.26583306e-02, 3.34343836e-02, 1.62644740e-02,\n",
" 1.19262813e-02, -2.92118434e-02, 2.73977909e-02, 3.44305374e-02,\n",
" 2.52833236e-02, 3.07514146e-02, 3.22558023e-02, -1.74628472e-08,\n",
" -1.52690308e-02, 5.37678273e-03, 1.41246557e-01, 5.08366488e-02,\n",
" 5.32255769e-02, 9.67938229e-02, 4.33674268e-02, -6.48313668e-03,\n",
" 1.58604365e-02, 4.05631550e-02, 6.94984868e-02, 6.04905970e-02,\n",
" -6.26188368e-02, -3.96144390e-02, 1.10648818e-01, 1.67735666e-02,\n",
" -7.68705085e-03, 2.59616226e-02, -5.28793186e-02, -2.22318973e-02,\n",
" 1.74595993e-02, 4.75340039e-02, 3.27674821e-02, -4.59684506e-02,\n",
" 2.01770663e-02, -1.60875786e-02, -1.58614144e-02, -1.66657437e-02,\n",
" -3.05246655e-02, -3.87907699e-02, -1.27654579e-02, 6.57610893e-02,\n",
" -2.22502407e-02, -9.44991410e-03, 2.32080836e-02, 2.66038626e-02,\n",
" 2.14203075e-02, -7.54577760e-03, 8.84752721e-02, -9.43514556e-02,\n",
" -5.74871078e-02, -7.77098387e-02, 1.95802352e-03, -1.50347678e-02,\n",
" -8.08493048e-03, 1.88217331e-02, 8.42517056e-03, -3.78592350e-02,\n",
" 1.24534788e-02, -7.94995204e-02, -2.15789191e-02, 1.20276539e-02,\n",
" 1.74870472e-02, 8.74479711e-02, 6.64091185e-02, 3.13737318e-02,\n",
" -1.00628426e-02, 2.07700673e-02, -5.20163700e-02, -8.91336147e-03,\n",
" 1.48542315e-01, -2.51266640e-03, 9.93156359e-02, 2.34929807e-02],\n",
" dtype=float32)"
]
},
"execution_count": 21,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"vector = model.encode([\"Well hi there\"])[0]\n",
"vector"
]
},
{
"cell_type": "code",
"execution_count": 22,
"id": "626add8d",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"def cosine_similarity(a, b):\n",
" return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))\n",
"\n",
"def how_similar(text1, text2):\n",
" vector1, vector2 = model.encode([text1, text2])\n",
" similarity = cosine_similarity(vector1, vector2)\n",
" print(f\"Similarity between {text1} and {text2} is {similarity*100:.1f}%\")"
]
},
{
"cell_type": "code",
"execution_count": 23,
"id": "7f1e7f04",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Similarity between Java and C++ is 50.7%\n",
"Similarity between Java and mug is 25.8%\n",
"Similarity between Cup of Java and mug is 49.3%\n"
]
}
],
"source": [
"how_similar(\"Java\", \"C++\")\n",
"how_similar(\"Java\", \"mug\")\n",
"how_similar(\"Cup of Java\", \"mug\")"
]
},
{
"cell_type": "code",
"execution_count": 24,
"id": "d60e8a2f",
"metadata": {},
"outputs": [],
"source": [
"def description(item):\n",
" text = item.prompt.replace(\"How much does this cost to the nearest dollar?\\n\\n\", \"\")\n",
" return text.split(\"\\n\\nPrice is $\")[0]"
]
},
{
"cell_type": "code",
"execution_count": 25,
"id": "37665268",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'Delphi FG0166 Fuel Pump Module\\nDelphi brings 80 years of OE Heritage into each Delphi pump, ensuring quality and fitment for each Delphi part. Part is validated, tested and matched to the right vehicle application Delphi brings 80 years of OE Heritage into each Delphi assembly, ensuring quality and fitment for each Delphi part Always be sure to check and clean fuel tank to avoid unnecessary returns Rigorous OE-testing ensures the pump can withstand extreme temperatures Brand Delphi, Fit Type Vehicle Specific Fit, Dimensions LxWxH 19.7 x 7.7 x 5.1 inches, Weight 2.2 Pounds, Auto Part Position Unknown, Operation Mode Mechanical, Manufacturer Delphi, Model FUEL PUMP, Dimensions 19.7'"
]
},
"execution_count": 25,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"description(train[0])"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a8592f7f",
"metadata": {},
"outputs": [],
"source": [
"NUMBER_OF_DOCUMENTS = len(train)\n",
"\n",
"# Uncomment if you'd rather not wait for the full 400,000\n",
"# NUMBER_OF_DOCUMENTS = 20000\n",
"\n",
"for i in tqdm(range(0, NUMBER_OF_DOCUMENTS, 1000)):\n",
" documents = [description(item) for item in train[i: i+1000]]\n",
" vectors = model.encode(documents).astype(float).tolist()\n",
" metadatas = [{\"category\": item.category, \"price\": item.price} for item in train[i: i+1000]]\n",
" ids = [f\"doc_{j}\" for j in range(i, i+len(documents))]\n",
" collection.add(\n",
" ids=ids,\n",
" documents=documents,\n",
" embeddings=vectors,\n",
" metadatas=metadatas\n",
" )"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.22"
}
},
"nbformat": 4,
"nbformat_minor": 5
}