Spaces:

yashb18
/

deepreads

Running

yash-b18

feat: data pipeline — download, preprocess, split, and feature engineering

fef33ce 7 months ago

7.31 kB

"""
Data preprocessing for Amazon Reviews 2023.
Loads raw JSONL files, filters low-activity users and items,
encodes user/item IDs as sequential integers, and saves clean
parquet files to data/processed/.
"""

from __future__ import annotations

import pickle
from pathlib import Path
from typing import Tuple

import pandas as pd
from sklearn.preprocessing import LabelEncoder

from src.config import DataConfig, PROCESSED_DIR
from src.utils.logging import get_logger

logger = get_logger(__name__)


def load_raw_reviews(path: Path) -> pd.DataFrame:
    """
    Load raw reviews JSONL into a DataFrame, keeping only the columns we need.
    Expected columns: user_id, asin (=parent_asin if available), rating, timestamp.
    Args:
        path: Path to raw reviews JSONL.
    Returns:
        DataFrame with columns [user_id, item_id, rating, timestamp].
    """
    logger.info("Loading reviews from %s", path)
    df = pd.read_json(path, lines=True)

    # Use parent_asin when available as the canonical item identifier
    if "parent_asin" in df.columns:
        df["item_id"] = df["parent_asin"].fillna(df["asin"])
    else:
        df["item_id"] = df["asin"]

    keep = ["user_id", "item_id", "rating", "timestamp"]
    df = df[keep].copy()

    # Normalise types
    df["rating"] = pd.to_numeric(df["rating"], errors="coerce")
    df["timestamp"] = pd.to_numeric(df["timestamp"], errors="coerce")

    # Drop rows with missing essentials
    df = df.dropna(subset=["user_id", "item_id", "rating"])
    df = df.drop_duplicates(subset=["user_id", "item_id"])
    logger.info("Loaded %d reviews (%d users, %d items)",
                len(df), df["user_id"].nunique(), df["item_id"].nunique())
    return df


def load_raw_metadata(path: Path) -> pd.DataFrame:
    """
    Load raw metadata JSONL into a DataFrame.
    Args:
        path: Path to raw metadata JSONL.
    Returns:
        DataFrame with item-level attributes.
    """
    logger.info("Loading metadata from %s", path)
    df = pd.read_json(path, lines=True)

    if "parent_asin" in df.columns:
        df["item_id"] = df["parent_asin"].fillna(df.get("asin", df["parent_asin"]))
    elif "asin" in df.columns:
        df["item_id"] = df["asin"]
    else:
        raise ValueError("Metadata JSONL has neither 'asin' nor 'parent_asin' column.")

    # Flatten categories: take first listed category as primary
    if "categories" in df.columns:
        df["category"] = df["categories"].apply(
            lambda x: x[0] if isinstance(x, list) and len(x) > 0 else "Unknown"
        )
    else:
        df["category"] = "Unknown"

    # Normalise price to float
    if "price" in df.columns:
        df["price"] = pd.to_numeric(df["price"], errors="coerce")
    else:
        df["price"] = float("nan")

    keep_cols = ["item_id", "title", "category", "price"]
    # Include brand/store if present
    for col in ("store", "brand"):
        if col in df.columns:
            df["brand"] = df[col].fillna("Unknown")
            break
    else:
        df["brand"] = "Unknown"
    keep_cols.append("brand")

    # Include description for text features
    for col in ("description", "features"):
        if col in df.columns:
            df["description"] = df[col].apply(
                lambda x: " ".join(x) if isinstance(x, list) else str(x)
            )
            break
    else:
        df["description"] = ""
    keep_cols.append("description")

    df = df[[c for c in keep_cols if c in df.columns]].drop_duplicates("item_id")
    logger.info("Loaded metadata for %d items", len(df))
    return df


def filter_interactions(
    df: pd.DataFrame,
    min_user_reviews: int,
    min_item_reviews: int,
) -> pd.DataFrame:
    """
    Iteratively filter users and items below minimum interaction counts.
    Args:
        df: Interactions DataFrame.
        min_user_reviews: Minimum reviews per user.
        min_item_reviews: Minimum reviews per item.
    Returns:
        Filtered DataFrame.
    """
    prev_size = len(df) + 1
    while len(df) < prev_size:
        prev_size = len(df)
        user_counts = df["user_id"].value_counts()
        df = df[df["user_id"].isin(user_counts[user_counts >= min_user_reviews].index)]
        item_counts = df["item_id"].value_counts()
        df = df[df["item_id"].isin(item_counts[item_counts >= min_item_reviews].index)]

    logger.info("After filtering: %d reviews, %d users, %d items",
                len(df), df["user_id"].nunique(), df["item_id"].nunique())
    return df


