Trainin' stuff

scripts/split_follows.py

···
 import os
 import random
 import sys
+from typing import Optional
 
 import numpy as np
 
+from scripts.utils import get_logger
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
+logger = get_logger(__name__)
 
-# Create formatter
-formatter = logging.Formatter(
-    "%(asctime)s | %(levelname)s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
-)
 
-# Console handler
-console_handler = logging.StreamHandler(sys.stdout)
-console_handler.setFormatter(formatter)
-logger.addHandler(console_handler)
-
-
-def main(follow_dir: str, output_dir: str, val_split: float, file_size: int):
+def main(
+    follow_dir: str,
+    output_dir: str,
+    val_split: float,
+    file_size: int,
+    max_accounts: Optional[int],
+    exclude_unfetched: bool = False,
+):
     if val_split >= 1.0 or val_split <= 0.0:
         raise ValueError("Validation split must be between 0 and 1 exclusive")
 
···
 
     did_id_map: dict[str, int] = dict()
 
+    files = os.listdir(follow_dir)
+    if max_accounts:
+        files = files[:max_accounts]
+
+    if exclude_unfetched:
+        for file in files:
+            did_id_map[file[-3]] = len(did_id_map)
+
     logger.info("Reading follows...")
     train_file_idx = 1
     val_file_idx = 1
···
     val_follows: list[tuple[int, int]] = []
     train_files = []
     val_files = []
-    for file in os.listdir(follow_dir):
+    for file in files:
         source_did = file[:-3]  # Remove .gz extension
+        if max_accounts and len(did_id_map) >= max_accounts:
+            continue
+
         if source_did not in did_id_map:
             did_id_map[source_did] = len(did_id_map)
 
         for line in gzip.open(follow_dir + file, "rt"):
             target_did = json.loads(line)["value"]["subject"]
+            if exclude_unfetched and target_did not in did_id_map:
+                continue
+
+            if max_accounts and len(did_id_map) >= max_accounts:
+                continue
+
             if target_did not in did_id_map:
                 did_id_map[target_did] = len(did_id_map)
 
···
         type=int,
         help="Max rows per output file",
     )
+    parser.add_argument(
+        "--exclude-unfetched",
+        dest="exclude_unfetched",
+        required=False,
+        type=bool,
+        help="Whether to include accounts whose follows haven't been retrieved",
+    )
+    parser.add_argument(
+        "--max-accounts",
+        dest="max_accounts",
+        required=False,
+        type=int,
+        help="Maximum number of accounts to include",
+    )
     args = parser.parse_args()
 
     main(
···
         output_dir=args.output_dir,
         val_split=args.val_split,
         file_size=args.file_size,
+        exclude_unfetched=args.exclude_unfetched,
+        max_accounts=args.max_accounts,
     )

scripts/training/data.py

···
+from functools import reduce
 import json
 import glob
+from multiprocessing import Pool
 import os
 import random
 
···
 from torch.utils.data import Dataset
 
 
+from scripts.utils import get_logger
+
+logger = get_logger(__name__)
+
+
+def get_frequencies(array: np.ndarray, user_count: int) -> np.ndarray:
+    frequencies = np.zeros((user_count, 2))
+    for source_id, target_id in array:
+        frequencies[source_id][0] += 1
+        frequencies[target_id][1] += 1
+    return frequencies
+
+
 class FollowDataset(Dataset):
     def __init__(
         self,
         dataset_path: str,
         split: str,
-        negative_sample_count: int,
+        negative_sample_chance: float,
     ):
         with open(os.path.join(dataset_path, "metadata.json"), "r") as in_file:
             metadata = json.load(in_file)
···
             raise ValueError(f"Could not find split in metadata file: {split}")
         split_files = metadata["splits"][split]
 
-        self.parquet_files: list[tuple[str, str, int, int]] = []
+        self.numpy_files: list[tuple[str, str, int, int]] = []
         for file in split_files:
             file_idx = file["filename"].split("_")[2].split(".")[0]
-            self.parquet_files.append((file["filename"], file["dtype"], file_idx, file["shape"][0]))  # type: ignore
+            self.numpy_files.append((file["filename"], file["dtype"], file_idx, file["shape"][0]))  # type: ignore
 
-        self.parquet_files.sort(key=lambda x: x[1])
+        self.numpy_files.sort(key=lambda x: x[1])
         self.dataframes: list[np.ndarray] = []
-        for filename, dtype, _, row_count in self.parquet_files:
+        for filename, dtype, _, row_count in self.numpy_files:
             self.dataframes.append(
                 np.memmap(
                     os.path.join(dataset_path, filename),
···
                 )
             )
 
-        self.negative_sample_count = negative_sample_count
+        logger.info("Calculating node frequency...")
+        with Pool(7) as p:
+            self.cumulative_freq = reduce(
+                np.add,
+                p.starmap(
+                    get_frequencies,
+                    [
+                        (dataframe, len(self.did_id_map))
+                        for dataframe in self.dataframes
+                    ],
+                ),
+            )
+            self.cumulative_freq[0] = self.cumulative_freq[0].cumsum()
+            self.cumulative_freq[1] = self.cumulative_freq[1].cumsum()
+
+        self.negative_sample_chance = negative_sample_chance
 
     def __len__(self) -> int:
