Update logic2.py

logic2.py

@@ -5,7 +5,7 @@ import itertools
 import requests
 import sys
 
-
+from pyvis.network import Network
 import torch
 import torch.nn.functional as F
 from torch.nn import Linear
@@ -100,4 +100,27 @@ def get_recommendation(model,data,user_id):
     top_ten_recs = [rec_movies for rec_movies in rec_movie_ids[0].tolist()[:10]] 
     top_ten_recs = [get_movie(movie_idx) for movie_idx in top_ten_recs]
     return {'user': user_id, 'rec_movies': top_ten_recs}
+
+def make_1_hop_graph(data,user_id):
+  a = data["user", "rates", "movie"].edge_index
+  b = data["user", "rates", "movie"].edge_label
+  idxs = (a[0] == user_id).nonzero(as_tuple=True)[0]
+  ratings = b[idxs]#.tolist()
+  movie_idxs = a[1][idxs]#.tolist()
+
+
+  n = len(ratings)
+  net = Network(cdn_resources="local",filter_menu=True)
+  for i in range(n):
+    #print(i)
+    Source = user_id
+    lab = get_title(movie_idxs[i])[0]
+    Target = movie_idxs[i] + 671 # Addition for sperating movie with user_id
+    weight = ratings[i].item()
+
+    net.add_node(Source, label=str(Source),color='#FF0000')
+    net.add_node(Target.item(), label=lab)
+    net.add_edge(Source, Target.item(), title=weight)
+  
+  net.show('index.html')