Fixing small bugs

app.py

@@ -33,14 +33,9 @@ torch.manual_seed(SEED)
 if torch.cuda.is_available():
     torch.cuda.manual_seed_all(SEED)
 
+# Global Variables
 global G, features, user_nodes, post_nodes, node2idx, pyg_data, trained_model
-G = None
-features = None
-user_nodes = None
-post_nodes = None
-node2idx = None
-pyg_data = None
-trained_model = None
+G = features = user_nodes = post_nodes = node2idx = pyg_data = trained_model = None
 
 SUPABASE_URL = os.getenv('supabaseUrl')
 SUPABASE_KEY = os.getenv('supabaseAnonKey')
@@ -48,41 +43,27 @@ SUPABASE_KEY = os.getenv('supabaseAnonKey')
 def get_supabase_client():
     return create_client(SUPABASE_URL, SUPABASE_KEY)
 
-def load_and_preprocess_data_for_posts():
+def load_and_preprocess_data():
     supabase = get_supabase_client()
 
-    profiles_response = supabase.table('profiles').select('id').execute()
-    df_profiles = pd.DataFrame(profiles_response.data)
-
-    posts_response = supabase.table('posts').select('id, author').execute()
-    df_posts = pd.DataFrame(posts_response.data)
-
-    likes_response = supabase.table('likes').select('user_id, post_id').execute()
-    df_likes = pd.DataFrame(likes_response.data)
+    profiles = pd.DataFrame(supabase.table('profiles').select('id').execute().data)
+    posts = pd.DataFrame(supabase.table('posts').select('id, author').execute().data)
+    likes = pd.DataFrame(supabase.table('likes').select('user_id, post_id').execute().data)
 
     bipartite = nx.DiGraph()
 
-    user_set = set(df_posts['author'].dropna().tolist()) | set(df_likes['user_id'].dropna().tolist())
-    post_set = set(df_posts['id'].tolist())
+    user_set = set(posts['author']) | set(likes['user_id'])
+    post_set = set(posts['id'])
 
     for user in user_set:
-        if user:
-            bipartite.add_node(user, type='user')
-
+        bipartite.add_node(user, type='user')
     for post in post_set:
         bipartite.add_node(post, type='post')
 
-    for _, row in df_posts.iterrows():
-        user = row['author']
-        post = row['id']
-        if user and post:
-            bipartite.add_edge(user, post)
-
-    for _, row in df_likes.iterrows():
-        user = row['user_id']
-        post = row['post_id']
-        if user and post:
-            bipartite.add_edge(user, post)
+    for _, row in posts.iterrows():
+        bipartite.add_edge(row['author'], row['id'])
+    for _, row in likes.iterrows():
+        bipartite.add_edge(row['user_id'], row['post_id'])
 
     return bipartite
 
@@ -101,21 +82,18 @@ class GraphRecommender(nn.Module):
 
 def prepare_training_data(G, node2idx, user_nodes, post_nodes):
     pos_edges = [(node2idx[u], node2idx[v]) for u, v in G.edges() if G.nodes[u]['type'] == 'user' and G.nodes[v]['type'] == 'post']
-    pos_edge_index = torch.tensor(pos_edges).T
 
     all_possible = [(node2idx[u], node2idx[p]) for u in user_nodes for p in post_nodes]
     pos_set = set(pos_edges)
     neg_candidates = [pair for pair in all_possible if pair not in pos_set]
+
     neg_sample_size = min(len(pos_edges), len(neg_candidates))
     neg_edges = random.sample(neg_candidates, neg_sample_size)
-    neg_edge_index = torch.tensor(neg_edges).T
 
-    return pos_edge_index, neg_edge_index
+    return torch.tensor(pos_edges).T, torch.tensor(neg_edges).T
 
 def train_model(model, data, pos_edges, neg_edges, epochs=200):
     optimizer = optim.Adam(model.parameters(), lr=0.005, weight_decay=1e-4)