-        return sum((row_count for _, _, _, row_count in self.parquet_files))
+        return sum((row_count for _, _, _, row_count in self.numpy_files))
 
     def num_users(self) -> int:
         return len(self.did_id_map)
···
         row_index_total = 0
         effective_idx = idx
         i = 0
-        for i, (_, _, _, row_count) in enumerate(self.parquet_files):
+        for i, (_, _, _, row_count) in enumerate(self.numpy_files):
             row_index_total += row_count
             if idx < row_index_total:
                 break
···
         row = self.dataframes[i][effective_idx]
         return (row[0].item(), row[1].item())
 
-    def __getitem__(self, idx: int) -> tuple[list[tuple[int, int]], list[int]]:
+    def __getitem__(self, idx: int) -> tuple[tuple[int, int], int]:
         """
-        Grab follow connection and create negative samples (follows that dont exist)
+        Grab follow connection and corrupt it at defined frequency to another id
+        weighted by that id's prevalence in the dataset
         """
-        samples = [self._idx_to_row(idx)]
-
-        for _ in range(self.negative_sample_count):
-            rand_idx = random.randrange(0, len(self.did_id_map))
-            # Assume that users arent connected (roughly 1/700 chance they are)
-            samples.append((samples[0][0], rand_idx))
-
-        labels = [-1] * len(samples)
-        labels[0] = 1
-        return samples, labels
+        sample = self._idx_to_row(idx)
+        return (sample, 1)
 
-    @staticmethod
     def collate_rows(
-        batch: list[tuple[list[tuple[int, int]], list[int]]],
+        self,
+        batch: list[tuple[tuple[int, int], int]],
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        follows = torch.concat(tuple(torch.IntTensor(follow) for follow, _ in batch))
-        labels = torch.concat(tuple(torch.IntTensor(label) for _, label in batch))
+        # Corrupt some rows into negative edges
+        corrupted_sources = []
+        corrupted_targets = []
+        for i in range(len(batch)):
+            if random.random() < self.negative_sample_chance:
+                if random.random() < 0.5:
+                    corrupted_sources.append(i)
+                else:
+                    corrupted_targets.append(i)
+
+        new_sources = self.cumulative_freq[:, 0].searchsorted(
+            np.random.sample(len(corrupted_sources)) * self.cumulative_freq.shape[0]
+        )
+        new_targets = self.cumulative_freq[:, 1].searchsorted(
+            np.random.sample(len(corrupted_targets)) * self.cumulative_freq.shape[0]
+        )
+        new_sources[new_sources >= self.cumulative_freq.shape[0]] = (
+            self.cumulative_freq.shape[0] - 1
+        )
+        new_targets[new_targets >= self.cumulative_freq.shape[0]] = (
+            self.cumulative_freq.shape[0] - 1
+        )
+
+        for i, idx in enumerate(corrupted_sources):
+            batch[idx] = ((new_sources[i], batch[idx][0][1]), -1)
+
+        for i, idx in enumerate(corrupted_targets):
+            batch[idx] = ((batch[idx][0][0], new_targets[i]), -1)
+
+        follows = torch.concat(tuple(torch.IntTensor([follow]) for follow, _ in batch))
+        labels = torch.concat(tuple(torch.IntTensor([label]) for _, label in batch))
         return (follows, labels)

scripts/training/models.py

···
 import torch
 from torch import nn
-from torch.utils.data import DataLoader
 import torch.nn.functional as F
 import lightning as pyl
-
-from scripts.training.data import FollowDataset
 
 
 class UserEmbedding(nn.Module):
···
 
 
 class FollowEmbedModule(pyl.LightningModule):
-    def __init__(self, source_embedding, target_embedding):
+    def __init__(self, num_embeds: int, embed_dim: int, learning_rate: float):
         super().__init__()
-        self.source_embed = source_embedding
-        self.target_embed = target_embedding
+        self.save_hyperparameters()
+        self.source_embed = UserEmbedding(num_embeds, embed_dim)
+        self.target_embed = UserEmbedding(num_embeds, embed_dim)
+        self.learning_rate = learning_rate
 
     def training_step(self, batch, _):
         x, y = batch
···
             print(f"\nTraining Loss: {avg_loss:.4f}\n")
 
     def configure_optimizers(self):
-        optimizers = torch.optim.Adagrad(self.parameters(), lr=1e-1)
+        optimizers = torch.optim.Adagrad(self.parameters(), lr=self.learning_rate)
         return optimizers

scripts/training/serve.py

···
+from argparse import ArgumentParser
+import json
+
+from scripts.training.models import UserEmbedding, FollowEmbedModule
+
+
+def main(checkpoint_path: str, id_map_path: str):
+    model = FollowEmbedModule.load_from_checkpoint(checkpoint_path)
+
+    print(f"Shape of source embed weights: {model.source_embed.embed.weight.shape}")
+    with open(id_map_path, "r") as in_file:
+        did_id_map = json.load(in_file)
+
+    did = "did:plc:uxmy3zztxyhfk6mxrkun5tpr"
+    user_id = did_id_map[did]
+    print(f"did: {did}")
+    print(f"id:  {did_id_map[did]}")
+
+    source_tensor = model.source_embed.embed.weight[user_id, :]
+    print(f"Embedding: {source_tensor}")
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--checkpoint",
+        required=True,
+        help="Path to model checkpoint file",
+    )
+    parser.add_argument(
+        "--did-id-map",
+        dest="did_id_map",
+        required=True,
+        help="Path to did_id_map file",
+    )
+    args = parser.parse_args()
+
+    main(checkpoint_path=args.checkpoint, id_map_path=args.did_id_map)

