small error handling and logging

app.py

@@ -1,7 +1,8 @@
 import os
 import random
+import asyncio
+import aiohttp
 import numpy as np
-import pandas as pd
 import networkx as nx
 import torch
 import torch.nn as nn
@@ -19,6 +20,7 @@ from apscheduler.schedulers.background import BackgroundScheduler
 from apscheduler.triggers.cron import CronTrigger
 import logging
 import pytz
+from collections import defaultdict
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -27,7 +29,6 @@ logger = logging.getLogger(__name__)
 load_dotenv()
 
 app = FastAPI()
-
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -44,8 +45,8 @@ 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 = features = user_nodes = post_nodes = node2idx = pyg_data = trained_model = None
+global G, features, user_nodes, post_nodes, node2idx, pyg_data, trained_model, post_embeddings
+G = features = user_nodes = post_nodes = node2idx = pyg_data = trained_model = post_embeddings = None
 
 SUPABASE_URL = os.getenv('supabaseUrl')
 SUPABASE_KEY = os.getenv('supabaseAnonKey')
@@ -53,39 +54,49 @@ SUPABASE_KEY = os.getenv('supabaseAnonKey')
 def get_supabase_client():
     return create_client(SUPABASE_URL, SUPABASE_KEY)
 
-def load_and_preprocess_data():
+async def fetch_chunk(session, table, columns, offset, chunk_size):
     supabase = get_supabase_client()
-    logger.info("Loading data from Supabase")
-
-    def fetch_table(table, columns, chunk_size=1000):
+    response = await asyncio.to_thread(
+        supabase.table(table).select(columns).range(offset, offset + chunk_size - 1).execute
+    )
+    return response.data
+
+async def fetch_table_async(table, columns, chunk_size=1000):
+    async with aiohttp.ClientSession() as session:
+        tasks = []
         offset = 0
-        all_data = []
-        while True:
-            response = supabase.table(table).select(columns).range(offset, offset + chunk_size - 1).execute()
-            data = response.data
-            if not data:
-                break
-            all_data.extend(data)
+        max_rows = 100000  # Adjust based on expected data size
+        while offset < max_rows:
+            tasks.append(fetch_chunk(session, table, columns, offset, chunk_size))
             offset += chunk_size
-        return all_data
+        chunks = await asyncio.gather(*tasks)
+        return [item for chunk in chunks if chunk for item in chunk]
 
-    profiles = fetch_table('profiles', 'id')
-    posts = fetch_table('posts', 'id, author')
-    likes = fetch_table('likes', 'user_id, post_id')
+def load_and_preprocess_data():
+    logger.info("Loading data from Supabase")
+    profiles = asyncio.run(fetch_table_async('profiles', 'id'))
+    posts = asyncio.run(fetch_table_async('posts', 'id, author'))
+    likes = asyncio.run(fetch_table_async('likes', 'user_id, post_id'))
 
-    bipartite = nx.DiGraph()
+    # Use adjacency list for graph
+    graph = defaultdict(set)
+    node_types = {}
 
     user_set = {p['author'] for p in posts} | {l['user_id'] for l in likes}
     post_set = {p['id'] for p in posts}
 
-    bipartite.add_nodes_from(user_set, type='user')
-    bipartite.add_nodes_from(post_set, type='post')
+    for user in user_set:
+        node_types[user] = 'user'
+    for post in post_set:
+        node_types[post] = 'post'
 
-    bipartite.add_edges_from((p['author'], p['id']) for p in posts)
-    bipartite.add_edges_from((l['user_id'], l['post_id']) for l in likes)
+    for p in posts:
+        graph[p['author']].add(p['id'])
+    for l in likes:
+        graph[l['user_id']].add(l['post_id'])
 
     logger.info(f"Loaded graph with {len(user_set)} users and {len(post_set)} posts")
-    return bipartite
+    return graph, node_types
 
 class GraphRecommender(nn.Module):
     def __init__(self, input_dim, hidden_dim=128, output_dim=64):
@@ -101,94 +112,107 @@ class GraphRecommender(nn.Module):
         return x
 
 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']
