added scheduler, logging and optimization updates

app.py

@@ -13,13 +13,21 @@ from torch_geometric.nn import SAGEConv
 from supabase import create_client
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
 from dotenv import load_dotenv
+import json
+from apscheduler.schedulers.background import BackgroundScheduler
+from apscheduler.triggers.cron import CronTrigger
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 load_dotenv()
 
 app = FastAPI()
 
-# Enable CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -37,11 +45,7 @@ if torch.cuda.is_available():
 
 # Global variables
 global G, features, user_ids, pyg_data, trained_model
-G = None
-features = None
-user_ids = None
-pyg_data = None
-trained_model = None
+G = features = user_ids = pyg_data = trained_model = None
 
 SUPABASE_URL = os.getenv('supabaseUrl')
 SUPABASE_KEY = os.getenv('supabaseAnonKey')
@@ -51,34 +55,60 @@ def get_supabase_client():
 
 def load_and_preprocess_data():
     supabase = get_supabase_client()
-    followers_response = supabase.table('followers').select('*').execute()
-    users_response = supabase.table('profiles').select('id').execute()
-
-    followers = pd.DataFrame(followers_response.data)
-    users = pd.DataFrame(users_response.data)
-
-    # 'id' is the followed user, 'following' is the follower
-    merged = followers.merge(users, left_on='id', right_on='id', how='left')
-    merged = merged.rename(columns={'following': 'follower_id', 'id': 'followed_id'})
-    merged = merged[['follower_id', 'followed_id']].dropna()
-    return merged[(merged['follower_id'] != '') & (merged['followed_id'] != '')]
-
-def create_graph_dataframe(merged_df):
-    # Edge from follower to followed
-    G = nx.from_pandas_edgelist(merged_df, source='follower_id', target='followed_id', create_using=nx.DiGraph())
+    logger.info("Loading data from Supabase")
+
+    def fetch_table(table, columns, chunk_size=1000):
+        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)
+            offset += chunk_size
+        return all_data
+
+    followers = fetch_table('followers', 'id, following')
+    users = fetch_table('profiles', 'id')
+
+    # Use native Python for merging instead of pandas
+    follower_dict = {f['id']: f['following'] for f in followers}
+    user_set = {u['id'] for u in users}
+    merged = [
+        {'follower_id': follower_dict[fid], 'followed_id': fid}
+        for fid in follower_dict if fid in user_set and follower_dict[fid] != '' and fid != ''
+    ]
+    logger.info(f"Loaded {len(merged)} follower relationships")
+    return merged
+
+def create_graph_dataframe(merged_data):
+    G = nx.DiGraph()
+    edges = [(d['follower_id'], d['followed_id']) for d in merged_data]
+    G.add_edges_from(edges)
     user_ids = sorted(G.nodes())
-    return G, torch.eye(len(user_ids)), user_ids
+    
+    # Use sparse identity matrix for features
+    features = torch.sparse_coo_tensor(
+        torch.arange(len(user_ids)).repeat(2, 1),
+        torch.ones(len(user_ids)),
+        (len(user_ids), len(user_ids))
+    )
+    logger.info(f"Created graph with {len(user_ids)} nodes")
+    return G, features, user_ids
 
 def prepare_training_data(G, user_ids):
     pos_edges = [(user_ids.index(u), user_ids.index(v)) for u, v in G.edges()]
-    pos_edge_index = torch.tensor(pos_edges).T
+    pos_edge_index = torch.tensor(pos_edges, dtype=torch.long).T
 
     num_nodes = len(user_ids)
     all_possible_edges = set(itertools.permutations(range(num_nodes), 2))
     existing_edges = set(zip(pos_edge_index[0].tolist(), pos_edge_index[1].tolist()))
-    negative_edges = random.sample(list(all_possible_edges - existing_edges), len(pos_edges))
+    neg_sample_size = len(pos_edges)
+    negative_edges = random.sample(list(all_possible_edges - existing_edges), neg_sample_size)
 
-    return pos_edge_index, torch.tensor(negative_edges).T
+    logger.info(f"Prepared {len(pos_edges)} positive and {len(negative_edges)} negative edges")
+    return pos_edge_index, torch.tensor(negative_edges, dtype=torch.long).T
 
 class GraphRecommender(nn.Module):
     def __init__(self, input_dim, hidden_dim=128, output_dim=64):
@@ -93,11 +123,18 @@ class GraphRecommender(nn.Module):
         x = self.conv2(x, edge_index)
         return x
 
