import os import os.path as osp import gzip import pickle import json import torch import pandas as pd import numpy as np from collections import Counter from tqdm import tqdm from huggingface_hub import hf_hub_download import zipfile from ogb.utils.url import download_url from src.benchmarks.semistruct.knowledge_base import SemiStructureKB from src.tools.process_text import clean_data, compact_text from src.tools.node import df_row_to_dict, Node, register_node from src.tools.io import save_files, load_files PROCESSED_DATASET = { "repo": "snap-stanford/stark", "file": "skb/amazon/processed.zip", } class AmazonSemiStruct(SemiStructureKB): REVIEW_CATEGORIES = set(['Amazon_Fashion','All_Beauty','Appliances', 'Arts_Crafts_and_Sewing','Automotive','Books', 'CDs_and_Vinyl','Cell_Phones_and_Accessories', 'Clothing_Shoes_and_Jewelry','Digital_Music', 'Electronics','Gift_Cards','Grocery_and_Gourmet_Food', 'Home_and_Kitchen','Industrial_and_Scientific', 'Kindle_Store', 'Luxury_Beauty','Magazine_Subscriptions', 'Movies_and_TV', 'Musical_Instruments', 'Office_Products','Patio_Lawn_and_Garden', 'Pet_Supplies','Prime_Pantry','Software','Sports_and_Outdoors', 'Tools_and_Home_Improvement','Toys_and_Games','Video_Games']) # single answers QA_CATEGORIES = set(['Appliances','Arts_Crafts_and_Sewing', 'Automotive', 'Baby','Beauty','Cell_Phones_and_Accessories', 'Clothing_Shoes_and_Jewelry','Electronics', 'Grocery_and_Gourmet_Food','Health_and_Personal_Care', 'Home_and_Kitchen','Musical_Instruments','Office_Products', 'Patio_Lawn_and_Garden','Pet_Supplies','Sports_and_Outdoors', 'Tools_and_Home_Improvement','Toys_and_Games','Video_Games']) COMMON = set(['Appliances', 'Arts_Crafts_and_Sewing', 'Automotive', 'Cell_Phones_and_Accessories', 'Clothing_Shoes_and_Jewelry', 'Electronics', 'Grocery_and_Gourmet_Food', 'Home_and_Kitchen', 'Musical_Instruments', 'Office_Products', 'Patio_Lawn_and_Garden', 'Pet_Supplies', 'Sports_and_Outdoors', 'Tools_and_Home_Improvement', 'Toys_and_Games', 'Video_Games']) sub_category = 'data/amazon/category_list.json' SUB_CATEGORIES = set(json.load(open(sub_category, 'r'))) link_columns = ['also_buy', 'also_view'] review_columns = ['reviewerID', 'summary', 'style', 'reviewText', 'vote', 'overall', 'verified', 'reviewTime'] qa_columns = ['questionType', 'answerType', 'question', 'answer', 'answerTime'] meta_columns = ['asin', 'title', 'global_category', 'category', 'price', 'brand', 'feature', 'rank', 'details', 'description'] candidate_types = ['product'] node_attr_dict = {'product': ['title', 'dimensions', 'weight', 'description', 'features', 'reviews', 'Q&A'], 'brand': ['brand_name'], 'category': ['category_name'], 'color': ['color_name']} def __init__(self, root, categories: list, meta_link_types=['brand', 'category', 'color'], max_entries=25, download_processed=True, **kwargs): ''' Args: root (str): root directory to store the data categories (list): product categories meta_link_types (list): a list which may contain entries in node info that used to consruct meta links, e.g. ['category', 'brand'] will construct entity nodes of catrgory and brand which link to corresponding nodes max_entries (int): maximum number of review & qa entries to show in the description indirected (bool): make the graph indirected ''' self.root = root self.max_entries = max_entries self.raw_data_dir = osp.join(root, 'raw') self.processed_data_dir = osp.join(root, 'processed') os.makedirs(self.raw_data_dir, exist_ok=True) os.makedirs(self.processed_data_dir, exist_ok=True) # construct the graph based on link info in the raw data cache_path = None if meta_link_types is None else \ osp.join(self.processed_data_dir, 'cache', '-'.join(meta_link_types)) if not osp.exists(osp.join(cache_path, 'node_info.pkl')) and download_processed: print('Downloading processed data...') processed_path = hf_hub_download( PROCESSED_DATASET["repo"], PROCESSED_DATASET["file"], repo_type="dataset" ) with zipfile.ZipFile(processed_path, 'r') as zip_ref: zip_ref.extractall(self.root) os.remove(processed_path) print('Downloaded processed data!') if not (cache_path is None) and osp.exists(cache_path): print(f'Load cached graph with meta link types {meta_link_types}') processed_data = load_files(cache_path) else: print(f'Start processing raw data...') print(f'{meta_link_types=}') processed_data = self._process_raw(categories) if meta_link_types: # customize the graph by adding meta links processed_data = self.post_process(processed_data, meta_link_types=meta_link_types, cache_path=cache_path) super(AmazonSemiStruct, self).__init__(**processed_data, **kwargs) def __getitem__(self, idx): idx = int(idx) node_info = self.node_info[idx] node = Node() register_node(node, node_info) return node def get_chunk_info(self, idx, attribute): if not hasattr(self[idx], attribute): return '' node_attr = getattr(self[idx], attribute) if 'feature' in attribute: features = [] if len(node_attr): for feature_idx, feature in enumerate(node_attr): if feature == '': continue if 'asin' in feature.lower(): continue features.append(feature) chunk = ' '.join(features) elif 'review' in attribute: chunk = '' if len(node_attr): scores = [0 if pd.isnull(review['vote']) else int(review['vote'].replace(",","")) for review in node_attr] ranks = np.argsort(-np.array(scores)) for idx, review_idx in enumerate(ranks): review = node_attr[review_idx] chunk += 'The review \"' + str(review['summary']) + '\"' chunk += 'states that \"' + str(review['reviewText']) + '\". ' if idx > self.max_entries: break elif 'qa' in attribute: chunk = '' if len(node_attr): for idx, question in enumerate(node_attr): chunk += 'The question is \"' + str(question['question']) + '\", ' chunk += 'and the answer is \"' + str(question['answer']) + '\". ' if idx > self.max_entries: break elif 'description' in attribute and len(node_attr): chunk = " ".join(node_attr) else: chunk = node_attr return chunk def get_doc_info(self, idx, add_rel=True, compact=False): if self.node_type_dict[int(self.node_types[idx])] == 'brand': return f'brand name: {self[idx].brand_name}' if self.node_type_dict[int(self.node_types[idx])] == 'category': return f'category name: {self[idx].category_name}' if self.node_type_dict[int(self.node_types[idx])] == 'color': return f'color name: {self[idx].color_name}' node = self[idx] doc = f'- product: {node.title}\n' if hasattr(node, 'brand'): doc += f'- brand: {node.brand}\n' try: dimensions, weight = node.details.dictionary.product_dimensions.split(' ; ') doc += (f'- dimensions: {dimensions}\n' f'- weight: {weight}\n') except: pass if len(node.description): description = " ".join(node.description).strip(" ") if len(description) > 0: doc += f'- description: {description}\n' feature_text = f'- features: \n' if len(node.feature): for feature_idx, feature in enumerate(node.feature): if feature == '': continue if 'asin' in feature.lower(): continue feature_text += (f'#{feature_idx + 1}: {feature}\n') else: feature_text = '' if len(node.review): review_text = f'- reviews: \n' scores = [0 if pd.isnull(review['vote']) else int(review['vote'].replace(",","")) for review in node.review] ranks = np.argsort(-np.array(scores)) for i, review_idx in enumerate(ranks): review = node.review[review_idx] review_text += (f'#{review_idx + 1}:\n' f'summary: {review["summary"]}\n' f'text: "{review["reviewText"]}"\n') if i > self.max_entries: break else: review_text = '' if len(node.qa): qa_text = f'- Q&A: \n' for qa_idx, qa in enumerate(node.qa): qa_text += (f'#{qa_idx + 1}:\n' f'question: "{qa["question"]}"\n' f'answer: "{qa["answer"]}"\n') if qa_idx > self.max_entries: break else: qa_text = '' doc += feature_text + review_text + qa_text if add_rel: doc += self.get_rel_info(idx) if compact: doc = compact_text(doc) return doc def get_rel_info(self, idx, rel_types=None, n_rel=-1): doc = '' rel_types = self.rel_type_lst() if rel_types is None else rel_types n_also_buy = self.get_neighbor_nodes(idx, 'also_buy') n_also_view = self.get_neighbor_nodes(idx, 'also_view') n_has_brand = self.get_neighbor_nodes(idx, 'has_brand') str_also_buy = [f"#{idx + 1}: " + self[i].title + '\n' for idx, i in enumerate(n_also_buy)] str_also_view = [f"#{idx + 1}: " + self[i].title + '\n' for idx, i in enumerate(n_also_view)] if len(str_also_buy) == 0: str_also_buy = '' if len(str_also_view) == 0: str_also_view = '' str_has_brand = '' if len(n_has_brand): str_has_brand = f' brand: {self[n_has_brand[0]].brand_name}\n' str_also_buy = ''.join(str_also_buy) str_also_view = ''.join(str_also_view) if len(str_also_buy): doc += f' products also purchased: \n{str_also_buy}' if len(str_also_view): doc += f' products also viewed: \n{str_also_view}' if len(n_has_brand): doc += str_has_brand if len(doc): doc = '- relations:\n' + doc return doc def _process_raw(self, categories): if 'all' in categories: review_categories = self.REVIEW_CATEGORIES qa_categories = self.QA_CATEGORIES else: qa_categories = review_categories = categories assert len(set(categories) - self.COMMON) == 0, f'invalid categories exist' if osp.exists(osp.join(self.processed_data_dir, 'node_info.pkl')): print(f'Load processed data from {self.processed_data_dir}') loaded_files = load_files(self.processed_data_dir) loaded_files.update( {'node_types': torch.zeros(len(loaded_files['node_info'])), 'node_type_dict': {0: 'product'}}) return loaded_files print(f'Check data downloading...') for category in review_categories: review_header = 'https://datarepo.eng.ucsd.edu/mcauley_group/data/amazon_v2' if not os.path.exists(osp.join(self.raw_data_dir, f'{category}.json.gz')): print(f'Downloading {category} data...') download_url(f'{review_header}/categoryFiles/{category}.json.gz', self.raw_data_dir) download_url(f'{review_header}/metaFiles2/meta_{category}.json.gz', self.raw_data_dir) for category in qa_categories: qa_header = 'https://datarepo.eng.ucsd.edu/mcauley_group/data/amazon/qa' if not os.path.exists(osp.join(self.raw_data_dir, f'qa_{category}.json.gz')): print(f'Downloading {category} QA data...') download_url(f'{qa_header}/qa_{category}.json.gz', self.raw_data_dir) if not osp.exists(osp.join(self.processed_data_dir, 'node_info.pkl')): ckt_path = 'data/amazon/intermediate' print('Loading data... It might take a while') # read amazon QA data df_qa_path = os.path.join(ckt_path, 'df_qa.pkl') if os.path.exists(df_qa_path): df_qa = pd.read_pickle(df_qa_path) else: df_qa = pd.concat([read_qa(osp.join(self.raw_data_dir, f'qa_{category}.json.gz')) for category in qa_categories])[['asin'] + self.qa_columns] df_qa.to_pickle(df_qa_path) print('df_qa loaded') # read amazon review data df_review_path = os.path.join(ckt_path, 'df_review.pkl') if os.path.exists(df_review_path): df_review = pd.read_pickle(df_review_path) else: df_review = pd.concat([read_review(osp.join(self.raw_data_dir, f'{category}.json.gz')) for category in review_categories])[['asin'] + self.review_columns] df_review.to_pickle(df_review_path) print('df_review loaded') # read amazon meta data from amazon review & amazon kdd df_ucsd_meta_path = os.path.join(ckt_path, 'df_ucsd_meta.pkl') if os.path.exists(df_ucsd_meta_path): df_ucsd_meta = pd.read_pickle(df_ucsd_meta_path) else: meta_df_lst = [] for category in review_categories: cat_review = read_review(osp.join(self.raw_data_dir, f'meta_{category}.json.gz')) cat_review.insert(0, 'global_category', category.replace('_', ' ')) meta_df_lst.append(cat_review) df_ucsd_meta = pd.concat(meta_df_lst) df_ucsd_meta.to_pickle(df_ucsd_meta_path) print('df_ucsd_meta loaded') print('Preprocessing data...') df_ucsd_meta = df_ucsd_meta.drop_duplicates(subset='asin', keep='first') df_meta = df_ucsd_meta[self.meta_columns + self.link_columns] # Merge dataframes df_review_meta = df_review.merge(df_meta, left_on='asin', right_on='asin') unique_asin = np.unique(np.array(df_review_meta['asin'])) # Filer items with both meta and review data df_qa_reduced = df_qa[df_qa['asin'].isin(unique_asin)] df_review_reduced = df_review[df_review['asin'].isin(unique_asin)] df_meta_reduced = df_meta[df_meta['asin'].isin(unique_asin)].reset_index() def get_map(df): asin2id, id2asin = {}, {} for idx in range(len(df)): asin2id[df['asin'][idx]] = idx id2asin[idx] = df['asin'][idx] return asin2id, id2asin print('Construct node info and graph...') # get mapping from asin to node id and its reversed mapping self.asin2id, self.id2asin = get_map(df_meta_reduced) node_info = self.construct_raw_node_info(df_meta_reduced, df_review_reduced, df_qa_reduced) edge_index, edge_types = self.create_raw_product_graph(df_meta_reduced, columns=self.link_columns) edge_type_dict = {0: 'also_buy', 1: 'also_view'} processed_data = { 'node_info': node_info, 'edge_index': edge_index, 'edge_types': edge_types, 'edge_type_dict': edge_type_dict} print(f'Saving to {self.processed_data_dir}...') save_files(save_path=self.processed_data_dir, **processed_data) processed_data.update({'node_types': torch.zeros(len(processed_data['node_info'])), 'node_type_dict': {0: 'product'}}) return processed_data def post_process(self, raw_info, meta_link_types, cache_path=None): print(f'Adding meta link types {meta_link_types}') node_info = raw_info['node_info'] edge_type_dict = raw_info['edge_type_dict'] node_type_dict = raw_info['node_type_dict'] node_types = raw_info['node_types'].tolist() edge_index = raw_info['edge_index'].tolist() edge_types = raw_info['edge_types'].tolist() n_e_types, n_n_types = len(edge_type_dict), len(node_type_dict) for i, link_type in enumerate(meta_link_types): if link_type == 'brand': values = np.array([node_info_i[link_type] for node_info_i in node_info.values() if link_type in node_info_i.keys()]) indices = np.array([idx for idx, node_info_i in enumerate(node_info.values()) if link_type in node_info_i.keys()]) elif link_type in ['category', 'color']: value_list = [] indice_list = [] for idx, node_info_i in enumerate(node_info.values()): if link_type in node_info_i.keys(): value_list.extend(node_info_i[link_type]) indice_list.extend([idx for _ in range(len(node_info_i[link_type]))]) values = np.array(value_list) indices = np.array(indice_list) else: raise Exception(f'Invalid meta link type {link_type}') cur_n_nodes = len(node_info) node_type_dict[n_n_types + i] = link_type edge_type_dict[n_e_types + i] = "has_" + link_type unique = np.unique(values) for j, unique_j in tqdm(enumerate(unique)): node_info[cur_n_nodes + j] = {link_type + '_name': unique_j} ids = indices[np.array(values == unique_j)] edge_index[0].extend(list(ids)) edge_index[1].extend([cur_n_nodes + j for _ in range(len(ids))]) edge_types.extend([i + n_e_types for _ in range(len(ids))]) node_types.extend([n_n_types + i for _ in range(len(unique))]) print(f'finished adding {link_type}') edge_index = torch.LongTensor(edge_index) edge_types = torch.LongTensor(edge_types) node_types = torch.LongTensor(node_types) files = {'node_info': node_info, 'edge_index': edge_index, 'edge_types': edge_types, 'edge_type_dict': edge_type_dict, 'node_type_dict': node_type_dict, 'node_types': node_types } if cache_path is not None: save_files(cache_path, **files) return files def _process_brand(self, brand): brand = brand.strip(" \".*+,-_!@#$%^&*();\/|<>\'\t\n\r\\") if len(brand) > 3 and brand[:3] == 'by ': brand = brand[3:] if len(brand) > 4 and brand[-4:] == '.com': brand = brand[:-4] if len(brand) > 4 and brand[:4] == 'www.': brand = brand[4:] if len(brand) > 100: brand = brand.split(' ')[0] return brand def construct_raw_node_info(self, df_meta, df_review, df_qa): node_info = {} for idx, asin in self.id2asin.items(): node_info[idx] = {} node_info[idx]['review'] = [] node_info[idx]['qa'] = [] ###################### Assign color ######################## def assign_colors(df_review, lower_limit=20): # asign to color df_review = df_review[['asin', 'style']] df_review = df_review.dropna(subset=['style']) raw_color_dict = {} for idx, row in tqdm(df_review.iterrows()): asin, style = row['asin'], row['style'] for key in style.keys(): if 'color' in key.lower(): try: raw_color_dict[asin] except: raw_color_dict[asin] = [] raw_color_dict[asin].append( style[key].strip().lower() if isinstance(style[key], str) else style[key][0].strip()) all_color_values = [] for asin in raw_color_dict.keys(): raw_color_dict[asin] = list(set(raw_color_dict[asin])) all_color_values.extend(raw_color_dict[asin]) print('number of all colors', len(all_color_values)) color_counter = Counter(all_color_values) print('number of unique colors', len(color_counter)) color_counter = {k: v for k, v in sorted(color_counter.items(), key=lambda item: item[1], reverse=True)} selected_colors = [] for color, number in color_counter.items(): if number > lower_limit and len(color) > 2 and len(color.split(' ')) < 5 and color.isnumeric() is False: selected_colors.append(color) print('number of selected colors', len(selected_colors)) filtered_color_dict = {} total_color_connections = 0 for asin in raw_color_dict.keys(): filtered_color_dict[asin] = [] for value in raw_color_dict[asin]: if value in selected_colors: filtered_color_dict[asin].append(value) total_color_connections += len(filtered_color_dict[asin]) print('number of linked products', len(filtered_color_dict)) print('number of total connections', total_color_connections) return filtered_color_dict filtered_color_dict_path = os.path.join('data/amazon/intermediate', 'filtered_color_dict.pkl') if os.path.exists(filtered_color_dict_path): with open(filtered_color_dict_path, 'rb') as f: filtered_color_dict = pickle.load(f) else: filtered_color_dict = assign_colors(df_review) with open(filtered_color_dict_path, 'wb') as f: pickle.dump(filtered_color_dict, f) for i in tqdm(range(len(df_meta))): df_meta_i = df_meta.iloc[i] asin = df_meta_i['asin'] idx = self.asin2id[asin] try: color = filtered_color_dict[asin] if len(color): node_info[idx]['color'] = color except: pass print('loaded color') #################################################################### for i in tqdm(range(len(df_meta))): df_meta_i = df_meta.iloc[i] asin = df_meta_i['asin'] idx = self.asin2id[asin] for column in self.meta_columns: if column == 'brand': brand = self._process_brand(clean_data(df_meta_i[column])) if len(brand) > 1: node_info[idx]['brand'] = brand elif column == 'category': category_list = [] for category in df_meta_i[column]: category = category.lower() if category in self.SUB_CATEGORIES: category_list.append(category) if len(category_list) > 0: node_info[idx]['category'] = category_list else: node_info[idx][column] = clean_data(df_meta_i[column]) review_columns = self.review_columns review_columns.remove('style') for name, df in zip(['review', 'qa'], [df_review, df_qa]): for i in tqdm(range(len(df))): df_i = df.iloc[i] asin = df_i['asin'] idx = self.asin2id[asin] node_info[idx][name].append( df_row_to_dict(df_i, colunm_names=self.review_columns \ if name == 'review' else self.qa_columns)) import pdb; pdb.set_trace() return node_info def create_raw_product_graph(self, df, columns): edge_types = [] edge_index = [[], []] for idx in range(len(df)): out_node = self.asin2id[df['asin'].iloc[idx]] for edge_type_id, edge_type in enumerate(columns): in_nodes = [] if not isinstance(df[edge_type].iloc[idx], list): continue for i in df[edge_type].iloc[idx]: try: in_nodes.append(self.asin2id[i]) except KeyError: continue edge_types.extend([edge_type_id for _ in range(len(in_nodes))]) edge_index[0].extend([out_node for _ in range(len(in_nodes))]) edge_index[1].extend(in_nodes) return torch.LongTensor(edge_index), torch.LongTensor(edge_types) def has_brand(self, idx, brand): try: b = self[idx].brand if len(b) > 4 and b[-4:] == '.com': b = b[:-4] if len(brand) > 4 and brand[-4:] == '.com': brand = brand[:-4] return b.lower().strip("\"") == brand.lower().strip("\"") except: return False def has_also_buy(self, idx, also_buy_item): try: also_buy_lst = self.get_neighbor_nodes(idx, 'also_buy') return also_buy_item in also_buy_lst except: return False def has_also_view(self, idx, also_view_item): try: also_buy_lst = self.get_neighbor_nodes(idx, 'also_view') return also_view_item in also_buy_lst except: return False # read review files def read_review(path): def parse(path): g = gzip.open(path, 'rb') for l in g: yield json.loads(l) def getDF(path): i = 0 df = {} for d in parse(path): df[i] = d i += 1 return pd.DataFrame.from_dict(df, orient='index') return getDF(path) # read qa files def read_qa(path): def parse(path): g = gzip.open(path, 'rb') for l in g: yield eval(l) def getDF(path): i = 0 df = {} for d in parse(path): df[i] = d i += 1 return pd.DataFrame.from_dict(df, orient='index') return getDF(path)