-
-    all_possible = [(node2idx[u], node2idx[p]) for u in user_nodes for p in post_nodes]
+    pos_edges = [(node2idx[u], node2idx[v]) for u in user_nodes for v in G[u]]
     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_edges = []
+    num_neg = len(pos_edges)
+    post_indices = [node2idx[p] for p in post_nodes]
+    for _ in range(num_neg):
+        u_idx = random.choice([node2idx[u] for u in user_nodes])
+        p_idx = random.choice(post_indices)
+        while (u_idx, p_idx) in pos_set:
+            p_idx = random.choice(post_indices)
+        neg_edges.append((u_idx, p_idx))
 
     logger.info(f"Prepared {len(pos_edges)} positive and {len(neg_edges)} negative edges")
-    return torch.tensor(pos_edges, dtype=torch.long).T, torch.tensor(neg_edges, dtype=torch.long).T
+    return torch.tensor(pos_edges, dtype=torch.long), torch.tensor(neg_edges, dtype=torch.long)
 
-def train_model(model, data, pos_edges, neg_edges, epochs=200):
+def train_model(model, data, pos_edges, neg_edges, epochs=50, batch_size=1024):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model = model.to(device)
     data = data.to(device)
-    pos_edges = pos_edges.to(device)
-    neg_edges = neg_edges.to(device)
-
     optimizer = optim.Adam(model.parameters(), lr=0.005, weight_decay=1e-4)
 
+    pos_dataset = torch.utils.data.TensorDataset(pos_edges.T)
+    neg_dataset = torch.utils.data.TensorDataset(neg_edges.T)
+    pos_loader = torch.utils.data.DataLoader(pos_dataset, batch_size=batch_size, shuffle=True)
+    neg_loader = torch.utils.data.DataLoader(neg_dataset, batch_size=batch_size, shuffle=True)
+
     logger.info("Starting model training")
     for epoch in range(epochs):
         model.train()
-        optimizer.zero_grad()
-
-        embeddings = model(data.x, data.edge_index)
+        total_loss = 0
+        for (pos_batch,), (neg_batch,) in zip(pos_loader, neg_loader):
+            pos_batch, neg_batch = pos_batch.to(device), neg_batch.to(device)
+            optimizer.zero_grad()
 
-        pos_scores = (embeddings[pos_edges[0]] * embeddings[pos_edges[1]]).sum(1)
-        neg_scores = (embeddings[neg_edges[0]] * embeddings[neg_edges[1]]).sum(1)
+            embeddings = model(data.x, data.edge_index)
+            pos_scores = (embeddings[pos_batch[:, 0]] * embeddings[pos_batch[:, 1]]).sum(1)
+            neg_scores = (embeddings[neg_batch[:, 0]] * embeddings[neg_batch[:, 1]]).sum(1)
 
-        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))
+            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))
+            loss = pos_loss + neg_loss
 
-        total_loss = pos_loss + neg_loss
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
 
-        total_loss.backward()
-        optimizer.step()
+        logger.info(f"Epoch {epoch+1}, Loss: {total_loss}")
 
+    model.eval()
+    with torch.no_grad():
+        embeddings = model(data.x, data.edge_index).to('cpu')
     logger.info("Model training completed")
-    return model.to('cpu')
+    return model.to('cpu'), embeddings
 
 def rebuild_model():
-    global G, features, user_nodes, post_nodes, node2idx, pyg_data, trained_model
+    global G, features, user_nodes, post_nodes, node2idx, pyg_data, trained_model, post_embeddings
     logger.info("Starting model rebuild at 3:30 AM Pacific Time")
     try:
-        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')
+        G, node_types = load_and_preprocess_data()
+        user_nodes = sorted(n for n in G if node_types[n] == 'user')
+        post_nodes = sorted(n for n in node_types if node_types[n] == 'post')
 
         all_nodes = user_nodes + post_nodes
         node2idx = {node: i for i, node in enumerate(all_nodes)}
 