-def train_model(model, data, pos_edges, neg_edges, epochs=200):
+def train_model(model, data, pos_edges, neg_edges, epochs=200, patience=20):
+    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)
     best_loss = float('inf')
     patience_counter = 0
 
+    logger.info("Starting model training")
     for epoch in range(epochs):
         model.train()
         optimizer.zero_grad()
@@ -121,42 +158,48 @@ def train_model(model, data, pos_edges, neg_edges, epochs=200):
             patience_counter = 0
         else:
             patience_counter += 1
-            if patience_counter >= 20:
+            if patience_counter >= patience:
+                logger.info(f"Early stopping at epoch {epoch}")
                 break
 
-    return model
+    logger.info("Model training completed")
+    return model.to('cpu')  # Move back to CPU for inference
 
 def get_recommendations(user_id, model, data, G, user_ids, top_k=10):
     if user_id not in user_ids:
         return []
 
     user_idx = user_ids.index(user_id)
-    # Successors are the users this user follows
     current_follows = set(G.successors(user_id))
-    
-    # Exclude self and already-followed users
-    candidates = [u for u in user_ids if u != user_id and u not in current_follows]
-    
+    candidate_indices = [i for i, u in enumerate(user_ids) if u != user_id and u not in current_follows]
+
     with torch.no_grad():
         embeddings = model(data.x, data.edge_index)
-        user_embed = embeddings[user_idx]
-        candidate_indices = [user_ids.index(u) for u in candidates]
+        user_embed = embeddings[user_idx].unsqueeze(0)
         candidate_embeds = embeddings[candidate_indices]
-        scores = torch.mm(user_embed.view(1, -1), candidate_embeds.T).squeeze()
+        scores = torch.matmul(user_embed, candidate_embeds.T).squeeze()
 
     top_indices = scores.argsort(descending=True)[:top_k]
-    return [candidates[i] for i in top_indices]
+    recommendations = [user_ids[candidate_indices[i]] for i in top_indices]
+    logger.info(f"Generated {len(recommendations)} recommendations for user {user_id}")
+    return recommendations
 
 def rebuild_model():
     global G, features, user_ids, pyg_data, trained_model
-    merged_df = load_and_preprocess_data()
-    G, features, user_ids = create_graph_dataframe(merged_df)
-    pyg_data = from_networkx(G)
-    pyg_data.x = features
-
-    pos_edge_index, neg_edge_index = prepare_training_data(G, user_ids)
-    model = GraphRecommender(input_dim=len(user_ids))
-    trained_model = train_model(model, pyg_data, pos_edge_index, neg_edge_index)
+    logger.info("Starting model rebuild at 3:30 AM")
+    try:
+        merged_data = load_and_preprocess_data()
+        G, features, user_ids = create_graph_dataframe(merged_data)
+        pyg_data = from_networkx(G)
+        pyg_data.x = features
+
+        pos_edge_index, neg_edge_index = prepare_training_data(G, user_ids)
+        model = GraphRecommender(input_dim=len(user_ids))
+        trained_model = train_model(model, pyg_data, pos_edge_index, neg_edge_index)
+        logger.info("Model rebuild completed successfully")
+    except Exception as e:
+        logger.error(f"Error during model rebuild: {str(e)}")
+        raise
 
 @app.post("/rebuild")
 async def rebuild_handler():
@@ -169,10 +212,42 @@ async def get_recommendations_handler(user_id: str = Query(...)):
         raise HTTPException(status_code=500, detail="Model not initialized, please rebuild first.")
 
     recommendations = get_recommendations(user_id, trained_model, pyg_data, G, user_ids)
-    return {"status": "success", "recommendations": recommendations}
+    
+    # Stream the response
+    def generate():
+        yield '{"status": "success", "recommendations": ['
+        for i, rec in enumerate(recommendations):
+            yield json.dumps(rec)
+            if i < len(recommendations) - 1:
+                yield ','
+        yield ']}'
+
+    return StreamingResponse(generate(), media_type="application/json")
 
 @app.get("/")
 async def health_check():
     return {"status": "success", "message": "Recommendation service operational"}
 
-rebuild_model()
+# Scheduler setup
+scheduler = BackgroundScheduler()
+scheduler.add_job(
+    rebuild_model,
+    trigger=CronTrigger(hour=3, minute=30),  # Run at 3:30 AM every day
+    id='daily_model_rebuild',
+    replace_existing=True
+)
+
+@app.on_event("startup")
+async def startup_event():
+    rebuild_model()  # Initial build on startup
+    scheduler.start()
+    logger.info("Scheduler started, model will rebuild daily at 3:30 AM")
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    scheduler.shutdown()
+    logger.info("Scheduler shut down")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)