SalesAI / product_recommender.py
Zasha1's picture
Update product_recommender.py
c17360f verified
import pandas as pd
from sentence_transformers import SentenceTransformer
import faiss
class ProductRecommender:
def __init__(self, product_data_path):
try:
# Attempt to load the product data CSV
self.data = pd.read_csv(product_data_path)
print("Product data loaded successfully.")
except Exception as e:
print(f"Error loading product data: {e}")
self.data = pd.DataFrame() # Create an empty DataFrame if loading fails
return
try:
# Initialize the sentence transformer model
self.model = SentenceTransformer('all-MiniLM-L6-v2')
print("Model loaded successfully.")
except Exception as e:
print(f"Error loading SentenceTransformer model: {e}")
self.model = None # Set model to None if loading fails
return
try:
# Check if 'product_description' column exists
if 'product_description' not in self.data.columns:
print("Error: 'product_description' column is missing in the data.")
return
# Generate embeddings for the product descriptions
self.embeddings = self.model.encode(self.data['product_description'].tolist())
print(f"Embeddings generated successfully. Shape: {self.embeddings.shape}")
except Exception as e:
print(f"Error generating embeddings: {e}")
self.embeddings = None # Set embeddings to None if generation fails
return
try:
# Initialize FAISS index and add the embeddings
self.index = faiss.IndexFlatL2(self.embeddings.shape[1])
self.index.add(self.embeddings)
print("FAISS index created and embeddings added.")
except Exception as e:
print(f"Error creating FAISS index or adding embeddings: {e}")
self.index = None # Set index to None if creation fails
return
def get_recommendations(self, query, top_n=5):
if self.model is None or self.index is None:
print("Error: Model or FAISS index not initialized. Cannot make recommendations.")
return []
try:
# Generate the embedding for the query
query_embedding = self.model.encode([query])
print(f"Query embedding generated. Shape: {query_embedding.shape}")
except Exception as e:
print(f"Error generating query embedding: {e}")
return []
try:
# Search for top_n recommendations
distances, indices = self.index.search(query_embedding, top_n)
recommendations = []
for i in indices[0]:
recommendations.append(self.data.iloc[i]['product_title'] + ": " + self.data.iloc[i]['product_description'])
print(f"Recommendations generated successfully: {recommendations}")
return recommendations
except Exception as e:
print(f"Error during recommendation search: {e}")
return []