scripts/training/train.py

···
 from torchdata.stateful_dataloader import StatefulDataLoader
 
 from scripts.training.data import FollowDataset
-from scripts.training.models import UserEmbedding, FollowEmbedModule
+from scripts.training.models import FollowEmbedModule
+from scripts.utils import get_logger
 
+logger = get_logger(__name__)
 
 PROJECT = "goodposts-followers"
 SAVE_DIR = "./data/training"
 NAME = "following-embedding"
 
 
-def main(run_id: Optional[str], model_checkpoint: Optional[str]):
+def main(run_id: Optional[str], model_checkpoint: Optional[str], dataset_path: str):
     if run_id is None:
         run_id = NAME + str(int(datetime.now().timestamp()))
 
-    print("Loading validation dataset")
-    val_dataset = FollowDataset(
-        dataset_path="data/processed/numpy/",
-        split="validation",
-        negative_sample_count=0,
-    )
-
-    print("Loading training dataset")
+    logger.info("Loading training dataset")
     train_dataset = FollowDataset(
-        dataset_path="data/processed/numpy/",
+        dataset_path=dataset_path,
         split="training",
-        negative_sample_count=2,
+        negative_sample_chance=0.5,
     )
 
-    val_dataloader = StatefulDataLoader(
-        val_dataset,
-        pin_memory=True,
-        collate_fn=val_dataset.collate_rows,
-        batch_size=1024,
-        num_workers=7,
+    logger.info("Loading validation dataset")
+    val_dataset = FollowDataset(
+        dataset_path=dataset_path,
+        split="validation",
+        negative_sample_chance=0.0,
     )
 
     train_dataloader = StatefulDataLoader(
···
         collate_fn=train_dataset.collate_rows,
         batch_size=1024,
         num_workers=7,
+        shuffle=True,
     )
 
-    print("Instantiating embedding model")
+    val_dataloader = StatefulDataLoader(
+        val_dataset,
+        pin_memory=True,
+        collate_fn=val_dataset.collate_rows,
+        batch_size=1024,
+        num_workers=7,
+    )
+
+    logger.info("Instantiating embedding model")
     wandb_logger = WandbLogger(
         project=PROJECT,
         save_dir=SAVE_DIR,
···
         monitor="val_loss",
     )
 
-    embedding = FollowEmbedModule(
-        source_embedding=UserEmbedding(train_dataset.num_users(), 64),
-        target_embedding=UserEmbedding(train_dataset.num_users(), 64),
-    )
+    if model_checkpoint:
+        embedding = FollowEmbedModule.load_from_checkpoint(
+            checkpoint_path=model_checkpoint
+        )
+    else:
+        embedding = FollowEmbedModule(
+            num_embeds=train_dataset.num_users(),
+            embed_dim=256,
+            learning_rate=1e-1,
+        )
 
     torch.set_float32_matmul_precision("medium")
     trainer = pyl.Trainer(
···
         accelerator="gpu",
         devices=1,
         precision="bf16",
-        max_epochs=2,
-        limit_train_batches=0.01,
-        limit_val_batches=1_000,
-        val_check_interval=0.05,
+        gradient_clip_val=0.5,
+        max_epochs=10,
+        val_check_interval=0.5,
     )
     trainer.fit(
         model=embedding,
···
         dest="model_checkpoint",
         help="Path to checkpoint with run data",
     )
+    parser.add_argument(
+        "-d",
+        "--dataset-path",
+        dest="dataset_path",
+        help="Path to dataset directory",
+    )
 
     args = parser.parse_args()
-    main(run_id=args.run_id, model_checkpoint=args.model_checkpoint)
+    main(
+        run_id=args.run_id,
+        model_checkpoint=args.model_checkpoint,
+        dataset_path=args.dataset_path,
+    )

scripts/utils.py

···
     accts = [(acct, follows) for acct, follows in to_explore.items()]
     accts.sort(key=lambda x: -1 * x[1])
     return accts
+
+
+def get_logger(name: str):
+
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO)
+
+    # Create formatter
+    formatter = logging.Formatter(
+        "%(asctime)s | %(levelname)s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+    )
+
+    # Console handler
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+    return logger