create_preference_data_samples.py

import pandas as pd
from pathlib import Path
import random

# Set random seed for reproducibility
random.seed(42)

# Create output directory
output_dir = Path("data_samples")
output_dir.mkdir(parents=True, exist_ok=True)

# Train/test split ratio
TRAIN_RATIO = 0.5

# Number of popularity buckets for stratified sampling
NUM_POPULARITY_BUCKETS = 10


def calculate_item_popularity(user_items_dict, all_items):
    """Calculate popularity (interaction count) for each item.

    Args:
        user_items_dict: Dictionary mapping user_id to set of items they interacted with
        all_items: Set of all items in the dataset

    Returns:
        Dictionary mapping each item to its popularity count
    """
    # Initialize all items with 0 count
    item_popularity = {item: 0 for item in all_items}

    # Iterate through users and increment counts - much faster than iterating items
    for user_items in user_items_dict.values():
        for item in user_items:
            if item in item_popularity:
                item_popularity[item] += 1

    return item_popularity


def print_dataset_statistics(user_items_dict, all_items, item_popularity, item_to_bucket):
    """Print helpful statistics about the dataset and bucket distribution."""
    total_items = len(all_items)
    total_interactions = sum(len(items) for items in user_items_dict.values())
    total_users = len(user_items_dict)

    print(f"  Total items: {total_items:,}")
    print(f"  Total users: {total_users:,}")
    print(f"  Total interactions: {total_interactions:,}")
    print(f"  Avg interactions per item: {total_interactions / total_items:.2f}")

    # Bucket statistics
    num_buckets = max(item_to_bucket.values()) + 1
    bucket_stats = {i: {'items': 0, 'interactions': 0, 'min_pop': float('inf'), 'max_pop': 0}
                    for i in range(num_buckets)}

    for item, bucket_id in item_to_bucket.items():
        pop = item_popularity[item]
        bucket_stats[bucket_id]['items'] += 1
        bucket_stats[bucket_id]['interactions'] += pop
        bucket_stats[bucket_id]['min_pop'] = min(bucket_stats[bucket_id]['min_pop'], pop)
        bucket_stats[bucket_id]['max_pop'] = max(bucket_stats[bucket_id]['max_pop'], pop)

    print(f"\n  Bucket distribution ({num_buckets} buckets):")
    print(f"  {'Bucket':<8} {'Items':<10} {'Interactions':<15} {'Pop Range':<20} {'% of Total':<12}")
    print(f"  {'-'*8} {'-'*10} {'-'*15} {'-'*20} {'-'*12}")

    for bucket_id in range(num_buckets):
        stats = bucket_stats[bucket_id]
        pct = (stats['interactions'] / total_interactions * 100) if total_interactions > 0 else 0
        pop_range = f"{stats['min_pop']}-{stats['max_pop']}"
        print(f"  {bucket_id:<8} {stats['items']:<10,} {stats['interactions']:<15,} {pop_range:<20} {pct:<12.2f}%")
    print()


def create_popularity_buckets(user_items_dict, all_items, num_buckets=NUM_POPULARITY_BUCKETS):
    """Create popularity buckets for items based on interaction counts.

    Uses logarithmic bucketing to handle power-law distributions common in
    recommendation datasets. Items in the same bucket will have popularity
    within a small multiplicative factor of each other.

    Args:
        user_items_dict: Dictionary mapping user_id to set of items they interacted with
        all_items: Set of all items in the dataset
        num_buckets: Number of buckets to create based on log-popularity

    Returns:
        Tuple of (item_to_bucket dict, item_popularity dict)
    """
    import numpy as np

    # Count how many users interacted with each item
    item_popularity = calculate_item_popularity(user_items_dict, all_items)

    # Get min and max popularity (add 1 to avoid log(0))
    popularity_values = [pop for pop in item_popularity.values() if pop > 0]

    if not popularity_values:
        return {item: 0 for item in all_items}, item_popularity

    min_pop = min(popularity_values)
    max_pop = max(popularity_values)

    # Create logarithmic bucket boundaries
    # Items with popularity 0 go to bucket 0
    # Other items are distributed across buckets based on log(popularity)
    if min_pop == max_pop:
        # All items have same popularity - put them all in the same bucket
        item_to_bucket = {item: 0 for item in item_popularity.keys()}
    else:
        # Use log scale to create bucket boundaries
        log_min = np.log10(max(1, min_pop))
        log_max = np.log10(max_pop)

        # Create evenly spaced boundaries in log space
        log_boundaries = np.linspace(log_min, log_max, num_buckets + 1)

        # Convert back to linear space
        boundaries = [10 ** x for x in log_boundaries]

        # Assign items to buckets
        item_to_bucket = {}
        for item, pop in item_popularity.items():
            if pop == 0:
                item_to_bucket[item] = 0
            else:
                # Find the bucket for this popularity value
                bucket_id = num_buckets - 1
                for i in range(num_buckets):
                    if pop <= boundaries[i + 1]:
                        bucket_id = i
                        break
                item_to_bucket[item] = bucket_id

    return item_to_bucket, item_popularity


