Update logic.py

logic.py

@@ -424,9 +424,33 @@ def make_pyg_graph(movie_rec_db):
         rev_edge_types=[('movie', 'rev_rates', 'user')],
     )(data)
 
-    return train_data, val_data, test_data
+    return data,train_data, val_data, test_data
 
 
+def load_model(train_data, val_data, test_data):
+    model = Model(hidden_channels=32)
+    with torch.no_grad():
+        model.encoder(train_data.x_dict, train_data.edge_index_dict)
+    model.load_state_dict(torch.load('model.pt'))
+    model.eval()
+    return model
+
+def get_recommendation(model,data,user_id):
+
+    movies = movie_rec_db.collection('Movie')
+    total_movies = len(movies)
+    
+    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}
+    
+
 
 def train(train_data, val_data, test_data):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')