Update logic.py

logic.py

@@ -107,6 +107,50 @@ def create_ratings_graph(user_id, movie_id, ratings):
                 print("Inserting batch the last batch!")
                 edge_collection.import_bulk(batch)
 
+def create_pyg_edges(rating_docs):
+    src = []
+    dst = []
+    ratings = []
+    for doc in rating_docs:
+        _from = int(doc['_from'].split('/')[1])
+        _to   = int(doc['_to'].split('/')[1])
+         
+        src.append(_from)
+        dst.append(_to)
+        ratings.append(int(doc['_rating']))
+        
+    edge_index = torch.tensor([src, dst])
+    edge_attr = torch.tensor(ratings)
+
+    return edge_index, edge_attr 
+
+def SequenceEncoder(movie_docs , model_name=None):
+    movie_titles = [doc['movie_title'] for doc in movie_docs]
+    model = SentenceTransformer(model_name, device=device)
+    title_embeddings = model.encode(movie_titles, show_progress_bar=True,
+                              convert_to_tensor=True, device=device)
+    
+    return title_embeddings
+
+def GenresEncoder(movie_docs):
+    gen = []
+    #sep = '|'
+    for doc in movie_docs:
+        gen.append(doc['genres'])
+        #genre = doc['movie_genres']
+        #gen.append(genre.split(sep))
+    
+    # getting unique genres
+    unique_gen = set(list(itertools.chain(*gen)))
+    print("Number of unqiue genres we have:", unique_gen)
+    
+    mapping = {g: i for i, g in enumerate(unique_gen)}
+    x = torch.zeros(len(gen), len(mapping))
+    for i, m_gen in enumerate(gen):
+        for genre in m_gen:
+            x[i, mapping[genre]] = 1
+    return x.to(device)
+
 
 #-------------------------------------------------------------------------------------------
 
@@ -268,9 +312,43 @@ def load_data_to_ArangoDB(login):
 
 
     return movie_rec_db
+
+
+def make_pyg_graph(movie_rec_db):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    users = movie_rec_db.collection('Users')
+    movies = movie_rec_db.collection('Movie')
+    ratings_graph = movie_rec_db.collection('Ratings')
+
+    edge_index, edge_label = create_pyg_edges(movie_rec_db.aql.execute('FOR doc IN Ratings RETURN doc'))
     
+    title_emb = SequenceEncoder(movie_rec_db.aql.execute('FOR doc IN Movie RETURN doc'), model_name='all-MiniLM-L6-v2')
+    encoded_genres = GenresEncoder(movie_rec_db.aql.execute('FOR doc IN Movie RETURN doc'))
+    movie_x = torch.cat((title_emb, encoded_genres), dim=-1)
+
+    data = HeteroData()
+    data['user'].num_nodes = len(users)  # Users do not have any features.
+    data['movie'].x = movie_x
+    data['user', 'rates', 'movie'].edge_index = edge_index
+    data['user', 'rates', 'movie'].edge_label = edge_label
     
-
+    # Add user node features for message passing:
+    data['user'].x = torch.eye(data['user'].num_nodes, device=device)
+    del data['user'].num_nodes
+    data = ToUndirected()(data)
+    del data['movie', 'rev_rates', 'user'].edge_label  # Remove "reverse" label.
+
+    data = data.to(device)
+
+    train_data, val_data, test_data = T.RandomLinkSplit(
+        num_val=0.1,
+        num_test=0.1,
+        neg_sampling_ratio=0.0,
+        edge_types=[('user', 'rates', 'movie')],
+        rev_edge_types=[('movie', 'rev_rates', 'user')],
+    )(data)
+
+