def get_users_with_min_items(user_items_dict, min_items=50, num_users=10):
    """Get users with at least min_items in their history."""
    eligible_users = [user for user, items in user_items_dict.items() if len(items) >= min_items]

    if len(eligible_users) < num_users:
        print(f"Warning: Only {len(eligible_users)} users with >= {min_items} items found")
        return eligible_users

    return random.sample(eligible_users, num_users)


def save_samples(samples, file_path):
    """Save samples to a CSV file."""
    if not samples:
        return

    df = pd.DataFrame(samples)
    df.to_csv(file_path, index=False)


def create_samples_from_items(user_id, positive_items_subset, negative_items, domain, item_to_bucket=None, item_names=None, num_samples=25):
    """Create preference questions from a subset of positive items.

    This ensures no data leakage - items in positive_items_subset won't appear
    in questions generated from a different subset.

    Args:
        item_to_bucket: Optional dictionary mapping items to popularity buckets.
                        If provided, negative samples will be chosen from the same
                        popularity bucket as the positive sample.
    """
    samples = []

    positive_items_list = list(positive_items_subset)
    negative_items_list = list(negative_items)

    if not negative_items_list or not positive_items_list:
        return []

    # Group negative items by popularity bucket if buckets are provided
    negative_items_by_bucket = {}
    if item_to_bucket is not None:
        for item in negative_items_list:
            bucket = item_to_bucket.get(item)
            if bucket is not None:
                if bucket not in negative_items_by_bucket:
                    negative_items_by_bucket[bucket] = []
                negative_items_by_bucket[bucket].append(item)

    # Create question based on domain
    domain_verbs = {
        'anime': 'watch',
        'book': 'read',
        'movie': 'watch',
        'steam': 'buy'
    }
    domain_nouns = {
        'anime': 'anime',
        'book': 'book',
        'movie': 'movie',
        'steam': 'game'
    }

    verb = domain_verbs[domain]
    noun = domain_nouns[domain]

    for _ in range(num_samples):
        # Get a random positive item from this subset only
        positive_item = random.choice(positive_items_list)

        # Get a random negative item from the same popularity bucket if possible
        if item_to_bucket is not None:
            positive_bucket = item_to_bucket.get(positive_item)
            if positive_bucket is not None and positive_bucket in negative_items_by_bucket and negative_items_by_bucket[positive_bucket]:
                negative_item = random.choice(negative_items_by_bucket[positive_bucket])
            else:
                # Fallback to random if bucket is empty or item not in bucket mapping
                negative_item = random.choice(negative_items_list)
                print("Fallback, couldn't find negative sample in bucket.")
        else:
            # Original behavior: random negative item
            negative_item = random.choice(negative_items_list)

        # Randomly assign positions
        if random.random() < 0.5:
            item_1, item_2 = positive_item, negative_item
            y_true = 0  # Positive item is in position 0
        else:
            item_1, item_2 = negative_item, positive_item
            y_true = 1  # Positive item is in position 1

        # Get item names if available
        if item_names is not None:
            item_1_name = item_names.get(item_1, item_1)
            item_2_name = item_names.get(item_2, item_2)
        else:
            item_1_name = item_1
            item_2_name = item_2

        prompt = (f'Is user {user_id} more likely to {verb} the {noun} "{item_1_name}" or "{item_2_name}"? '
                  f'Respond only in JSON in the format {{"choice": 0}}, where 0 represents the first {noun}, '
                  f'and 1 represents the second {noun}.')

        samples.append({'prompt': prompt, 'y_true': y_true})

    return samples


