| """Shared constants and helpers for the Decoding Parliament RAG pipeline.""" |
| import json, os, re |
|
|
| DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data") |
| CHUNKS = os.path.join(DATA, "chunks.jsonl") |
| INDEX_DIR = os.path.join(DATA, "index") |
|
|
| EMBED_MODEL = "BAAI/bge-m3" |
| MAX_SEQ_LEN = 1024 |
| ITEM_TITLE_CAP = 200 |
|
|
| HOUSE_NAMES = {"lower": "Pyithu Hluttaw", "upper": "Amyotha Hluttaw", "union": "Pyidaungsu Hluttaw"} |
| HOUSE_NAMES_MY = {"lower": "ပြည်သူ့လွှတ်တော်", "upper": "အမျိုးသားလွှတ်တော်", "union": "ပြည်ထောင်စုလွှတ်တော်"} |
|
|
| |
| THEME_KEYWORDS = ["ကျန်းမာရေး", "ပညာရေး", "ဆေးရုံ", "ကျောင်း", "ဆရာ", "ဆေး", "လူမှုဖူလုံရေး", "ဆရာဝန်", |
| "ကျောင်းသား", "တက္ကသိုလ်", "ကလေး", "မိခင်", "ပညာသင်", "သူနာပြု"] |
|
|
| def cap_title(t): |
| t = t or "" |
| return t if len(t) <= ITEM_TITLE_CAP else t[:ITEM_TITLE_CAP] + "…" |
|
|
| def header_of(r): |
| """Citation/embedding header: house | date | section type | section title | item title | speaker.""" |
| return " | ".join(x for x in [HOUSE_NAMES.get(r["house"], r["house"]), r["date"], r["section_type"], |
| cap_title(r.get("section_title")), cap_title(r.get("item_title")), |
| r.get("speaker_name") or r.get("role")] if x) |
|
|
| def embed_text_of(r): |
| return header_of(r) + "\n" + r["text"] |
|
|
| def is_themed(line_or_rec): |
| s = line_or_rec if isinstance(line_or_rec, str) else json.dumps(line_or_rec, ensure_ascii=False) |
| return any(k in s for k in THEME_KEYWORDS) |
|
|
| PARTY_SHORT = {"ပြည်ထောင်စုကြံ့ခိုင်ရေးနှင့်ဖွံ့ဖြိုးရေးပါတီ": "USDP", "အမျိုးသားဒီမိုကရေစီအဖွဲ့ချုပ်": "NLD", |
| "တိုင်းရင်းသားစည်းလုံးညီညွတ်ရေးပါတီ": "NUP", "ရှမ်းတိုင်းရင်းသားများဒီမိုကရက်တစ်ပါတီ": "SNDP", |
| "ရခိုင်တိုင်းရင်းသားများတိုးတက်ရေးပါတီ": "RNDP", "ရခိုင်အမျိုးသားပါတီ": "ANP", "အမျိုးသားဒီမိုကရေစီအင်အားစုပါတီ": "NDF", |
| "တစ်သီးပုဂ္ဂလ": "Independent"} |
| _MP = None |
| def party_of(r): |
| """Party of the speaking MP in the term of the sitting, from the member register (public political affiliation). |
| Returns e.g. 'အမျိုးသားဒီမိုကရေစီအဖွဲ့ချုပ် (NLD)' or '' when unknown / not an MP.""" |
| global _MP |
| if r.get("speaker_kind") != "mp" or not r.get("speaker_id"): return "" |
| if _MP is None: |
| p = os.path.join(DATA, "mp_lookup.json") |
| _MP = json.load(open(p, encoding="utf-8")) if os.path.exists(p) else {} |
| m = _MP.get(str(r["speaker_id"])) or {} |
| party = m.get("party_t2" if r.get("term") == 2 else "party_t1") or "" |
| if not party: return "" |
| short = PARTY_SHORT.get(party) |
| return f"{party} ({short})" if short else party |
|
|
| def citation_of(r): |
| """Human-readable citation for answers: house, date, section, speaker (+constituency, party if known).""" |
| who = r.get("speaker_name") or {"speaker": "Speaker", "vicespeaker": "Deputy Speaker", "mc": "MC"}.get(r.get("role"), r.get("role")) |
| party = party_of(r) |
| if who and party: who = f"{who} · {party}" |
| title = cap_title(r.get('item_title') or r.get('section_title')) |
| parts = [f"{HOUSE_NAMES.get(r['house'], r['house'])}, {r['date']}, {r['section_type']}"] |
| if title: parts.append(title) |
| if who: parts.append(who) |
| return " — ".join(parts) |
|
|
| def meta_of(r): |
| """Slim metadata kept beside each embedding (everything needed to display and cite a hit).""" |
| return {k: r.get(k) for k in ("chunk_id", "speech_id", "doc_id", "house", "date", "term", "session", "sitting_day", |
| "section_type", "section_title", "item_title", "role", "speaker_kind", "speaker_id", |
| "speaker_name", "part", "n_parts", "text")} |
|
|
| def load_embedder(device=None, fp16=True, model="bge"): |
| """Load an embedding model from the registry in embed_models.py (default: BGE-M3). |
| Returns (SentenceTransformer, device, spec) — spec carries the query/passage prefixes and index dir.""" |
| import torch |
| from sentence_transformers import SentenceTransformer |
| from embed_models import spec |
| sp = spec(model) |
| if device is None: |
| device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu") |
| m = SentenceTransformer(sp["hf"], device=device, trust_remote_code=True) |
| m.max_seq_length = sp["max_seq"] |
| if fp16 and device != "cpu": m.half() |
| return m, device, sp |
|
|
