Updated experiments code

app/database.py

@@ -1,3 +1,4 @@
+import functools
 import io
 import sqlite3
 
@@ -26,8 +27,8 @@ class ItemDatabase:
             rows = c.fetchall()[:n_items]
             return [row[0] for row in rows]
 
+    @functools.lru_cache(maxsize=2**14)
     def get_item(self, item_id):
-
         with self._connect() as conn:
             c = conn.cursor()
             c.row_factory = sqlite3.Row

exp/deepwalk.py

@@ -1,80 +0,0 @@
-import argparse
-import os
-
-import numpy as np
-import pandas as pd
-import dgl
-import torch
-import wandb
-from tqdm.auto import tqdm
-
-from utils import prepare_graphs, extract_item_embeddings, normalize_embeddings
-
-
-def prepare_deepwalk_embeddings(
-        items_path,
-        ratings_path, 
-        embeddings_savepath,
-        emb_dim, 
-        window_size, 
-        batch_size, 
-        lr, 
-        num_epochs, 
-        device, 
-        wandb_name, 
-        use_wandb
-):
-    ### Prepare graph
-    bipartite_graph, graph = prepare_graphs(items_path, ratings_path)
-    bipartite_graph = bipartite_graph.to(device)
-    graph = graph.to(device)
-
-    ### Run deepwalk
-    if use_wandb:
-        wandb.init(project="graph-recs-deepwalk", name=wandb_name)
-
-    model = dgl.nn.DeepWalk(graph, emb_dim=emb_dim, window_size=window_size)
-    model = model.to(device)
-    dataloader = torch.utils.data.DataLoader(
-        torch.arange(graph.num_nodes()), 
-        batch_size=batch_size, 
-        shuffle=True, 
-        collate_fn=model.sample)
-    
-    optimizer = torch.optim.SparseAdam(model.parameters(), lr=lr)
-    for epoch in range(num_epochs):
-        for batch_walk in tqdm(dataloader):
-            loss = model(batch_walk)
-            if use_wandb:
-                wandb.log({"loss": loss.item()})
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-         
-    if use_wandb:
-        wandb.finish()
-    
-    node_embeddings = model.node_embed.weight.detach().to(device)
-
-    ### Extract & save item embeddings
-    item_embeddings = extract_item_embeddings(node_embeddings, bipartite_graph, graph)
-    item_embeddings = normalize_embeddings(item_embeddings)
-    np.save(embeddings_savepath, item_embeddings)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Prepare DeepWalk embeddings.")
-    parser.add_argument("--items_path", type=str, required=True, help="Path to the items file.")
-    parser.add_argument("--ratings_path", type=str, required=True, help="Path to the ratings file.")
-    parser.add_argument("--embeddings_savepath", type=str, required=True, help="Path to the file where embeddings will be saved.")
-    parser.add_argument("--emb_dim", type=int, default=384, help="Dimensionality of the embeddings.")
-    parser.add_argument("--window_size", type=int, default=4, help="Window size for the DeepWalk algorithm.")
-    parser.add_argument("--batch_size", type=int, default=512, help="Batch size for training.")
-    parser.add_argument("--lr", type=float, default=1e-2, help="Learning rate for training.")
-    parser.add_argument("--num_epochs", type=int, default=2, help="Number of epochs for training.")
-    parser.add_argument("--device", type=str, default="cpu", help="Device to use for training (cpu or cuda).")
-    parser.add_argument("--wandb_name", type=str, help="Name for WandB run.")
-    parser.add_argument("--no_wandb", action="store_false", dest="use_wandb", help="Disable WandB logging")
-    args = parser.parse_args()
-
-    prepare_deepwalk_embeddings(**vars(args))

exp/evaluate.py

@@ -1,5 +1,6 @@
 import argparse
 import json
+from tqdm.auto import tqdm
 
 import pandas as pd
 import numpy as np
@@ -9,16 +10,26 @@ from app.recommendations import RecommenderSystem
 
 def precision_at_k(recommended_items, relevant_items, k):
     recommended_at_k = set(recommended_items[:k])
