Update logic2.py

logic2.py

@@ -74,4 +74,18 @@ def load_model(data):
     model.load_state_dict(torch.load('model.pt',map_location=torch.device('cpu')))
     model.eval()
     return model
+
+def get_recommendation(model,data,user_id):
+
+    total_movies = 9025
+    
+    user_row = torch.tensor([user_id] * total_movies)
+    all_movie_ids = torch.arange(total_movies)
+    edge_label_index = torch.stack([user_row, all_movie_ids], dim=0)
+    pred = model(data.x_dict, data.edge_index_dict,edge_label_index)
+    pred = pred.clamp(min=0, max=5)
+    # we will only select movies for the user where the predicting rating is =5
+    rec_movie_ids = (pred == 5).nonzero(as_tuple=True)
+    top_ten_recs = [rec_movies for rec_movies in rec_movie_ids[0].tolist()[:10]] 
+    return {'user': user_id, 'rec_movies': top_ten_recs}