File size: 3,471 Bytes
97de7f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
import pandas as pd
from arango import ArangoClient
from tqdm import tqdm
import numpy as np
import itertools
import requests
import sys
import oasis
from arango import ArangoClient
import torch
import torch.nn.functional as F
from torch.nn import Linear
from arango import ArangoClient
import torch_geometric.transforms as T
from torch_geometric.nn import SAGEConv, to_hetero
from torch_geometric.transforms import RandomLinkSplit, ToUndirected
from sentence_transformers import SentenceTransformer
from torch_geometric.data import HeteroData
import yaml
#-------------------------------------------------------------------------------------------
# Functions
# performs user and movie mappings
def node_mappings(path, index_col):
df = pd.read_csv(path, index_col=index_col)
mapping = {index: i for i, index in enumerate(df.index.unique())}
return mapping
def convert_int(x):
try:
return int(x)
except:
return np.nan
def remove_movies(m_id):
'''
# Remove ids which dont have meta data information
'''
no_metadata = []
for idx in range(len(m_id)):
tmdb_id = id_map.loc[id_map['movieId'] == m_id[idx]]
if tmdb_id.size == 0:
no_metadata.append(m_id[idx])
#print('No Meta data information at:', m_id[idx])
return no_metadata
#-------------------------------------------------------------------------------------------
def make_graph():
metadata_path = './sampled_movie_dataset/movies_metadata.csv'
df = pd.read_csv(metadata_path)
df = df.drop([19730, 29503, 35587])
df['id'] = df['id'].astype('int')
links_small = pd.read_csv('./sampled_movie_dataset/links_small.csv')
links_small = links_small[links_small['tmdbId'].notnull()]['tmdbId'].astype('int') # selecting tmdbId coloumn from links_small file
sampled_md = df[df['id'].isin(links_small)]
sampled_md['tagline'] = sampled_md['tagline'].fillna('')
sampled_md['description'] = sampled_md['overview'] + sampled_md['tagline']
sampled_md['description'] = sampled_md['description'].fillna('')
sampled_md = sampled_md.reset_index()
indices = pd.Series(sampled_md.index, index=sampled_md['title'])
ind_gen = pd.Series(sampled_md.index, index=sampled_md['genres'])
ratings_path = './sampled_movie_dataset/ratings_small.csv'
ratings_df = pd.read_csv(ratings_path)
m_id = ratings_df['movieId'].tolist()
m_id = list(dict.fromkeys(m_id))
user_mapping = node_mappings(ratings_path, index_col='userId')
movie_mapping = node_mappings(ratings_path, index_col='movieId')
id_map = pd.read_csv('./sampled_movie_dataset/links_small.csv')[['movieId', 'tmdbId']]
id_map['tmdbId'] = id_map['tmdbId'].apply(convert_int)
id_map.columns = ['movieId', 'id']
id_map = id_map.merge(sampled_md[['title', 'id']], on='id').set_index('title') # tmbdid is same (of links_small) as of id in sampled_md
indices_map = id_map.set_index('id')
no_metadata = remove_movies(m_id)
## remove ids which dont have meta data information
for element in no_metadata:
if element in m_id:
print("ids with no metadata information:",element)
m_id.remove(element)
# create new movie_mapping dict with only m_ids having metadata information
movie_mappings = {}
for idx, m in enumerate(m_id):
movie_mappings[m] = idx
return movie_mappings
def load_data_to_ArangoDB():
|