File size: 8,342 Bytes
c38bb1a d91ce3e c38bb1a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 | #!/usr/bin/env python3
"""
Shared helpers for reading MSMARCO-XI. One place, so every script agrees.
THE ACTUAL PASSAGES LAYOUT (confirmed from the real file)
--------------------------------------------------------
passages: Struct({
'English_passages' : List(String),
'Translated_passages': List(String),
'is_selected' : List(Int64),
})
This is a STRUCT OF PARALLEL LISTS, not a list of structs. One row holds one
dict whose three values are equal-length lists, aligned by position:
row["English_passages"][i] <-> row["Translated_passages"][i]
<-> row["is_selected"][i]
That index `i` is exactly the `passage_index` half of the canonical ID, so the
same passage keeps the same identity in every language file. There is NO url
field, so pseudo-document grouping falls back to query_id.
`iter_passages` normalises both layouts so callers never branch on it.
"""
from __future__ import annotations
import os
from pathlib import Path
from typing import Iterator
# MSMARCO-XI language codes are ISO-639-3 + script, e.g. "asm_Beng", "npi_Deva".
ISO3_TO_ISO2 = {
"asm": "as", "ben": "bn", "guj": "gu", "hin": "hi", "kan": "kn",
"mal": "ml", "mar": "mr",
"nep": "ne", "npi": "ne", # dataset uses npi_Deva for Nepali
"ory": "or", "ori": "or",
"pan": "pa", "san": "sa", "tam": "ta", "tel": "te", "urd": "ur",
"eng": "en",
}
LANG_NAMES = {
"as": "Assamese", "bn": "Bengali", "gu": "Gujarati", "hi": "Hindi",
"kn": "Kannada", "ml": "Malayalam", "mr": "Marathi", "ne": "Nepali",
"or": "Odia", "pa": "Punjabi", "sa": "Sanskrit", "ta": "Tamil",
"te": "Telugu", "ur": "Urdu", "en": "English",
}
def norm_lang(raw) -> str:
"""'asm_Beng' -> 'as'. Unknown 2-letter codes pass through."""
if not raw:
return "en"
code = str(raw).split("_")[0].lower()
return ISO3_TO_ISO2.get(code, code if len(code) == 2 else code)
def iter_passages(row, t_key: str, en_key: str | None,
sel_key: str | None, url_key: str | None
) -> Iterator[tuple[int, str | None, str | None, int, str | None]]:
"""Yield (index, text, text_en, is_selected, url) for either layout.
Layout A - struct of parallel lists (MSMARCO-XI):
{"Translated_passages": [...], "English_passages": [...], "is_selected": [...]}
Layout B - list of structs (the common alternative):
[{"passage_text": ..., "is_selected": ...}, ...]
"""
def at(seq, i):
return seq[i] if isinstance(seq, (list, tuple)) and i < len(seq) else None
if isinstance(row, dict):
texts = row.get(t_key)
if isinstance(texts, (list, tuple)): # layout A
ens = row.get(en_key) if en_key else None
sels = row.get(sel_key) if sel_key else None
urls = row.get(url_key) if url_key else None
for i, txt in enumerate(texts):
yield (i, txt, at(ens, i), int(at(sels, i) or 0), at(urls, i))
else: # single struct
yield (0, texts,
row.get(en_key) if en_key else None,
int(row.get(sel_key) or 0) if sel_key else 0,
row.get(url_key) if url_key else None)
elif isinstance(row, (list, tuple)): # layout B
for i, item in enumerate(row):
if isinstance(item, dict):
yield (i, item.get(t_key),
item.get(en_key) if en_key else None,
int(item.get(sel_key) or 0) if sel_key else 0,
item.get(url_key) if url_key else None)
def detect_layout(row) -> str:
if isinstance(row, dict):
if any(isinstance(v, (list, tuple)) for v in row.values()):
return "struct_of_lists"
return "single_struct"
if isinstance(row, (list, tuple)) and row and isinstance(row[0], dict):
return "list_of_structs"
return "unknown"
def default_root() -> Path:
"""Data root, resolved without relying on the environment.
Order: $VOICERAG_ROOT -> <repo>/../voicerag_data
The second form is derived from this file's location, so it is correct on
jupyter-pod (/workspace/carbine/anurag) and kls-headnode
(/home/dgx-i-carbine/carbine/anurag) alike, and in a fresh shell where
nothing has been exported.
"""
env = os.environ.get("VOICERAG_ROOT")
if env:
return Path(env).expanduser().resolve()
return (Path(__file__).resolve().parents[2] / "voicerag_data").resolve()
def resolve_files(report: dict, root: Path) -> tuple[list[Path], list[str]]:
"""Map file paths recorded in schema_report.json onto the CURRENT data root.
The same NFS export is mounted at different points on the two hosts:
jupyter-pod /workspace/carbine/anurag/voicerag_data
kls-headnode /home/dgx-i-carbine/carbine/anurag/voicerag_data
A report written on one host therefore holds paths that do not exist on the
other. Rather than force a re-run of schema inspection, we rebase: try the
recorded path, then root-relative, then splice the path from its `hf_cache`
segment onto the current root.
Returns (existing_paths, unresolved_strings).
"""
found: list[Path] = []
missing: list[str] = []
for f in report.get("files", []):
p = Path(f)
if p.exists():
found.append(p)
continue
cand = root / f # relative form (written by newer runs)
if cand.exists():
found.append(cand)
continue
parts = p.parts # absolute form from the other host
if "hf_cache" in parts:
cand = root.joinpath(*parts[parts.index("hf_cache"):])
if cand.exists():
found.append(cand)
continue
missing.append(f)
return found, missing
def load_report(root: Path) -> dict:
"""Read schema_report.json and rebase its file list onto `root`."""
import json
rp = root / "data" / "schema_report.json"
if not rp.exists():
raise SystemExit(
f"\nMissing prerequisite: {rp}\n\n"
"Run first: python scripts/02_inspect_schema.py --root " + str(root) + "\n"
)
rep = json.loads(rp.read_text())
files, missing = resolve_files(rep, root)
if not files:
raise SystemExit(
f"\nNone of the {len(rep.get('files', []))} parquet paths in {rp.name} exist "
f"under {root}.\n"
f" first recorded path: {rep.get('files', ['<none>'])[0]}\n\n"
"The report was probably written on the other host. Re-run:\n"
f" python scripts/02_inspect_schema.py --root {root}\n"
)
if missing:
print(f" note: {len(missing)} of {len(rep['files'])} recorded paths "
f"could not be resolved (using {len(files)})")
rep["files"] = [str(f) for f in files]
return rep
def pick_device(explicit: str | None = None) -> str:
"""Which torch device to use. $VOICERAG_DEVICE wins.
It has to win, because on a ZeroGPU Space autodetection is actively wrong.
Free-tier Gradio Spaces are ZeroGPU-only, and OUTSIDE a @spaces.GPU function
ZeroGPU enables a CUDA *emulation* so models can be placed on 'cuda' during
startup. torch.cuda.is_available() therefore returns True while no real GPU
is attached, and autodetect would put bge-m3 on a device that cannot run a
forward pass.
app.py declares one @spaces.GPU function -- ZeroGPU refuses to start a Space
without one -- and never calls it. Quota is billed per call, so the Space
consumes none of the free tier's 5-minute daily allowance and can serve
unlimited requests, instead of going dark when a judge exhausts it.
"""
import os
if explicit:
return explicit
env = os.environ.get("VOICERAG_DEVICE", "").strip().lower()
if env:
return env
import torch
return "cuda" if torch.cuda.is_available() else "cpu"
|