""" Synthetic Pinterest-style dataset generator. Generates: - users.parquet : user profiles with interest vectors - pins.parquet : pin metadata with category embeddings - interactions.parquet : save/click/close-up events with timestamps """ import os import yaml import numpy as np import pandas as pd from pathlib import Path from loguru import logger # ─── Pinterest-like Categories ─────────────────────────────────────────────── CATEGORIES = [ "home_decor", "fashion", "food_recipes", "travel", "fitness", "beauty", "art_crafts", "photography", "wedding", "parenting", "gardening", "technology", "architecture", "quotes", "hair", "tattoos", "cars", "animals", "music", "books", "outdoors", "sports", "education", "business", "minimalism", "vintage", "diy", "jewelry", "movies", "skincare", ] INTERACTION_TYPES = ["save", "click", "close_up", "hide"] INTERACTION_WEIGHTS = [0.35, 0.45, 0.15, 0.05] # realistic skew def load_config(config_path: str = "config.yaml") -> dict: with open(config_path) as f: return yaml.safe_load(f) def generate_users(num_users: int, num_categories: int, seed: int) -> pd.DataFrame: """Generate user profiles with latent interest vectors per category.""" rng = np.random.default_rng(seed) cat_names = CATEGORIES[:num_categories] # Each user has a sparse interest profile (Dirichlet → sparse via top-k) raw_interests = rng.dirichlet(np.ones(num_categories) * 0.5, size=num_users) # Sparsify: keep only top-5 categories per user sparse = np.zeros_like(raw_interests) top_k_idx = np.argsort(raw_interests, axis=1)[:, -5:] for i, idx in enumerate(top_k_idx): sparse[i, idx] = raw_interests[i, idx] sparse[i] /= sparse[i].sum() # re-normalize interest_df = pd.DataFrame(sparse, columns=[f"interest_{c}" for c in cat_names]) users = pd.DataFrame({ "user_id": np.arange(num_users), "account_age_days": rng.integers(1, 2000, size=num_users), "num_boards": rng.integers(1, 50, size=num_users), "num_pins_saved": rng.integers(10, 5000, size=num_users), "is_mobile": rng.choice([0, 1], size=num_users, p=[0.3, 0.7]), }) return pd.concat([users, interest_df], axis=1) def generate_pins(num_pins: int, num_categories: int, seed: int) -> pd.DataFrame: """Generate pin metadata with category distributions and feature vectors.""" rng = np.random.default_rng(seed + 1) cat_names = CATEGORIES[:num_categories] # Each pin belongs primarily to one category but has soft multi-label primary_cat = rng.integers(0, num_categories, size=num_pins) cat_features = np.zeros((num_pins, num_categories)) for i, pc in enumerate(primary_cat): cat_features[i, pc] = rng.uniform(0.6, 1.0) # add noise from 1-2 secondary categories sec = rng.choice([c for c in range(num_categories) if c != pc], size=rng.integers(1, 3), replace=False) cat_features[i, sec] = rng.uniform(0.05, 0.3, size=len(sec)) cat_feat_df = pd.DataFrame(cat_features, columns=[f"cat_{c}" for c in cat_names]) # Visual embedding (128-dim, simulating image encoder output) visual_embs = rng.standard_normal((num_pins, 128)).astype(np.float32) # Inject category signal: pins in same category cluster together for cat_id in range(num_categories): mask = primary_cat == cat_id centroid = rng.standard_normal(128) * 2 visual_embs[mask] += centroid visual_emb_df = pd.DataFrame( visual_embs, columns=[f"visual_{i}" for i in range(128)] ) pins = pd.DataFrame({ "pin_id": np.arange(num_pins), "primary_category": [cat_names[c] for c in primary_cat], "primary_category_id": primary_cat, "num_saves": rng.integers(0, 100000, size=num_pins), "num_clicks": rng.integers(0, 500000, size=num_pins), "is_promoted": rng.choice([0, 1], size=num_pins, p=[0.9, 0.1]), "has_price": rng.choice([0, 1], size=num_pins, p=[0.7, 0.3]), "content_length": rng.integers(10, 500, size=num_pins), }) return pd.concat([pins, cat_feat_df, visual_emb_df], axis=1) def generate_interactions( users: pd.DataFrame, pins: pd.DataFrame, num_interactions: int, num_categories: int, seed: int, ) -> pd.DataFrame: """ Generate user-pin interactions with interest-driven sampling. Users are more likely to interact with pins matching their interests. """ rng = np.random.default_rng(seed + 2) cat_names = CATEGORIES[:num_categories] num_users = len(users) num_pins = len(pins) user_interests = users[[f"interest_{c}" for c in cat_names]].values # (U, C) pin_cats = pins[[f"cat_{c}" for c in cat_names]].values # (P, C) # Build affinity matrix in chunks to avoid OOM rows, cols, types, ts = [], [], [], [] chunk = 5000 logger.info(f"Sampling {num_interactions} interactions...") sampled = 0 while sampled < num_interactions: n = min(chunk, num_interactions - sampled) uid = rng.integers(0, num_users, size=n) # Compute affinity scores for sampled users vs all pins scores = user_interests[uid] @ pin_cats.T # (n, P) scores = np.clip(scores, 1e-8, None) probs = scores / scores.sum(axis=1, keepdims=True) # Sample one pin per user proportional to affinity pid = np.array([rng.choice(num_pins, p=p) for p in probs]) itype = rng.choice( INTERACTION_TYPES, size=n, p=INTERACTION_WEIGHTS ) timestamp = rng.integers( pd.Timestamp("2023-01-01").value // 10**9, pd.Timestamp("2024-12-31").value // 10**9, size=n, ) rows.extend(uid.tolist()) cols.extend(pid.tolist()) types.extend(itype.tolist()) ts.extend(timestamp.tolist()) sampled += n interactions = pd.DataFrame({ "user_id": rows, "pin_id": cols, "interaction_type": types, "timestamp": ts, }) # Assign weights: save=3, click=1, close_up=2, hide=-1 weight_map = {"save": 3, "click": 1, "close_up": 2, "hide": -1} interactions["weight"] = interactions["interaction_type"].map(weight_map) # Drop duplicate (user, pin) keeping max weight interactions = ( interactions.sort_values("weight", ascending=False) .drop_duplicates(subset=["user_id", "pin_id"], keep="first") .reset_index(drop=True) ) # Keep only positive interactions for training interactions = interactions[interactions["weight"] > 0].reset_index(drop=True) logger.info(f"Generated {len(interactions)} positive interactions") return interactions def train_val_test_split( interactions: pd.DataFrame, train: float, val: float, seed: int ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: """ Temporal split: last interaction per user → test, second-last → val, rest → train. Mirrors real evaluation protocol for recommendation systems. """ interactions = interactions.sort_values(["user_id", "timestamp"]) test_idx = interactions.groupby("user_id").tail(1).index remaining = interactions.drop(test_idx) val_idx = remaining.groupby("user_id").tail(1).index train_df = remaining.drop(val_idx).reset_index(drop=True) val_df = interactions.loc[val_idx].reset_index(drop=True) test_df = interactions.loc[test_idx].reset_index(drop=True) logger.info( f"Split → train: {len(train_df)} | val: {len(val_df)} | test: {len(test_df)}" ) return train_df, val_df, test_df def main(config_path: str = "config.yaml"): cfg = load_config(config_path) dc = cfg["data"] out = Path(cfg["paths"]["data_dir"]) raw = out / "raw" proc = out / "processed" raw.mkdir(parents=True, exist_ok=True) proc.mkdir(parents=True, exist_ok=True) logger.info("Generating users...") users = generate_users(dc["num_users"], dc["num_categories"], dc["seed"]) users.to_parquet(raw / "users.parquet", index=False) logger.info(f" → {len(users)} users saved") logger.info("Generating pins...") pins = generate_pins(dc["num_pins"], dc["num_categories"], dc["seed"]) pins.to_parquet(raw / "pins.parquet", index=False) logger.info(f" → {len(pins)} pins saved") logger.info("Generating interactions...") interactions = generate_interactions( users, pins, dc["num_interactions"], dc["num_categories"], dc["seed"] ) interactions.to_parquet(raw / "interactions.parquet", index=False) logger.info("Splitting dataset...") train_df, val_df, test_df = train_val_test_split( interactions, dc["train_split"], dc["val_split"], dc["seed"] ) train_df.to_parquet(proc / "train.parquet", index=False) val_df.to_parquet(proc / "val.parquet", index=False) test_df.to_parquet(proc / "test.parquet", index=False) logger.info("✅ Dataset generation complete.") logger.info(f" Users: {len(users)} | Pins: {len(pins)} | Interactions: {len(interactions)}") if __name__ == "__main__": main()