def encode_ids(
    df: pd.DataFrame,
    user_encoder_path: Path,
    item_encoder_path: Path,
) -> Tuple[pd.DataFrame, LabelEncoder, LabelEncoder]:
    """
    Encode user_id and item_id as zero-based integers and save encoders.
    Args:
        df: Interactions DataFrame with string user_id and item_id.
        user_encoder_path: Where to save the user LabelEncoder.
        item_encoder_path: Where to save the item LabelEncoder.
    Returns:
        Tuple of (encoded_df, user_encoder, item_encoder).
    """
    PROCESSED_DIR.mkdir(parents=True, exist_ok=True)

    user_enc = LabelEncoder()
    item_enc = LabelEncoder()

    df = df.copy()
    df["user_idx"] = user_enc.fit_transform(df["user_id"])
    df["item_idx"] = item_enc.fit_transform(df["item_id"])

    with open(user_encoder_path, "wb") as f:
        pickle.dump(user_enc, f)
    with open(item_encoder_path, "wb") as f:
        pickle.dump(item_enc, f)

    logger.info("Encoded %d users and %d items", len(user_enc.classes_), len(item_enc.classes_))
    return df, user_enc, item_enc


def preprocess(cfg: DataConfig) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Full preprocessing pipeline: load → filter → encode → save.
    Args:
        cfg: DataConfig.
    Returns:
        Tuple of (interactions_df, metadata_df).
    """
    interactions = load_raw_reviews(cfg.raw_reviews_path)
    metadata = load_raw_metadata(cfg.raw_metadata_path)

    # Optional debug subsetting before filtering
    if cfg.debug:
        top_users = interactions["user_id"].value_counts().head(cfg.debug_n_users).index
        interactions = interactions[interactions["user_id"].isin(top_users)].copy()
        logger.info("Debug mode: trimmed to %d users", interactions["user_id"].nunique())

    interactions = filter_interactions(
        interactions, cfg.min_user_reviews, cfg.min_item_reviews
    )

    interactions, user_enc, item_enc = encode_ids(
        interactions, cfg.user_encoder_path, cfg.item_encoder_path
    )

    # Keep only metadata for items present after filtering
    valid_items = set(interactions["item_id"].unique())
    metadata = metadata[metadata["item_id"].isin(valid_items)].copy()

    # Add item_idx to metadata
    item_id_to_idx = dict(zip(user_enc.classes_, range(len(user_enc.classes_))))
    item_id_to_idx = dict(zip(item_enc.classes_, range(len(item_enc.classes_))))
    metadata["item_idx"] = metadata["item_id"].map(item_id_to_idx)

    # Save
    PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
    interactions.to_parquet(cfg.interactions_path, index=False)
    metadata.to_parquet(cfg.metadata_path, index=False)
    logger.info("Saved interactions → %s", cfg.interactions_path)
    logger.info("Saved metadata     → %s", cfg.metadata_path)

    return interactions, metadata

	"""
	Data preprocessing for Amazon Reviews 2023.

	Loads raw JSONL files, filters low-activity users and items,
	encodes user/item IDs as sequential integers, and saves clean
	parquet files to data/processed/.
	"""

	from __future__ import annotations

	import pickle
	from pathlib import Path
	from typing import Tuple

	import pandas as pd
	from sklearn.preprocessing import LabelEncoder

	from src.config import DataConfig, PROCESSED_DIR
	from src.utils.logging import get_logger

	logger = get_logger(__name__)


	def load_raw_reviews(path: Path) -> pd.DataFrame:
	"""
	Load raw reviews JSONL into a DataFrame, keeping only the columns we need.

	Expected columns: user_id, asin (=parent_asin if available), rating, timestamp.

	Args:
	path: Path to raw reviews JSONL.

	Returns:
	DataFrame with columns [user_id, item_id, rating, timestamp].
	"""
	logger.info("Loading reviews from %s", path)
	df = pd.read_json(path, lines=True)

	# Use parent_asin when available as the canonical item identifier
	if "parent_asin" in df.columns:
	df["item_id"] = df["parent_asin"].fillna(df["asin"])
	else:
	df["item_id"] = df["asin"]

	keep = ["user_id", "item_id", "rating", "timestamp"]
	df = df[keep].copy()

	# Normalise types
	df["rating"] = pd.to_numeric(df["rating"], errors="coerce")
	df["timestamp"] = pd.to_numeric(df["timestamp"], errors="coerce")

	# Drop rows with missing essentials
	df = df.dropna(subset=["user_id", "item_id", "rating"])
	df = df.drop_duplicates(subset=["user_id", "item_id"])
	logger.info("Loaded %d reviews (%d users, %d items)",
	len(df), df["user_id"].nunique(), df["item_id"].nunique())
	return df


	def load_raw_metadata(path: Path) -> pd.DataFrame:
	"""
	Load raw metadata JSONL into a DataFrame.

	Args:
	path: Path to raw metadata JSONL.

	Returns:
	DataFrame with item-level attributes.
	"""
	logger.info("Loading metadata from %s", path)
	df = pd.read_json(path, lines=True)

	if "parent_asin" in df.columns:
	df["item_id"] = df["parent_asin"].fillna(df.get("asin", df["parent_asin"]))
	elif "asin" in df.columns:
	df["item_id"] = df["asin"]
	else:
	raise ValueError("Metadata JSONL has neither 'asin' nor 'parent_asin' column.")

	# Flatten categories: take first listed category as primary
	if "categories" in df.columns:
	df["category"] = df["categories"].apply(
	lambda x: x[0] if isinstance(x, list) and len(x) > 0 else "Unknown"
	)
	else:
	df["category"] = "Unknown"

	# Normalise price to float
	if "price" in df.columns:
	df["price"] = pd.to_numeric(df["price"], errors="coerce")
	else:
	df["price"] = float("nan")

	keep_cols = ["item_id", "title", "category", "price"]
	# Include brand/store if present
	for col in ("store", "brand"):
	if col in df.columns:
	df["brand"] = df[col].fillna("Unknown")
	break
	else:
	df["brand"] = "Unknown"
	keep_cols.append("brand")

	# Include description for text features
	for col in ("description", "features"):
	if col in df.columns:
	df["description"] = df[col].apply(
	lambda x: " ".join(x) if isinstance(x, list) else str(x)
	)
	break
	else:
	df["description"] = ""
	keep_cols.append("description")

	df = df[[c for c in keep_cols if c in df.columns]].drop_duplicates("item_id")
	logger.info("Loaded metadata for %d items", len(df))
	return df


	def filter_interactions(
	df: pd.DataFrame,
	min_user_reviews: int,
	min_item_reviews: int,
	) -> pd.DataFrame:
	"""
	Iteratively filter users and items below minimum interaction counts.

	Args:
	df: Interactions DataFrame.
	min_user_reviews: Minimum reviews per user.
	min_item_reviews: Minimum reviews per item.

	Returns:
	Filtered DataFrame.
	"""
	prev_size = len(df) + 1
	while len(df) < prev_size:
	prev_size = len(df)
	user_counts = df["user_id"].value_counts()
	df = df[df["user_id"].isin(user_counts[user_counts >= min_user_reviews].index)]
	item_counts = df["item_id"].value_counts()
	df = df[df["item_id"].isin(item_counts[item_counts >= min_item_reviews].index)]

	logger.info("After filtering: %d reviews, %d users, %d items",
	len(df), df["user_id"].nunique(), df["item_id"].nunique())
	return df


	def encode_ids(
	df: pd.DataFrame,
	user_encoder_path: Path,
	item_encoder_path: Path,
	) -> Tuple[pd.DataFrame, LabelEncoder, LabelEncoder]:
	"""
	Encode user_id and item_id as zero-based integers and save encoders.

	Args:
	df: Interactions DataFrame with string user_id and item_id.
	user_encoder_path: Where to save the user LabelEncoder.
	item_encoder_path: Where to save the item LabelEncoder.

	Returns:
	Tuple of (encoded_df, user_encoder, item_encoder).
	"""
	PROCESSED_DIR.mkdir(parents=True, exist_ok=True)

	user_enc = LabelEncoder()
	item_enc = LabelEncoder()

	df = df.copy()
	df["user_idx"] = user_enc.fit_transform(df["user_id"])
	df["item_idx"] = item_enc.fit_transform(df["item_id"])

	with open(user_encoder_path, "wb") as f:
	pickle.dump(user_enc, f)
	with open(item_encoder_path, "wb") as f:
	pickle.dump(item_enc, f)

	logger.info("Encoded %d users and %d items", len(user_enc.classes_), len(item_enc.classes_))
	return df, user_enc, item_enc


	def preprocess(cfg: DataConfig) -> Tuple[pd.DataFrame, pd.DataFrame]:
	"""
	Full preprocessing pipeline: load → filter → encode → save.

	Args:
	cfg: DataConfig.

	Returns:
	Tuple of (interactions_df, metadata_df).
	"""
	interactions = load_raw_reviews(cfg.raw_reviews_path)
	metadata = load_raw_metadata(cfg.raw_metadata_path)

	# Optional debug subsetting before filtering
	if cfg.debug:
	top_users = interactions["user_id"].value_counts().head(cfg.debug_n_users).index
	interactions = interactions[interactions["user_id"].isin(top_users)].copy()
	logger.info("Debug mode: trimmed to %d users", interactions["user_id"].nunique())

	interactions = filter_interactions(
	interactions, cfg.min_user_reviews, cfg.min_item_reviews
	)

	interactions, user_enc, item_enc = encode_ids(
	interactions, cfg.user_encoder_path, cfg.item_encoder_path
	)

	# Keep only metadata for items present after filtering
	valid_items = set(interactions["item_id"].unique())
	metadata = metadata[metadata["item_id"].isin(valid_items)].copy()

	# Add item_idx to metadata
	item_id_to_idx = dict(zip(user_enc.classes_, range(len(user_enc.classes_))))
	item_id_to_idx = dict(zip(item_enc.classes_, range(len(item_enc.classes_))))
	metadata["item_idx"] = metadata["item_id"].map(item_id_to_idx)

	# Save
	PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
	interactions.to_parquet(cfg.interactions_path, index=False)
	metadata.to_parquet(cfg.metadata_path, index=False)
	logger.info("Saved interactions → %s", cfg.interactions_path)
	logger.info("Saved metadata → %s", cfg.metadata_path)

	return interactions, metadata