-    relevant_set = set(relevant_items)
-    return len(recommended_at_k & relevant_set) / k
+    return len(recommended_at_k & relevant_items) / k
+
+
+def average_precision_at_k(recommended_items, relevant_items, k):
+    relevant_items = set(relevant_items)
+
+    apk_sum = 0.0
+    for m, item in enumerate(recommended_items):
+        if item in relevant_items:
+            apk_sum += precision_at_k(recommended_items, relevant_items, m+1)
+    
+    return apk_sum / min(k, len(relevant_items))
 
 
 def evaluate_recsys(
-    metrics_savepath,
     val_ratings_path, 
     faiss_index_path, 
     db_path,
-    n_recommend_items,
+    n_recommend_items=10,
+    metrics_savepath=None
 ):
     recsys = RecommenderSystem(
         faiss_index_path=faiss_index_path, 
@@ -30,16 +41,14 @@ def evaluate_recsys(
 
 
     metric_arrays = {
-        "precision@1": [],
-        "precision@3": [],
-        "precision@10": []
+        "ap@5": [],
     }
 
-    for item_group in grouped_items:
+    for item_group in tqdm(grouped_items):
         if len(item_group) == 1:
             continue
 
-        ### Precision@k is computed for each edge.
+        ### Metrics are computed for each edge.
         ### We will first aggregate it over all edges for user
         ### And after that - aggregate over all users
         user_metric_arrays = dict()
@@ -50,12 +59,8 @@ def evaluate_recsys(
             recommend_items = list(recsys.recommend_items(item, n_recommend_items))
             relevant_items = set(item_group) - {item}
 
-            user_metric_arrays["precision@1"].append(
-                precision_at_k(recommend_items, relevant_items, k=1))
-            user_metric_arrays["precision@3"].append(
-                precision_at_k(recommend_items, relevant_items, k=3))
-            user_metric_arrays["precision@10"].append(
-                precision_at_k(recommend_items, relevant_items, k=10))
+            user_metric_arrays["ap@5"].append(
+                average_precision_at_k(recommend_items, relevant_items, k=5))
 
         for metric in metric_arrays.keys():
             user_metric = np.mean(user_metric_arrays[metric])
@@ -65,9 +70,11 @@ def evaluate_recsys(
     for metric, array in metric_arrays.items():
         metrics[metric] = np.mean(array)
 
-    with open(metrics_savepath, "w") as f:
-        json.dump(metrics, f)
-    print(f"Saved metrics to {metrics_savepath}")
+    if metrics_savepath is not None:
+        with open(metrics_savepath, "w") as f:
+            json.dump(metrics, f)
+        print(f"Saved metrics to {metrics_savepath}")
+    return metrics
 
 
 if __name__ == "__main__":

exp/gnn.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+import tempfile
 
 import numpy as np
 import pandas as pd
@@ -8,7 +9,9 @@ import torch
 import wandb
 from tqdm.auto import tqdm
 
-from utils import prepare_graphs, normalize_embeddings, LRSchedule
+from exp.utils import prepare_graphs, normalize_embeddings, LRSchedule
+from exp.prepare_recsys import prepare_recsys
+from exp.evaluate import evaluate_recsys
 
 
 class GNNLayer(torch.nn.Module):
@@ -41,7 +44,6 @@ class GNNModel(torch.nn.Module):
             self,
             bipartite_graph,
             text_embeddings,
-            deepwalk_embeddings,
             num_layers,
             hidden_dim,
             aggregator_type,
@@ -56,14 +58,12 @@ class GNNModel(torch.nn.Module):
 
         self._bipartite_graph = bipartite_graph
         self._text_embeddings = text_embeddings
-        self._deepwalk_embeddings = deepwalk_embeddings
 
         self._sampler = dgl.sampling.PinSAGESampler(
             bipartite_graph, "Item", "User", num_traversals, 
             termination_prob, num_random_walks, num_neighbor)
 
         self._text_encoder = torch.nn.Linear(text_embeddings.shape[-1], hidden_dim)
-        self._deepwalk_encoder = torch.nn.Linear(deepwalk_embeddings.shape[-1], hidden_dim)
 
         self._layers = torch.nn.ModuleList()
         for _ in range(num_layers):
@@ -92,13 +92,10 @@ class GNNModel(torch.nn.Module):
         sampled_subgraph = self._sample_subraph(ids)
         sampled_subgraph = dgl.compact_graphs(sampled_subgraph, always_preserve=ids)
 
-        ### Encode text & DeepWalk embeddings
+        ### Encode text embeddings
         text_embeddings = self._text_embeddings[
             sampled_subgraph.ndata[dgl.NID]]
-        deepwalk_embeddings = self._deepwalk_embeddings[
-            sampled_subgraph.ndata[dgl.NID]]
-        features = self._text_encoder(text_embeddings) \
-            + self._deepwalk_encoder(deepwalk_embeddings)
+        features = self._text_encoder(text_embeddings)
 
         ### GNN goes brr...
         for layer in self._layers:
@@ -142,12 +139,27 @@ def sample_item_batch(user_batch, bipartite_graph):
     return item_batch
 
 
+@torch.no_grad()
+def inference_model(model, bipartite_graph, batch_size, hidden_dim, device):
+    model.eval()
+    item_embeddings = torch.zeros(bipartite_graph.num_nodes("Item"), hidden_dim).to(device)
+    for items_batch in tqdm(torch.utils.data.DataLoader(
+            torch.arange(bipartite_graph.num_nodes("Item")), 
+            batch_size=batch_size, 
+            shuffle=True
+    )):
+        item_embeddings[items_batch] = model(items_batch.to(device))
+
+    item_embeddings = normalize_embeddings(item_embeddings.cpu().numpy())
+    return item_embeddings
+
+
 def prepare_gnn_embeddings(
         # Paths
         items_path,
-        ratings_path,
+        train_ratings_path,
+        val_ratings_path,
         text_embeddings_path,
-        deepwalk_embeddings_path,
         embeddings_savepath, 
         # Learning hyperparameters
         temperature,
@@ -165,12 +177,13 @@ def prepare_gnn_embeddings(
         num_random_walks, 
         num_neighbor,
         # Misc
+        validate_every_n_epoch,
         device, 
         wandb_name, 
         use_wandb,
 ):
     ### Prepare graph
-    bipartite_graph, _ = prepare_graphs(items_path, ratings_path)
+    bipartite_graph, _ = prepare_graphs(items_path, train_ratings_path)
     bipartite_graph = bipartite_graph.to(device)
 
     ### Init wandb
@@ -179,11 +192,9 @@ def prepare_gnn_embeddings(
 
     ### Prepare model
     text_embeddings = torch.tensor(np.load(text_embeddings_path)).to(device)
-    deepwalk_embeddings = torch.tensor(np.load(deepwalk_embeddings_path)).to(device)
     model = GNNModel(
         bipartite_graph=bipartite_graph, 
         text_embeddings=text_embeddings, 
-        deepwalk_embeddings=deepwalk_embeddings,
         num_layers=num_layers,
         hidden_dim=hidden_dim,
         aggregator_type=aggregator_type,
@@ -214,6 +225,7 @@ def prepare_gnn_embeddings(
     ### Train loop
     model.train()
     for epoch in range(num_epochs):
+        ### Train
         for user_batch in tqdm(dataloader):
             item_batch = sample_item_batch(user_batch, bipartite_graph)  # (2, |user_batch|)
             item_batch = item_batch.reshape(-1)  # (2 * |user_batch|)
@@ -226,24 +238,29 @@ def prepare_gnn_embeddings(
             loss.backward()
             optimizer.step()
             lr_scheduler.step()
+        ### Validation
+        if (validate_every_n_epoch is not None) and (((epoch + 1) % validate_every_n_epoch) == 0):
+            item_embeddings = inference_model(
+                model, bipartite_graph, batch_size, hidden_dim, device)
+            with tempfile.TemporaryDirectory() as tmp_dir_name:
+                tmp_embeddings_path = os.path.join(tmp_dir_name, "embeddings.npy")
+                np.save(tmp_embeddings_path, item_embeddings)
+                prepare_recsys(items_path, tmp_embeddings_path, tmp_dir_name)
+                metrics = evaluate_recsys(
+                    val_ratings_path, 
+                    os.path.join(tmp_dir_name, "index.faiss"),
+                    os.path.join(tmp_dir_name, "items.db"))
+                print(f"Epoch {epoch + 1} / {num_epochs}. {metrics}")
+                if use_wandb:
+                    wandb.log(metrics)
+            
+        
          
     if use_wandb:
         wandb.finish()
     
     ### Process full dataset
-    model.eval()
-    with torch.no_grad():
-        hidden_dim = text_embeddings.shape[-1]
-        item_embeddings = torch.zeros(bipartite_graph.num_nodes("Item"), hidden_dim).to(device)
-        for items_batch in tqdm(torch.utils.data.DataLoader(
-                torch.arange(bipartite_graph.num_nodes("Item")), 
-                batch_size=batch_size, 
-                shuffle=True
-        )):
-            item_embeddings[items_batch] = model(items_batch.to(device))
-
-    ### Extract & save item embeddings
-    item_embeddings = normalize_embeddings(item_embeddings.cpu().numpy())
+    item_embeddings = inference_model(model, bipartite_graph, batch_size, hidden_dim, device)
     np.save(embeddings_savepath, item_embeddings)
 
 
@@ -252,9 +269,9 @@ if __name__ == "__main__":
 
     # Paths
     parser.add_argument("--items_path", type=str, required=True, help="Path to the items file")
-    parser.add_argument("--ratings_path", type=str, required=True, help="Path to the ratings file")
+    parser.add_argument("--train_ratings_path", type=str, required=True, help="Path to the train ratings file")
+    parser.add_argument("--val_ratings_path", type=str, required=True, help="Path to the validation ratings file")
     parser.add_argument("--text_embeddings_path", type=str, required=True, help="Path to the text embeddings file")
-    parser.add_argument("--deepwalk_embeddings_path", type=str, required=True, help="Path to the deepwalk embeddings file")
     parser.add_argument("--embeddings_savepath", type=str, required=True, help="Path to the file where gnn embeddings will be saved")
 
     # Learning hyperparameters
@@ -265,7 +282,7 @@ if __name__ == "__main__":
 
     # Model hyperparameters
     parser.add_argument("--num_layers", type=int, default=2, help="Number of layers in the model")
-    parser.add_argument("--hidden_dim", type=int, default=384, help="Hidden dimension size")
+    parser.add_argument("--hidden_dim", type=int, default=64, help="Hidden dimension size")
     parser.add_argument("--aggregator_type", type=str, default="mean", help="Type of aggregator in SAGEConv")
     parser.add_argument("--no_skip_connection", action="store_false", dest="skip_connection", help="Disable skip connections")
     parser.add_argument("--no_bidirectional", action="store_false", dest="bidirectional", help="Do not use reversed edges in convolution")
@@ -275,6 +292,7 @@ if __name__ == "__main__":
     parser.add_argument("--num_neighbor", type=int, default=3, help="Number of neighbors in PinSAGE-like sampler")
 
     # Misc
+    parser.add_argument("--validate_every_n_epoch", type=int, default=2, help="Perform RecSys validation every n train epochs.")
     parser.add_argument("--device", type=str, default="cpu", help="Device to run the model on (cpu or cuda)")
     parser.add_argument("--wandb_name", type=str, help="WandB run name")
     parser.add_argument("--no_wandb", action="store_false", dest="use_wandb", help="Disable WandB logging")

exp/prepare_embeddings.sh

@@ -10,14 +10,7 @@ echo Running on "$device".
 PYTHONPATH=. python exp/process_raw_data.py \
     --input_directory "$input_directory" \
     --save_directory "$save_directory" \
-    --create_train_val_split
-
-PYTHONPATH=. python exp/deepwalk.py \
-    --items_path "$save_directory/items.csv" \
-    --ratings_path "$save_directory/train_ratings.csv" \
-    --embeddings_savepath "$save_directory/deepwalk_embeddings.npy" \
-    --device $device \
-    --no_wandb
+    --create_train_val_test_split
 
 PYTHONPATH=. python exp/sbert.py \
     --items_path "$save_directory/items.csv" \
@@ -26,25 +19,21 @@ PYTHONPATH=. python exp/sbert.py \
 
 PYTHONPATH=. python exp/gnn.py \
     --items_path "$save_directory/items.csv" \
-    --ratings_path "$save_directory/train_ratings.csv" \
+    --train_ratings_path "$save_directory/train_ratings.csv" \
+    --val_ratings_path "$save_directory/val_ratings.csv" \
     --text_embeddings_path "$save_directory/text_embeddings.npy" \
-    --deepwalk_embeddings_path "$save_directory/deepwalk_embeddings.npy" \
     --embeddings_savepath "$save_directory/embeddings.npy"\
     --device $device \
-    --no_wandb 
-
-PYTHONPATH=. python exp/prepare_index.py \
-    --embeddings_path "$save_directory/embeddings.npy" \
-    --save_path "$save_directory/index.faiss"
+    --no_wandb
 
-PYTHONPATH=. python exp/prepare_db.py \
+PYTHONPATH=. python exp/prepare_recsys.py \
     --items_path "$save_directory/items.csv" \
     --embeddings_path "$save_directory/embeddings.npy" \
-    --db_path "$save_directory/items.db"
+    --save_directory "$save_directory"
 
 PYTHONPATH=. python exp/evaluate.py \
     --metrics_savepath "$save_directory/metrics.json" \
-    --val_ratings_path "$save_directory/val_ratings.csv" \
+    --val_ratings_path "$save_directory/test_ratings.csv" \
     --faiss_index_path "$save_directory/index.faiss" \
     --db_path "$save_directory/items.db"
 

exp/prepare_index.py

@@ -1,20 +0,0 @@
-import argparse
-
-import faiss
-import numpy as np
-
-
-def build_index(embeddings_path, save_path, n_neighbors):
-    embeddings = np.load(embeddings_path)
-    index = faiss.IndexHNSWFlat(embeddings.shape[-1], 32)
-    index.add(embeddings)
-    faiss.write_index(index, save_path)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Build an HNSW index from embeddings.")
-    parser.add_argument("--embeddings_path", required=True, type=str, help="Path to the embeddings file.")
-    parser.add_argument("--save_path", type=str, required=True, help="Path to save the built index.")
-    parser.add_argument("--n_neighbors", type=int, default=32, help="Number of neighbors for the index.")
-    args = parser.parse_args()
-    build_index(**vars(args))

exp/{prepare_db.py → prepare_recsys.py}

@@ -1,7 +1,9 @@
 import argparse
 import sqlite3
 import io
+import os
 
+import faiss
 import pandas as pd
 import numpy as np
 
@@ -23,11 +25,30 @@ def prepare_items_db(items_path, embeddings_path, db_path):
         items.to_sql("items", conn, if_exists="replace", index=False, dtype={"embedding": "embedding"})
 
 
+def build_index(embeddings_path, save_path, n_neighbors):
+    embeddings = np.load(embeddings_path)
+    index = faiss.IndexHNSWFlat(embeddings.shape[-1], n_neighbors)
+    index.add(embeddings)
+    faiss.write_index(index, save_path)
+
+
+def prepare_recsys(
+    items_path,
+    embeddings_path,
+    save_directory,
+    n_neighbors=32,
+):
+    prepare_items_db(items_path, embeddings_path, os.path.join(save_directory, "items.db"))
+    build_index(embeddings_path, os.path.join(save_directory, "index.faiss"), n_neighbors)
+
+
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Prepare items database from a CSV file.")
+    parser = argparse.ArgumentParser(description="Prepare items database and HNSW index from a CSV file and embeddings.")
+
     parser.add_argument("--items_path", required=True, type=str, help="Path to the CSV file containing items.")
     parser.add_argument("--embeddings_path", required=True, type=str, help="Path to the .npy file containing item embeddings.")
-    parser.add_argument("--db_path", required=True, type=str, help="Path to the SQLite database file.")
+    parser.add_argument("--save_directory", required=True, type=str, help="Path to the save directory.")
+    parser.add_argument("--n_neighbors", type=int, default=32, help="Number of neighbors for the index.")
 
     args = parser.parse_args()
-    prepare_items_db(**vars(args))
+    prepare_recsys(**vars(args))

exp/process_raw_data.py

@@ -80,37 +80,46 @@ def process_raw_data_goodreads(input_directory, save_directory, positive_rating_
     ratings.to_csv(os.path.join(save_directory, "ratings.csv"), index=False)
 
 
-def create_train_val_split(ratings_path, train_savepath, val_savepath, seed=42):
+def create_train_val_test_split(ratings_path, train_savepath, val_savepath, test_savepath, seed=42):
     ratings = pd.read_csv(ratings_path)
     user_ids = ratings["user_id"].unique()
 
     rng = np.random.default_rng(seed=seed)
-    train_size = int(len(user_ids) * 0.9)
-    train_indices = rng.choice(user_ids, size=train_size, replace=False)
+    train_size = int(len(user_ids) * 0.7)
+    val_size = int(len(user_ids) * 0.15)
+
+    indices = rng.permutation(user_ids)
+    train_indices = indices[:train_size]
+    val_indices = indices[train_size:train_size+val_size]
+    test_indices = indices[train_size+val_size:]
 
     train_data = ratings.loc[ratings["user_id"].isin(train_indices)]
-    val_data = ratings.loc[~ratings["user_id"].isin(train_indices)]
+    val_data = ratings.loc[ratings["user_id"].isin(val_indices)]
+    test_data = ratings.loc[ratings["user_id"].isin(test_indices)]
 
     print(f"Train size: {len(train_data)}.")
     print(f"Validation size: {len(val_data)}.")
+    print(f"Test size: {len(test_data)}.")
 
     train_data.to_csv(train_savepath, index=False)
     val_data.to_csv(val_savepath, index=False)
+    test_data.to_csv(test_savepath, index=False)
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Process raw data.")
     parser.add_argument("--input_directory", required=True, type=str, help="Directory containing the raw data.")
     parser.add_argument("--save_directory", required=True, type=str, help="Directory where processed data will be saved.")
-    parser.add_argument("--create_train_val_split", action="store_true", help="Flag to indicate whether to create a train-validation split.")
+    parser.add_argument("--create_train_val_test_split", action="store_true", help="Flag to indicate whether to create a train-validation split.")
     args = parser.parse_args()
 
     print("Processing raw data...")
     process_raw_data_goodreads(args.input_directory, args.save_directory)
-    if args.create_train_val_split:
-        create_train_val_split(
+    if args.create_train_val_test_split:
+        create_train_val_test_split(
             os.path.join(args.save_directory, "ratings.csv"),
             os.path.join(args.save_directory, "train_ratings.csv"),
-            os.path.join(args.save_directory, "val_ratings.csv")
+            os.path.join(args.save_directory, "val_ratings.csv"),
+            os.path.join(args.save_directory, "test_ratings.csv")
         )
     print("The raw data has been successfully processed.")

exp/requirements.txt

@@ -1,6 +1,7 @@
 -r ../requirements.txt  # install base requirements of app
-dgl==2.1.0
-torch==2.1.2
+-f https://data.dgl.ai/wheels/torch-2.2/repo.html
+dgl==1.1.3
+torch==2.2.1
 wandb==0.17.0
 tqdm==4.66.4
 pydantic==2.5.3

exp/requirements_gpu.txt

@@ -1,7 +1,7 @@
 -f https://data.dgl.ai/wheels/cu121/repo.html
 -r ../requirements.txt  # install base requirements of app
-dgl==2.1.0
-torch==2.1.0
+dgl==1.1.3
+torch==2.2.1
 wandb==0.17.0
 tqdm==4.66.4
 pydantic==2.5.3

exp/sbert.py

@@ -19,6 +19,7 @@ def prepare_sbert_embeddings(
     items = pd.read_csv(items_path).sort_values("item_id")
     sentences = items["description"].values
     model = SentenceTransformer(model_name).to(device)
+    model.eval()
     embeddings = []
     for start_index in tqdm(range(0, len(sentences), batch_size)):
         batch = sentences[start_index:start_index+batch_size]