utils/tfidf.py

import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import pairwise_distances
from typing import List, Dict
from utils.config import Config
import os


# Load the dataset (replace with the actual path to your dataset)
dataset_path = Config.read('app', 'dataset')

# Ensure the dataset exists
if not os.path.exists(dataset_path):
    raise FileNotFoundError(f"The dataset file at {dataset_path} was not found.")

# Load the dataset
data = pd.read_pickle(dataset_path)

# Ensure the dataset has the necessary columns: 'asin', 'title', 'brand', 'medium_image_url'
required_columns = ['asin', 'title', 'brand', 'medium_image_url']
for col in required_columns:
    if col not in data.columns:
        raise ValueError(f"Missing required column: {col} in the dataset")

# Set up the vectorizer and fit the model
tfidf_title_vectorizer = TfidfVectorizer(min_df = 0.0)
tfidf_title_features = tfidf_title_vectorizer.fit_transform(data['title'])

# Function to calculate the tf-idf model and return closest matches
def tfidf_model(input_text: str, num_results: int) -> List[Dict]:

    # Transform the input text to the same TF-IDF feature space
    query_vec = tfidf_title_vectorizer.transform([input_text])

    pairwise_dist = pairwise_distances(tfidf_title_features, query_vec)

    # np.argsort will return indices of 9 smallest distances
    indices = np.argsort(pairwise_dist.flatten())[0:num_results]

    #data frame indices of the 9 smallest distace's
    df_indices = list(data.index[indices])

    results = []
    for i in range(0,len(indices)):
        result = {
            'asin': data['asin'].loc[df_indices[i]],
            'brand': data['brand'].loc[df_indices[i]],
            'title': data['title'].loc[df_indices[i]],
            'url': data['medium_image_url'].loc[df_indices[i]]
        }
        results.append(result)

    return results