scripts/build_task_a_review_rag.py

#!/usr/bin/env python3

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Any


def repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def compact_attributes(attrs: Any, max_len: int = 400) -> str:
    if not attrs:
        return ""
    try:
        s = json.dumps(attrs, ensure_ascii=False, sort_keys=True)
    except (TypeError, ValueError):
        s = str(attrs)
    if len(s) > max_len:
        return s[: max_len - 1] + "…"
    return s


def business_context_line(b: dict[str, Any]) -> str:
    cats = b.get("categories") or ""
    parts = [
        f"name: {b.get('name', '')}",
        f"categories: {cats}",
        f"location: {b.get('city', '')}, {b.get('state', '')}",
        f"business_avg_stars: {b.get('stars', '')}",
        f"business_review_count: {b.get('review_count', '')}",
    ]
    attr = compact_attributes(b.get("attributes"))
    if attr:
        parts.append(f"attributes_json: {attr}")
    return "\n".join(parts)


def load_business_map(business_json: Path) -> dict[str, dict[str, Any]]:
    m: dict[str, dict[str, Any]] = {}
    with business_json.open(encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            b = json.loads(line)
            bid = b.get("business_id")
            if bid:
                m[str(bid)] = b
    return m


def review_text_for_embedding(review_excerpt: str, business_ctx: str) -> str:
    return f"{business_ctx}\nreview: {review_excerpt}"


def main() -> None:
    root = repo_root()
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--review-json",
        type=Path,
        default=root / "yelp_dataset" / "extracted" / "yelp_academic_dataset_review.json",
    )
    parser.add_argument(
        "--business-json",
        type=Path,
        default=root / "yelp_dataset" / "extracted" / "yelp_academic_dataset_business.json",
    )
    parser.add_argument("--output", type=Path, default=root / "data" / "task_a_reviews_embedded.jsonl")
    parser.add_argument("--max-rows", type=int, default=12_000)
    parser.add_argument("--review-text-chars", type=int, default=480)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument(
        "--model",
        type=str,
        default=os.environ.get("TASK_B_LOCAL_EMBEDDING_MODEL", "all-MiniLM-L6-v2"),
    )
    args = parser.parse_args()

    if not args.review_json.is_file():
        raise SystemExit(f"Missing review JSON: {args.review_json}")

    biz: dict[str, dict[str, Any]] = {}
    if args.business_json.is_file():
        biz = load_business_map(args.business_json)
    else:
        print("build_task_a_review_rag: business JSON not found — using review text only for embedding.")

    rows_raw: list[dict[str, Any]] = []
    with args.review_json.open(encoding="utf-8", errors="replace") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            r = json.loads(line)
            text = (r.get("text") or "").strip()
            if not text:
                continue
            bid = str(r.get("business_id") or "")
            uid = str(r.get("user_id") or "")
            stars = r.get("stars")
            try:
                stars_i = int(stars) if stars is not None else None
            except (TypeError, ValueError):
                stars_i = None
            if stars_i is not None:
                stars_i = max(1, min(5, stars_i))
            excerpt = text[: args.review_text_chars]
            bctx = ""
            if bid and bid in biz:
                bctx = business_context_line(biz[bid])
            emb_src = review_text_for_embedding(excerpt, bctx) if bctx else excerpt
            rows_raw.append(
                {
                    "user_id": uid,
                    "business_id": bid,
                    "stars": stars_i,
                    "review_excerpt": excerpt,
                    "business_context": bctx,
                    "embedding_source": emb_src,
                }
            )
            if len(rows_raw) >= args.max_rows:
                break

    if not rows_raw:
        raise SystemExit("No reviews ingested — check review JSON format.")

    try:
        from sentence_transformers import SentenceTransformer  # type: ignore[import-untyped]
    except ImportError as e:
        raise SystemExit("pip install sentence-transformers") from e

    texts = [r["embedding_source"] for r in rows_raw]
    model = SentenceTransformer(args.model)
    mat = model.encode(
        texts,
        batch_size=args.batch_size,
        convert_to_numpy=True,
        normalize_embeddings=False,
        show_progress_bar=len(texts) > args.batch_size,
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as fout:
        for rec, vec in zip(rows_raw, mat, strict=True):
            row_out = {k: v for k, v in rec.items() if k != "embedding_source"}
            row_out["embedding"] = vec.astype(float).tolist()
            fout.write(json.dumps(row_out, ensure_ascii=False) + "\n")

    print(f"Wrote {len(rows_raw)} embedded reviews -> {args.output}")


if __name__ == "__main__":
    main()