def create_train_test_samples(user_id, positive_items, all_items, domain, item_to_bucket=None, item_names=None, num_samples_per_split=25):
    """Create train and test samples with no data leakage.

    Splits BOTH positive and negative items to ensure zero overlap in
    user-item combinations between train and test sets.

    Args:
        item_to_bucket: Optional dictionary mapping items to popularity buckets.
                        If provided, negative samples will be chosen from the same
                        popularity bucket as the positive sample.
    """
    # Get negative items (items user did not interact with)
    negative_items = all_items - positive_items

    positive_items_list = list(positive_items)
    negative_items_list = list(negative_items)

    if len(positive_items_list) < 2 or len(negative_items_list) < 2:
        return [], []

    # Shuffle and split positive items into train and test sets
    random.shuffle(positive_items_list)
    split_idx = int(len(positive_items_list) * TRAIN_RATIO)

    train_positive_items = set(positive_items_list[:split_idx])
    test_positive_items = set(positive_items_list[split_idx:])

    # Shuffle and split negative items into train and test sets
    random.shuffle(negative_items_list)
    split_idx = int(len(negative_items_list) * TRAIN_RATIO)

    train_negative_items = set(negative_items_list[:split_idx])
    test_negative_items = set(negative_items_list[split_idx:])

    # Generate train samples using only train items (no overlap with test)
    train_samples = create_samples_from_items(
        user_id, train_positive_items, train_negative_items, domain, item_to_bucket, item_names, num_samples_per_split
    )

    # Generate test samples using only test items (no overlap with train)
    test_samples = create_samples_from_items(
        user_id, test_positive_items, test_negative_items, domain, item_to_bucket, item_names, num_samples_per_split
    )

    return train_samples, test_samples


def process_anime():
    """Process anime preference dataset."""
    print("Processing anime dataset...")

    # Load data
    ratings = pd.read_csv("datasets/anime_pref/rating.csv")
    anime = pd.read_csv("datasets/anime_pref/anime.csv")

    # Filter out negative ratings (consider only positive interactions)
    positive_ratings = ratings[ratings['rating'] > 0]

    # Create user -> items mapping (vectorized - much faster!)
    user_items = positive_ratings.groupby('user_id')['anime_id'].apply(set).to_dict()

    # Get all anime IDs that have at least one interaction (not all anime in catalog)
    all_anime = set(positive_ratings['anime_id'].unique())

    # Create anime name mapping
    anime_names = dict(zip(anime['anime_id'], anime['name']))

    # Create popularity buckets
    print("  Creating popularity buckets...")
    item_to_bucket, item_popularity = create_popularity_buckets(user_items, all_anime)

    # Print dataset statistics
    print_dataset_statistics(user_items, all_anime, item_popularity, item_to_bucket)

    # Get eligible users
    selected_users = get_users_with_min_items(user_items, min_items=50, num_users=10)

    # Generate samples
    for idx, user_id in enumerate(selected_users):
        train_samples, test_samples = create_train_test_samples(
            user_id,
            user_items[user_id],
            all_anime,
            'anime',
            item_to_bucket,
            anime_names,
            num_samples_per_split=25
        )

        sample_dir = output_dir / f"anime_sample_{idx+1}"
        sample_dir.mkdir(parents=True, exist_ok=True)

        save_samples(train_samples, sample_dir / "train.csv")
        save_samples(test_samples, sample_dir / "test.csv")

        print(f"  Created anime_sample_{idx+1}: train ({len(train_samples)} questions), test ({len(test_samples)} questions) for user {user_id}")


def process_book():
    """Process book preference dataset."""
    print("Processing book dataset...")

    # Load data
    ratings = pd.read_csv("datasets/book_pref/Ratings.csv", sep=';')
    books = pd.read_csv("datasets/book_pref/Books.csv", sep=';', low_memory=False)

    # Filter positive ratings (rating > 0)
    positive_ratings = ratings[ratings['Rating'] > 0]

    # Create user -> items mapping (vectorized - much faster!)
    user_items = positive_ratings.groupby('User-ID')['ISBN'].apply(set).to_dict()

    # Get all ISBNs that have at least one interaction (not all books in catalog)
    all_books = set(positive_ratings['ISBN'].unique())

    # Create book name mapping
    book_names = dict(zip(books['ISBN'], books['Title']))

    # Create popularity buckets
    print("  Creating popularity buckets...")
    item_to_bucket, item_popularity = create_popularity_buckets(user_items, all_books)

    # Print dataset statistics
    print_dataset_statistics(user_items, all_books, item_popularity, item_to_bucket)

    # Get eligible users
    selected_users = get_users_with_min_items(user_items, min_items=50, num_users=10)

    # Generate samples
    for idx, user_id in enumerate(selected_users):
        train_samples, test_samples = create_train_test_samples(
            user_id,
            user_items[user_id],
            all_books,
            'book',
            item_to_bucket,
            book_names,
            num_samples_per_split=25
        )

        sample_dir = output_dir / f"book_sample_{idx+1}"
        sample_dir.mkdir(parents=True, exist_ok=True)

        save_samples(train_samples, sample_dir / "train.csv")
        save_samples(test_samples, sample_dir / "test.csv")

        print(f"  Created book_sample_{idx+1}: train ({len(train_samples)} questions), test ({len(test_samples)} questions) for user {user_id}")