-    best_loss = float('inf')
-    patience_counter = 0
 
     for epoch in range(epochs):
         model.train()
@@ -128,29 +106,21 @@ def train_model(model, data, pos_edges, neg_edges, epochs=200):
 
         pos_loss = F.binary_cross_entropy_with_logits(pos_scores, torch.ones_like(pos_scores))
         neg_loss = F.binary_cross_entropy_with_logits(neg_scores, torch.zeros_like(neg_scores))
-        reg_loss = torch.norm(embeddings, p=2)
 
-        total_loss = pos_loss + neg_loss + 0.001 * reg_loss
+        total_loss = pos_loss + neg_loss
 
         total_loss.backward()
         optimizer.step()
 
-        if total_loss < best_loss:
-            best_loss = total_loss
-            patience_counter = 0
-        else:
-            patience_counter += 1
-            if patience_counter >= 20:
-                break
-
     return model
 
 def rebuild_model():
     global G, features, user_nodes, post_nodes, node2idx, pyg_data, trained_model
-    G = load_and_preprocess_data_for_posts()
 
-    user_nodes = sorted(n for n, attr in G.nodes(data=True) if attr.get('type') == 'user')
-    post_nodes = sorted(n for n, attr in G.nodes(data=True) if attr.get('type') == 'post')
+    G = load_and_preprocess_data()
+    user_nodes = sorted(n for n, attr in G.nodes(data=True) if attr['type'] == 'user')
+    post_nodes = sorted(n for n, attr in G.nodes(data=True) if attr['type'] == 'post')
+
     all_nodes = user_nodes + post_nodes
     node2idx = {node: i for i, node in enumerate(all_nodes)}
 
@@ -158,26 +128,44 @@ def rebuild_model():
     pyg_data = from_networkx(G)
     pyg_data.x = features
 
-    pos_edge_index, neg_edge_index = prepare_training_data(G, node2idx, user_nodes, post_nodes)
+    pos_edges, neg_edges = prepare_training_data(G, node2idx, user_nodes, post_nodes)
 
     input_dim = features.shape[1]
-    model = GraphRecommender(input_dim=input_dim)
-    trained_model = train_model(model, pyg_data, pos_edge_index, neg_edge_index)
+    model = GraphRecommender(input_dim)
+    trained_model = train_model(model, pyg_data, pos_edges, neg_edges)
+
+def get_recommendations(user_id, model, data, G, user_nodes, post_nodes, node2idx, top_k=10):
+    if user_id not in user_nodes:
+        return []
+
+    user_idx = node2idx[user_id]
+
+    user_interacted = {v for _, v in G.out_edges(user_id) if G.nodes[v]['type'] == 'post'}
+
+    with torch.no_grad():
+        embeddings = model(data.x, data.edge_index)
+    user_embed = embeddings[user_idx]
+
+    scores = [(post, torch.dot(user_embed, embeddings[node2idx[post]]).item()) for post in post_nodes if post not in user_interacted]
+    scores = sorted(scores, key=lambda x: x[1], reverse=True)
+
+    return [post for post, _ in scores[:top_k]]
 
 @app.post("/rebuild")
 async def rebuild_handler():
     rebuild_model()
-    return {"status": "success", "message": "Model and data rebuilt successfully"}
+    return {"status": "success", "message": "Model rebuilt successfully"}
 
 @app.get("/recommend/feed")
 async def get_recommendations_handler(user_id: str = Query(...)):
     if trained_model is None:
         raise HTTPException(status_code=500, detail="Model not initialized, please rebuild first.")
+
     recs = get_recommendations(user_id, trained_model, pyg_data, G, user_nodes, post_nodes, node2idx)
     return {"status": "success", "recommendations": recs}
 
 @app.get("/")
 async def health_check():
-    return {"status": "success", "message": "Recommendation service operational"}
+    return {"status": "success", "message": "Service operational"}
 
-rebuild_model()
+rebuild_model()