| from __future__ import annotations |
|
|
| import json |
| import logging |
| import math |
| import re |
| from collections import defaultdict |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def extract_user_id(persona: str) -> str | None: |
| for pat in ( |
| r"(?im)^user_id:\s*(\S+)", |
| r"(?im)^User ID:\s*(\S+)", |
| r"(?i)\buser_id\s*=\s*(\S+)", |
| ): |
| m = re.search(pat, persona) |
| if m: |
| return m.group(1).strip() |
| return None |
|
|
|
|
| class TaskAReviewRagIndex: |
| def __init__(self, path: Path) -> None: |
| self.path = path |
| self._rows: list[dict[str, Any]] = [] |
| self._mat: np.ndarray | None = None |
| self._by_user: dict[str, list[int]] = defaultdict(list) |
| self._loaded = False |
|
|
| @property |
| def loaded(self) -> bool: |
| return self._loaded |
|
|
| def load(self) -> None: |
| if self._loaded: |
| return |
| if not self.path.is_file(): |
| logger.warning("Task A RAG file missing — run scripts/build_task_a_review_rag.py or Docker build (%s)", self.path) |
| self._rows = [] |
| self._mat = None |
| self._loaded = True |
| return |
|
|
| embeddings: list[list[float]] = [] |
| with self.path.open(encoding="utf-8", errors="replace") as f: |
| for line in f: |
| line = line.strip() |
| if not line: |
| continue |
| row = json.loads(line) |
| emb = row.get("embedding") |
| if not emb: |
| continue |
| uid = row.get("user_id") or "" |
| clean = {k: v for k, v in row.items() if k != "embedding"} |
| idx = len(self._rows) |
| self._rows.append(clean) |
| embeddings.append(emb) |
| if uid: |
| self._by_user[str(uid)].append(idx) |
|
|
| if not embeddings: |
| logger.warning("Task A RAG file empty or invalid: %s", self.path) |
| self._mat = None |
| self._loaded = True |
| return |
|
|
| self._mat = np.asarray(embeddings, dtype=np.float32) |
| norms = np.linalg.norm(self._mat, axis=1, keepdims=True) |
| norms[norms == 0] = 1.0 |
| self._mat = self._mat / norms |
| self._loaded = True |
| logger.info("Task A RAG loaded %d snippets from %s", len(self._rows), self.path) |
|
|
| def _encode_query(self, embedder: Any, text: str) -> np.ndarray: |
| t = text.replace("\n", " ")[:8000] |
| vec = embedder.encode([t], convert_to_numpy=True, normalize_embeddings=False)[0] |
| q = np.asarray(vec, dtype=np.float32) |
| nq = np.linalg.norm(q) |
| if nq == 0: |
| q = np.ones_like(q) / math.sqrt(len(q)) |
| else: |
| q = q / nq |
| return q |
|
|
| def retrieve( |
| self, |
| persona: str, |
| product: str, |
| embedder: Any, |
| top_k: int, |
| user_examples_cap: int = 3, |
| ) -> list[dict[str, Any]]: |
| self.load() |
| if self._mat is None or not self._rows: |
| return [] |
|
|
| k = max(1, top_k) |
| uid = extract_user_id(persona) |
| qtext = f"{persona.strip()}\n{product.strip()}" |
| q = self._encode_query(embedder, qtext) |
| scores = self._mat @ q |
| order = np.argsort(-scores) |
|
|
| picked_idx: list[int] = [] |
| seen: set[int] = set() |
|
|
| if uid and uid in self._by_user: |
| uidx = self._by_user[uid] |
| ranked_u = sorted(uidx, key=lambda i: float(scores[i]), reverse=True) |
| for i in ranked_u[: min(user_examples_cap, k)]: |
| picked_idx.append(i) |
| seen.add(i) |
|
|
| for i in order.tolist(): |
| if len(picked_idx) >= k: |
| break |
| if i in seen: |
| continue |
| picked_idx.append(i) |
| seen.add(i) |
|
|
| out: list[dict[str, Any]] = [] |
| for i in picked_idx[:k]: |
| r = self._rows[i] |
| out.append( |
| { |
| "stars": r.get("stars"), |
| "review_excerpt": r.get("review_excerpt", ""), |
| "business_context": r.get("business_context", ""), |
| "same_user_as_persona": bool(uid and r.get("user_id") == uid), |
| } |
| ) |
| return out |
|
|
