File size: 5,517 Bytes
099b63e 935f6d4 099b63e 935f6d4 099b63e 935f6d4 099b63e | 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 | """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") # embedding shards + metadata live here
EMBED_MODEL = "BAAI/bge-m3" # baseline; other candidates in embed_models.py
MAX_SEQ_LEN = 1024 # Burmese ≈ 3.3 chars/token; 1,800-char chunks fit, 512 would truncate 74%
ITEM_TITLE_CAP = 200 # some QA <title>s hold the whole question text
HOUSE_NAMES = {"lower": "Pyithu Hluttaw", "upper": "Amyotha Hluttaw", "union": "Pyidaungsu Hluttaw"}
HOUSE_NAMES_MY = {"lower": "ပြည်သူ့လွှတ်တော်", "upper": "အမျိုးသားလွှတ်တော်", "union": "ပြည်ထောင်စုလွှတ်တော်"}
# health / education / social-welfare screen used to order the embedding run (themed chunks first)
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
|