word2vec get similar product

get_similar_product.py

def get_similar_product(model:Word2Vec, df:pd.DataFrame, product_id:int)->pd.DataFrame:
    """
    Parameters
    ----------
    model : instance of Word2Vec
    df : dataframe with preprocessed data
    product_id : int
        unique product id for which we need to find similar product ids
    Returns
    -------
    dataframe with similar products
    """
    try:
        sim_product = model.wv.most_similar(positive=[str(product_id)])
        return df.loc[df['product_id'].isin([int(word[0]) 
                                             for word in sim_product])][[
                                                     'category_code',
                                                     'brand',
                                                     'product_id']].drop_duplicates()
    except KeyError:
        return f"Cannot find the specified product with id {product_id}"