Merge branch 'main' into shkolar/fix-mac

Author: yuvalshkolar

nemo_retriever/harness/HANDOFF.md

@@ -3,6 +3,7 @@
 active:
   dataset: jp20
   preset: single_gpu
+  run_mode: batch
   query_csv: data/jp20_query_gt.csv
   input_type: pdf
   recall_required: true

nemo_retriever/harness/test_configs.yaml

@@ -10,6 +10,7 @@
 
 import pandas as pd
 
+from nemo_retriever.io.image_store import resolve_image_b64
 from nemo_retriever.ocr.ocr import _crop_b64_image_by_norm_bbox
 from nemo_retriever.params.models import IMAGE_MODALITIES
 
@@ -71,7 +72,7 @@ def explode_content_to_rows(
         batch_df = batch_df.copy()
         if text_mod in IMAGE_MODALITIES and "page_image" in batch_df.columns:
             batch_df["_image_b64"] = batch_df["page_image"].apply(
-                lambda page_image: page_image.get("image_b64") if isinstance(page_image, dict) else None
+                lambda page_image: resolve_image_b64(page_image) if isinstance(page_image, dict) else None
             )
         batch_df["_embed_modality"] = text_mod
         return batch_df
@@ -84,7 +85,7 @@ def explode_content_to_rows(
         page_image = row_dict.get("page_image")
         page_image_b64: Optional[str] = None
         if any_images and isinstance(page_image, dict):
-            page_image_b64 = page_image.get("image_b64")
+            page_image_b64 = resolve_image_b64(page_image)
 
         page_text = row_dict.get(text_column)
         if isinstance(page_text, str) and page_text.strip():
@@ -103,6 +104,7 @@ def explode_content_to_rows(
             for item in content_list:
                 if not isinstance(item, dict):
                     continue
+                item_b64 = resolve_image_b64(item) if struct_mod in IMAGE_MODALITIES else None
                 # Emit rows for text and (optionally) caption fields.
                 for field, content_type in [("text", column), ("caption", f"{column}_caption")]:
                     value = item.get(field, "")
@@ -112,15 +114,18 @@ def explode_content_to_rows(
                     content_row[text_column] = value.strip()
                     content_row["_embed_modality"] = struct_mod
                     content_row["_content_type"] = content_type
-                    if struct_mod in IMAGE_MODALITIES and page_image_b64:
-                        bbox = item.get("bbox_xyxy_norm")
-                        if bbox and len(bbox) == 4:
-                            cropped_b64, _ = _crop_b64_image_by_norm_bbox(page_image_b64, bbox_xyxy_norm=bbox)
-                            content_row["_image_b64"] = cropped_b64
+                    if struct_mod in IMAGE_MODALITIES:
+                        if item_b64:
+                            content_row["_image_b64"] = item_b64
+                        elif page_image_b64:
+                            bbox = item.get("bbox_xyxy_norm")
+                            if bbox and len(bbox) == 4:
+                                cropped_b64, _ = _crop_b64_image_by_norm_bbox(page_image_b64, bbox_xyxy_norm=bbox)
+                                content_row["_image_b64"] = cropped_b64
+                            else:
+                                content_row["_image_b64"] = page_image_b64
                         else:
-                            content_row["_image_b64"] = page_image_b64
-                    elif struct_mod in IMAGE_MODALITIES:
-                        content_row["_image_b64"] = None
+                            content_row["_image_b64"] = None
                     new_rows.append(content_row)
                     exploded_any = True
 
@@ -155,7 +160,7 @@ def collapse_content_to_page_rows(
     if modality in IMAGE_MODALITIES:
         if "page_image" in batch_df.columns:
             batch_df["_image_b64"] = batch_df["page_image"].apply(
-                lambda page_image: page_image.get("image_b64") if isinstance(page_image, dict) else None
+                lambda page_image: resolve_image_b64(page_image) if isinstance(page_image, dict) else None
             )
         else:
             batch_df["_image_b64"] = None

nemo_retriever/src/nemo_retriever/graph/content_transforms.py

@@ -5,34 +5,116 @@
 from __future__ import annotations
 
 import json
+import re
 import subprocess
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any
 
 NEMO_RETRIEVER_ROOT = Path(__file__).resolve().parents[3]
 DEFAULT_ARTIFACTS_ROOT = NEMO_RETRIEVER_ROOT / "artifacts"
+_COMMIT_RE = re.compile(r"^[0-9a-fA-F]{7,40}$")
 
 
 def now_timestr() -> str:
     return datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S_UTC")
 
 
+def _normalize_commit(value: str | None) -> str | None:
+    text = (value or "").strip()
+    if not _COMMIT_RE.match(text):
+        return None
+    return text[:7]
+
+
+def _resolve_git_dir(repo_root: Path) -> Path | None:
+    dot_git = repo_root / ".git"
+    if dot_git.is_dir():
+        return dot_git
+    if not dot_git.is_file():
+        return None
+    try:
+        raw = dot_git.read_text(encoding="utf-8").strip()
+    except Exception:
+        return None
+    if not raw.startswith("gitdir:"):
+        return None
+    gitdir_text = raw.split(":", 1)[1].strip()
+    git_dir = Path(gitdir_text).expanduser()
+    if not git_dir.is_absolute():
+        git_dir = (repo_root / git_dir).resolve()
+    return git_dir
+
+
+def _read_packed_ref(git_dir: Path, ref_name: str) -> str | None:
+    packed_refs = git_dir / "packed-refs"
+    if not packed_refs.exists():
+        return None
+    try:
+        for line in packed_refs.read_text(encoding="utf-8").splitlines():
+            line = line.strip()
+            if not line or line.startswith("#") or line.startswith("^"):
+                continue
+            commit, _sep, ref = line.partition(" ")
+            if ref.strip() == ref_name:
+                normalized = _normalize_commit(commit)
+                if normalized is not None:
+                    return normalized
+    except Exception:
+        return None
+    return None
+
+
+def _read_head_commit(repo_root: Path) -> str | None:
+    git_dir = _resolve_git_dir(repo_root)
+    if git_dir is None:
+        return None
+
+    head_path = git_dir / "HEAD"
+    if not head_path.exists():
+        return None
+    try:
+        head_value = head_path.read_text(encoding="utf-8").strip()
+    except Exception:
+        return None
+
+    if head_value.startswith("ref:"):
+        ref_name = head_value.split(":", 1)[1].strip()
+        ref_path = git_dir / ref_name
+        if ref_path.exists():
+            try:
+                normalized = _normalize_commit(ref_path.read_text(encoding="utf-8"))
+                if normalized is not None:
+                    return normalized
+            except Exception:
+                pass
+        return _read_packed_ref(git_dir, ref_name)
+
+    return _normalize_commit(head_value)
+
+
 def last_commit() -> str:
+    repo_root = NEMO_RETRIEVER_ROOT.parent
     try:
         result = subprocess.run(
             ["git", "rev-parse", "--short", "HEAD"],
-            cwd=str(NEMO_RETRIEVER_ROOT.parent),
+            cwd=str(repo_root),
             check=False,
             capture_output=True,
             text=True,
         )
     except Exception:
-        return "unknown"
+        result = None
+
+    if result is not None and result.returncode == 0:
+        normalized = _normalize_commit(result.stdout)
+        if normalized is not None:
+            return normalized
 
-    if result.returncode != 0:
-        return "unknown"
-    return (result.stdout or "").strip() or "unknown"
+    fallback = _read_head_commit(repo_root)
+    if fallback is not None:
+        return fallback
+    return "unknown"
 
 
 def get_artifacts_root(base_dir: str | None = None) -> Path:

nemo_retriever/src/nemo_retriever/harness/artifacts.py

@@ -15,6 +15,7 @@
 REPO_ROOT = NEMO_RETRIEVER_ROOT.parent
 DEFAULT_TEST_CONFIG_PATH = NEMO_RETRIEVER_ROOT / "harness" / "test_configs.yaml"
 DEFAULT_NIGHTLY_CONFIG_PATH = NEMO_RETRIEVER_ROOT / "harness" / "nightly_config.yaml"
+VALID_RUN_MODES = {"batch", "inprocess"}
 VALID_EVALUATION_MODES = {"recall", "beir"}
 VALID_RECALL_ADAPTERS = {"none", "page_plus_one", "financebench_json"}
 VALID_BEIR_LOADERS = {"vidore_hf"}
@@ -54,6 +55,7 @@ class HarnessConfig:
     dataset_dir: str
     dataset_label: str
     preset: str
+    run_mode: str = "batch"
 
     query_csv: str | None = None
     input_type: str = "pdf"
@@ -114,6 +116,9 @@ def validate(self) -> list[str]:
         if self.query_csv is not None and not Path(self.query_csv).exists():
             errors.append(f"query_csv does not exist: {self.query_csv}")
 
+        if self.run_mode not in VALID_RUN_MODES:
+            errors.append(f"run_mode must be one of {sorted(VALID_RUN_MODES)}")
+
         if self.evaluation_mode not in VALID_EVALUATION_MODES:
             errors.append(f"evaluation_mode must be one of {sorted(VALID_EVALUATION_MODES)}")
 
@@ -263,6 +268,7 @@ def _apply_env_overrides(config_dict: dict[str, Any]) -> None:
         "HARNESS_DATASET": ("dataset", str),
         "HARNESS_DATASET_DIR": ("dataset_dir", str),
         "HARNESS_PRESET": ("preset", str),
+        "HARNESS_RUN_MODE": ("run_mode", str),
         "HARNESS_QUERY_CSV": ("query_csv", str),
         "HARNESS_INPUT_TYPE": ("input_type", str),
         "HARNESS_RECALL_REQUIRED": ("recall_required", _parse_bool),