jia_embedding_e5.py

#!/usr/bin/env python3
"""Generate multilingual-e5 embeddings for content_cleaned and parent_url_cleaned columns and save them to *_e52.csv files."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Dict, List

import pandas as pd
from mlx_embeddings.utils import load


BASE_DIR = Path("/Users/jia/Library/CloudStorage/OneDrive-Personal/Projects/Web Archives for Social Sciences")
DEFAULT_INPUTS = [
    BASE_DIR / "problem3_2021_combined_uurl.csv",
    BASE_DIR / "problem3_2024_combined_uurl.csv",
]
DEFAULT_MODEL_PATH = BASE_DIR / "models/multilingual-e5-large-mlx"
MAX_SEQ_LEN = 512
OUTPUT_SUFFIX = "_e53"
EMBED_COLUMNS = {
    "summary": "embedding_summary",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--inputs",
        type=Path,
        nargs="+",
        default=DEFAULT_INPUTS,
        help="One or more source CSV files.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Directory for output CSVs. Defaults to each input file's directory.",
    )
    parser.add_argument("--model-path", type=Path, default=DEFAULT_MODEL_PATH, help="Path to the MLX checkpoint.")
    parser.add_argument(
        "--flush-interval",
        type=int,
        default=500,
        help="Write out progress every N newly computed embeddings.",
    )
    return parser.parse_args()


def embed_text(model, tokenizer, text: str) -> List[float]:
    encoded = tokenizer.encode(
        text,
        return_tensors="mlx",
        truncation=True,
        max_length=MAX_SEQ_LEN,
    )
    outputs = model(encoded)
    emb = outputs.text_embeds.tolist()
    return emb[0]


def load_existing(output_path: Path) -> Dict[str, Dict[str, str]]:
    """Load existing embeddings by uid."""
    if not output_path.exists():
        print("No existing embedding file found. Starting fresh.")
        return {}

    df = pd.read_csv(output_path)
    required_cols = {"uid", *EMBED_COLUMNS.values()}
    if not required_cols.issubset(df.columns):
        print(f"Warning: {output_path} is missing required embedding columns.")
        return {}

    processed: Dict[str, Dict[str, str]] = {}
    for row in df.itertuples(index=False):
        row_dict = row._asdict()
        uid = str(row.uid)

        embeddings: Dict[str, str] = {}
        for source_col, emb_col in EMBED_COLUMNS.items():
            value = row_dict.get(emb_col)
            if pd.isna(value):
                embeddings = {}
                break
            embeddings[source_col] = value

        if embeddings:
            processed[uid] = embeddings

    print(f"Loaded {len(processed)} existing embeddings from {output_path}")
    return processed


def write_progress(output_path: Path, df: pd.DataFrame, embeddings_dict: Dict[str, Dict[str, str]]) -> None:
    """Write progress by merging all original columns with embedding columns."""
    if not embeddings_dict:
        return

    # Create a copy of the dataframe
    df_output = df.copy()

    # Add embedding columns
    for source_col, emb_col in EMBED_COLUMNS.items():
        column_embeddings = {uid: emb_dict.get(source_col) for uid, emb_dict in embeddings_dict.items()}
        df_output[emb_col] = df_output["uid"].map(column_embeddings)

    # Only keep rows that have embeddings for all requested columns
    df_output = df_output.dropna(subset=list(EMBED_COLUMNS.values()))

    output_path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = output_path.with_suffix(output_path.suffix + ".tmp")
    df_output.to_csv(tmp_path, index=False)
    tmp_path.replace(output_path)
    print(f"Wrote {len(df_output)} rows with embeddings to {output_path}")


def process_file(
    input_path: Path,
    output_path: Path,
    model,
    tokenizer,
    flush_interval: int,
) -> None:
    print(f"\n{'='*80}")
    print(f"Processing {input_path.name}")
    print(f"{'='*80}")

    df = pd.read_csv(input_path)

    # Check for required columns
    missing_cols = [col for col in ["uid", *EMBED_COLUMNS.keys()] if col not in df.columns]
    if missing_cols:
        raise ValueError(f"Input missing required columns: {', '.join(missing_cols)}")

    total_rows = len(df)
    unique_ids = df["uid"].nunique()
    dup_rows = total_rows - unique_ids
    print(f"Input rows: {total_rows}. Unique uids: {unique_ids}. Duplicate rows: {dup_rows}.")
    print(f"Will embed columns: {', '.join(EMBED_COLUMNS.keys())}")

    # Load existing embeddings
    processed = load_existing(output_path)

    # Find rows that need embedding
    seen_ids = set(processed.keys())
    pending_rows = []
    for row in df.itertuples(index=False):
        row_uid = str(row.uid)
        if row_uid not in seen_ids:
            pending_rows.append(row)

    if not pending_rows:
        print("All rows already embedded. Nothing to do.")
        # Still write the final file with all columns
        write_progress(output_path, df, processed)
        return

    print(f"Need to compute embeddings for {len(pending_rows)} rows")

    since_flush = 0
    total_new = 0

    for row in pending_rows:
        row_dict = row._asdict()
        row_uid = str(row.uid)

        row_embeddings: Dict[str, str] = {}
        for source_col in EMBED_COLUMNS:
            text_value = "" if pd.isna(row_dict.get(source_col)) else str(row_dict[source_col]).strip()
            if text_value:
                vector = embed_text(model, tokenizer, text_value)
            else:
                vector = []
            row_embeddings[source_col] = json.dumps(vector)

        processed[row_uid] = row_embeddings

        total_new += 1
        since_flush += 1

        if since_flush >= flush_interval:
            write_progress(output_path, df, processed)
            since_flush = 0
            print(f"Progress: {total_new}/{len(pending_rows)} embeddings computed")

    # Final write
    if since_flush > 0 or total_new > 0:
        write_progress(output_path, df, processed)

    print(f"Computed {total_new} new embeddings. Total stored: {len(processed)}")
    print(f"Output saved to: {output_path}")


def main() -> None:
    args = parse_args()

    # Load model once for all files
    print(f"Loading model from {args.model_path}...")
    model, tokenizer = load(str(args.model_path))
    print("Model loaded successfully.")

    for input_path in args.inputs:
        if not input_path.exists():
            print(f"Warning: Input file not found: {input_path}. Skipping.")
            continue

        # Determine output path
        if args.output_dir:
            output_dir = args.output_dir.expanduser().resolve()
        else:
            output_dir = input_path.parent

        # Change output filename pattern
        # problem3_2021_combined_uurl.csv -> problem3_2021_combined_uurl_e52.csv
        output_filename = f"{input_path.stem}{OUTPUT_SUFFIX}{input_path.suffix}"
        output_path = output_dir / output_filename

        process_file(
            input_path=input_path,
            output_path=output_path,
            model=model,
            tokenizer=tokenizer,
            flush_interval=args.flush_interval,
        )


if __name__ == "__main__":
    main()