def process_movie():
    """Process movie preference dataset."""
    print("Processing movie dataset...")

    # Load data
    ratings = pd.read_csv("datasets/movie_pref/ratings.csv")
    movies = pd.read_csv("datasets/movie_pref/movies.csv")

    positive_ratings = ratings[ratings['rating'] >= 0]

    # Create user -> items mapping (vectorized - much faster!)
    user_items = positive_ratings.groupby('userId')['movieId'].apply(set).to_dict()

    # Get all movie IDs that have at least one interaction (not all movies in catalog)
    all_movies = set(positive_ratings['movieId'].unique())

    # Create movie name mapping
    movie_names = dict(zip(movies['movieId'], movies['title']))

    # Create popularity buckets
    print("  Creating popularity buckets...")
    item_to_bucket, item_popularity = create_popularity_buckets(user_items, all_movies)

    # Print dataset statistics
    print_dataset_statistics(user_items, all_movies, item_popularity, item_to_bucket)

    # Get eligible users
    selected_users = get_users_with_min_items(user_items, min_items=50, num_users=10)

    # Generate samples
    for idx, user_id in enumerate(selected_users):
        train_samples, test_samples = create_train_test_samples(
            user_id,
            user_items[user_id],
            all_movies,
            'movie',
            item_to_bucket,
            movie_names,
            num_samples_per_split=25
        )

        sample_dir = output_dir / f"movie_sample_{idx+1}"
        sample_dir.mkdir(parents=True, exist_ok=True)

        save_samples(train_samples, sample_dir / "train.csv")
        save_samples(test_samples, sample_dir / "test.csv")

        print(f"  Created movie_sample_{idx+1}: train ({len(train_samples)} questions), test ({len(test_samples)} questions) for user {user_id}")


def process_steam():
    """Process steam preference dataset."""
    print("Processing steam dataset...")

    # Load data
    steam_data = pd.read_csv("datasets/steam-200k.csv", header=None,
                              names=['user_id', 'game_name', 'action', 'hours', 'extra'])

    # Filter purchase actions
    purchases = steam_data[steam_data['action'] == 'purchase']

    # Create user -> items mapping (vectorized - much faster!)
    user_items = purchases.groupby('user_id')['game_name'].apply(set).to_dict()

    # Get all games that have at least one purchase (not all games in dataset)
    all_games = set(purchases['game_name'].unique())

    # Create popularity buckets
    print("  Creating popularity buckets...")
    item_to_bucket, item_popularity = create_popularity_buckets(user_items, all_games)

    # Print dataset statistics
    print_dataset_statistics(user_items, all_games, item_popularity, item_to_bucket)

    # Get eligible users
    selected_users = get_users_with_min_items(user_items, min_items=50, num_users=10)

    # Generate samples
    for idx, user_id in enumerate(selected_users):
        train_samples, test_samples = create_train_test_samples(
            user_id,
            user_items[user_id],
            all_games,
            'steam',
            item_to_bucket,
            item_names=None,  # Game names are already the IDs
            num_samples_per_split=25
        )

        sample_dir = output_dir / f"steam_sample_{idx+1}"
        sample_dir.mkdir(parents=True, exist_ok=True)

        save_samples(train_samples, sample_dir / "train.csv")
        save_samples(test_samples, sample_dir / "test.csv")

        print(f"  Created steam_sample_{idx+1}: train ({len(train_samples)} questions), test ({len(test_samples)} questions) for user {user_id}")


if __name__ == "__main__":
    print("Creating preference data samples...")
    print("=" * 50)

    process_anime()
    print()

    process_book()
    print()

    process_movie()
    print()

    process_steam()
    print()

    print("=" * 50)
    print(f"All samples created successfully in {output_dir}")