EcomShoppingBuddy / preprocess.py
RomyMy's picture
fix code 7
e8d1947
import numpy as np
import pandas as pd
import redis
from sentence_transformers import SentenceTransformer
from constants import (
DATA_PATH,
MAX_TEXT_LENGTH,
NUMBER_PRODUCTS,
TEXT_EMBEDDING_DIMENSION,
)
from database import create_redis
from utils import create_flat_index, load_vectors
def auto_truncate(text: str):
return text[0:MAX_TEXT_LENGTH]
def data_preprocessing_and_loading():
pool = create_redis()
redis_conn = redis.Redis(connection_pool=pool)
data = pd.read_csv(
DATA_PATH,
converters={"bullet_point": auto_truncate, "item_keywords": auto_truncate, "item_name": auto_truncate},
)
data["primary_key"] = data["item_id"] + "-" + data["domain_name"]
data.drop(columns=["item_id", "domain_name"], inplace=True)
data["item_keywords"].replace("", np.nan, inplace=True)
data.dropna(subset=["item_keywords"], inplace=True)
data.reset_index(drop=True, inplace=True)
data_metadata = data.head(NUMBER_PRODUCTS).to_dict(orient="index")
# generate embeddings (vectors) for the item keywords
embedding_model = SentenceTransformer("sentence-transformers/all-distilroberta-v1")
# get the item keywords attribute for each product and encode them into vector embeddings
item_keywords = [data_metadata[i]["item_keywords"] for i in data_metadata.keys()]
item_keywords_vectors = [embedding_model.encode(item) for item in item_keywords]
# flush all data
redis_conn.flushall()
# create flat index & load vectors
create_flat_index(redis_conn, NUMBER_PRODUCTS, TEXT_EMBEDDING_DIMENSION, "COSINE")
load_vectors(redis_conn, data_metadata, item_keywords_vectors)
if __name__ == "__main__":
data_preprocessing_and_loading()