-        # Use dense features instead of sparse (SAGEConv requires dense input)
+        # Convert to PyTorch Geometric format
+        nx_G = nx.DiGraph()  # Temporary for conversion
+        nx_G.add_nodes_from(user_nodes + post_nodes)
+        for u in G:
+            for v in G[u]:
+                nx_G.add_edge(u, v)
         features = torch.eye(len(all_nodes))
-        pyg_data = from_networkx(G)
+        pyg_data = from_networkx(nx_G)
         pyg_data.x = features
 
         pos_edges, neg_edges = prepare_training_data(G, node2idx, user_nodes, post_nodes)
-
         input_dim = features.shape[1]
         model = GraphRecommender(input_dim)
-        trained_model = train_model(model, pyg_data, pos_edges, neg_edges)
+        trained_model, post_embeddings = train_model(model, pyg_data, pos_edges, neg_edges)
         logger.info("Model rebuild completed successfully")
     except Exception as e:
         logger.error(f"Error during model rebuild: {str(e)}")
         raise
 
-def get_recommendations(user_id, model, data, G, user_nodes, post_nodes, node2idx):
+def get_recommendations(user_id, model, data, G, user_nodes, post_nodes, node2idx, precomputed_embeds):
     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'}
+    user_interacted = set(G[user_id])
     post_indices = [node2idx[p] for p in post_nodes if p not in user_interacted]
 
     with torch.no_grad():
-        embeddings = model(data.x, data.edge_index)
-        user_embed = embeddings[user_idx].unsqueeze(0)
-        post_embeds = embeddings[post_indices]
-
+        user_embed = model(data.x, data.edge_index)[user_idx].unsqueeze(0)
+        post_embeds = precomputed_embeds[post_indices]
         scores = torch.matmul(user_embed, post_embeds.T).squeeze(0)
 
-    # Create an inverse mapping from index to post ID
     idx2node = {idx: node for node, idx in node2idx.items()}
-    # Map post_indices back to original post IDs
     post_scores = [(idx2node[i], score.item()) for i, score in zip(post_indices, scores)]
-    post_scores = sorted(post_scores, key=lambda x: x[1], reverse=True)
+    post_scores = sorted(post_scores, key=lambda x: x[1], reverse=True)[:10]  # Top-10
 
     logger.info(f"Generated {len(post_scores)} recommendations for user {user_id}")
     return [{"post_id": post, "score": score} for post, score in post_scores]
@@ -197,16 +221,14 @@ def fetch_full_post_records(post_ids, batch_size=1000):
     supabase = get_supabase_client()
     if not post_ids:
         return []
-
     records = []
     for i in range(0, len(post_ids), batch_size):
         batch_ids = post_ids[i:i + batch_size]
-        response = supabase.table('posts').select('*').in_('id', batch_ids).execute()
+        response = supabase.table('posts').select('id, title, author').in_('id', batch_ids).execute()
         batch_records = response.data
         for record in batch_records:
             record['type'] = 'post'
         records.extend(batch_records)
-
     return records
 
 @app.post("/rebuild")
@@ -219,7 +241,7 @@ 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.")
 
-    recommended_posts = get_recommendations(user_id, trained_model, pyg_data, G, user_nodes, post_nodes, node2idx)
+    recommended_posts = get_recommendations(user_id, trained_model, pyg_data, G, user_nodes, post_nodes, node2idx, post_embeddings)
     if not recommended_posts:
         return {"status": "success", "recommendations": []}
 
@@ -235,7 +257,10 @@ async def get_recommendations_handler(user_id: str = Query(...)):
             post_record["score"] = post["score"]
             ordered_recommendations.append(post_record)
 
-    # Stream the response
+    if ordered_recommendations:
+        insert_position = random.randint(0, min(9, len(ordered_recommendations) - 1))
+        ordered_recommendations.insert(insert_position, {"type": "suggestedaccounts"})
+
     def generate():
         yield '{"status": "success", "recommendations": ['
         for i, rec in enumerate(ordered_recommendations):

requirements.txt

@@ -8,4 +8,5 @@ fastapi
 python-dotenv
 uvicorn
 apscheduler
-pytz
+pytz
+aiohttp