| """Grounded question answering: retrieve passages, then let Gemini answer ONLY from them, with citations. |
| |
| Setup: put the key for the provider you use in FinalProject.env (git-ignored): |
| GEMINI_API_KEY=... pip install google-genai --provider gemini (default) |
| ANTHROPIC_API_KEY=... pip install anthropic --provider anthropic |
| OPENAI_API_KEY=... pip install openai --provider openai |
| |
| python rag.py "ကျေးလက်ဒေသ ဆေးရုံများတွင် ဆရာဝန် မလုံလောက်မှုကို အစိုးရက ဘယ်လိုဖြေရှင်းမလဲ" |
| python rag.py --provider anthropic --batch data/pilot_questions.jsonl --out data/pilot_answers.jsonl |
| python rag.py --no-retrieval "..." # ablation: same model, no passages (for the faithfulness comparison) |
| python rag.py --embed sealion --mode hybrid --per-speech 2 --batch data/eval_questions.jsonl --out data/eval_answers.jsonl |
| python rag.py --provider anthropic --list-models / --probe |
| |
| Retrieval configuration: --embed <model key> (embed_models.py), --mode dense|bm25|hybrid, --k 5, --per-speech N |
| (at most N passages from the same speech — the pilot's top-k was filled by one long speech). |
| """ |
| import argparse, json, logging, os, sys, time |
| from dotenv import load_dotenv |
| logging.getLogger("google_genai").setLevel(logging.ERROR) |
| from common import citation_of |
| from retrieval import Retriever |
|
|
| PROVIDERS = { |
| |
| "gemini": ("GEMINI_API_KEY", "gemini-3.5-flash-lite", ["gemma-4-26b-a4b-it", "gemini-flash-latest", "gemini-3.8-flash"]), |
| "anthropic": ("ANTHROPIC_API_KEY", "claude-haiku-5-5", []), |
| "openai": ("OPENAI_API_KEY", "gpt-5-mini", []), |
| } |
| DEFAULT_PROVIDER = "gemini" |
| DEFAULT_MODEL = None |
| K = 5 |
| EMBED = "sealion"; MODE = "hybrid"; PER_SPEECH = 2 |
|
|
| SYSTEM = """You are an assistant that explains what was said in Myanmar's parliament (Hluttaw) between 2011 and 2020. |
| You will be given a question and numbered passages. The passages are official parliamentary transcripts written in |
| Burmese; each begins with a header giving the house, date, agenda section, item and speaker. |
| Rules: |
| 1. Answer ONLY from the passages. Do not add outside knowledge or general background. |
| 2. Answer in the language of the QUESTION: a Burmese question gets a Burmese answer; an English question gets an |
| English answer that reports what the Burmese passages say (translate the content; keep names as in the passages). |
| Never refuse or switch language because the passages are in Burmese — read them and report them. |
| 3. Write 3–8 plain sentences a non-expert can follow. End every sentence that states a fact with the number(s) of |
| the passage(s) it comes from, like [2] or [1][3]. A sentence without a citation must not contain a fact. |
| 4. Say who said it (the MP, minister or ministry, as named in the passage header) and on what date. |
| 5. Copy numbers, dates, place names and amounts exactly as they appear in the passages. Do not round, convert or |
| infer numbers that are not there. |
| 6. If the passages do not answer the question, say so in one sentence (in the question's language) and do not guess. |
| If they answer only part of it, answer that part and say what is missing. |
| 7. If the question asks for an overview of what a person, party or house discussed "most" or "overall", say in one |
| sentence that the answer rests only on the retrieved passages, then report what those passages show. |
| 8. Speaker headers give each MP's constituency and, where known, party (e.g. "· အမျိုးသားဒီမိုကရေစီအဖွဲ့ချုပ် (NLD)"). |
| For a question about a party or group, use those labels to attribute statements; if no passage carries that label, |
| say the retrieved record does not identify speakers from that party, then still report what the passages say on the topic.""" |
|
|
| LANG_RULE = {"my": "ANSWER LANGUAGE: Burmese. Write the whole answer in Burmese, whatever language the question is in.", |
| "en": "ANSWER LANGUAGE: English. Write the whole answer in English, reporting what the Burmese passages say " |
| "(translate the content; keep personal names as in the passages, with a Latin transliteration in brackets if helpful)."} |
|
|
| def build_prompt(question, hits, lang=None): |
| ctx = "\n\n".join(f"[{i}] {citation_of(m)}\n{m['text']}" for i, (s, m) in enumerate(hits, 1)) |
| rule = ("\n" + LANG_RULE[lang]) if lang in LANG_RULE else "" |
| return f"{SYSTEM}{rule}\n\nPASSAGES:\n{ctx}\n\nQUESTION: {question}\n\nANSWER:" |
|
|
| def build_prompt_no_retrieval(question): |
| return (f"Answer the following question about what was said in Myanmar's parliament (Hluttaw) between 2011 and 2020. " |
| f"Answer in the language of the question, in 3–8 plain sentences, naming who said what and when as precisely as you can.\n\n" |
| f"QUESTION: {question}\n\nANSWER:") |
|
|
| class Generator: |
| def __init__(self, provider=DEFAULT_PROVIDER, model=DEFAULT_MODEL, no_fallback=False): |
| load_dotenv("FinalProject.env") |
| env, default, fallbacks = PROVIDERS[provider] |
| if no_fallback: fallbacks = [] |
| key = os.environ.get(env) |
| if not key: sys.exit(f"{env} missing — add it to FinalProject.env") |
| self.provider, self.model = provider, (model or default) |
| self.fallbacks = [m for m in fallbacks if m != self.model] |
| if provider == "gemini": |
| from google import genai; self.client = genai.Client(api_key=key) |
| elif provider == "anthropic": |
| import anthropic; self.client = anthropic.Anthropic(api_key=key) |
| elif provider == "openai": |
| import openai; self.client = openai.OpenAI(api_key=key) |
|
|
| def _generate(self, model, prompt): |
| """One call to the provider; returns the answer text.""" |
| if self.provider == "gemini": |
| return self.client.models.generate_content(model=model, contents=prompt).text or "" |
| if self.provider == "anthropic": |
| r = self.client.messages.create(model=model, max_tokens=1500, messages=[{"role": "user", "content": prompt}]) |
| return "".join(b.text for b in r.content if getattr(b, "type", "") == "text") |
| if self.provider == "openai": |
| r = self.client.chat.completions.create(model=model, messages=[{"role": "user", "content": prompt}]) |
| return r.choices[0].message.content or "" |
|
|
| def probe(self, candidates): |
| """Try each model once with a tiny prompt; report which ones respond.""" |
| ok = [] |
| for m in candidates: |
| t = time.time() |
| try: |
| text = self._generate(m, "Reply with the single word OK.") |
| print(f" {m:32} OK ({time.time()-t:.1f}s) {text.strip()[:40]!r}"); ok.append(m) |
| except Exception as e: |
| msg = str(e); code = getattr(e, "code", None) or getattr(e, "status_code", None) |
| short = "overloaded (503)" if "UNAVAILABLE" in msg else "rate-limited (429)" if "RESOURCE_EXHAUSTED" in msg else f"error {code}: {msg[:600]}" |
| print(f" {m:32} FAIL {short}") |
| print("\nworking:", ", ".join(ok) if ok else "none — wait a while and rerun, or try another --provider") |
|
|
| def list_models(self): |
| print(f"models available to this {self.provider} key:") |
| if self.provider == "gemini": |
| for m in self.client.models.list(): |
| if "generateContent" in (getattr(m, "supported_actions", None) or []): print(" ", m.name.replace("models/", "")) |
| elif self.provider == "anthropic": |
| for m in self.client.models.list(limit=100): print(" ", m.id) |
| elif self.provider == "openai": |
| for m in sorted(x.id for x in self.client.models.list()): print(" ", m) |
|
|
| def __call__(self, prompt): |
| """Returns (text, model_used). Retries on rate limits / overload; after 3 overload failures switches to the next fallback model.""" |
| while True: |
| for attempt in range(3): |
| try: |
| return self._generate(self.model, prompt), self.model |
| except Exception as e: |
| msg = str(e)[:300] |
| code = getattr(e, "code", None) or getattr(e, "status_code", None) |
| retryable = code in (429, 500, 502, 503, 504, 529) or "RESOURCE_EXHAUSTED" in msg or "UNAVAILABLE" in msg or "overloaded" in msg.lower() |
| if "exceeded your current quota" in msg or "quota" in msg.lower() and code == 429: |
| print(f" daily quota of {self.model} used up — switching model or stop and rerun with --resume after the reset (midnight Pacific = 2 pm Thailand)", file=sys.stderr) |
| break |
| print(f" [{self.model}] error: {msg}", file=sys.stderr) |
| if not retryable: |
| if code == 404 or "NOT_FOUND" in msg: |
| print(f" model '{self.model}' not found for this key — run python rag.py --list-models and pass --model <name>", file=sys.stderr) |
| sys.exit(1) |
| wait = 10 * (attempt + 1); print(f" retry in {wait}s", file=sys.stderr); time.sleep(wait) |
| if not self.fallbacks: |
| return "[generation failed]", self.model |
| nxt = self.fallbacks.pop(0) |
| print(f" {self.model} still unavailable — switching to {nxt}", file=sys.stderr) |
| self.model = nxt |
|
|
| def is_burmese(text): |
| return any("\u1000" <= ch <= "\u109f" for ch in text) |
|
|
| def translate_to_burmese(question, gen): |
| """Query translation for cross-lingual retrieval: a non-Burmese question is rendered in Burmese so that the |
| keyword index (and the dense model, in-language) can match names, places and terms in the record.""" |
| try: |
| text, _ = gen(f"Translate the following question into natural Burmese as a Myanmar citizen would ask it about " |
| f"parliament. Keep personal names in their usual Burmese form (e.g. Aung San Suu Kyi → ဒေါ်အောင်ဆန်းစုကြည်). " |
| f"Output only the Burmese translation, nothing else.\n\n{question}") |
| text = text.strip().splitlines()[0].strip() if text.strip() else "" |
| return text if is_burmese(text) else "" |
| except Exception: |
| return "" |
|
|
| def answer(question, ret, gen, k=K, no_retrieval=False, per_speech=PER_SPEECH, lang=None, translate=True): |
| """lang: None/"auto" = language of the question (detected, then stated explicitly to the model); |
| "my" or "en" = force the answer language. translate: for non-Burmese questions, also retrieve with a Burmese translation.""" |
| if lang in (None, "auto"): lang = "my" if is_burmese(question) else "en" |
| query_my = "" |
| if not no_retrieval and translate and not is_burmese(question): |
| query_my = translate_to_burmese(question, gen) |
| hits = [] if no_retrieval else ret.search(question, k=k, per_speech=per_speech, extra_queries=[query_my] if query_my else ()) |
| prompt = build_prompt_no_retrieval(question) if no_retrieval else build_prompt(question, hits, lang) |
| t = time.time(); text, used = gen(prompt); dt = time.time() - t |
| return {"question": question, "answer": text.strip(), "provider": gen.provider, "model": used, "retrieval": not no_retrieval, |
| "retriever": None if no_retrieval else ret.name, "k": k, "per_speech": per_speech, "lang": lang, "query_my": query_my or None, "gen_seconds": round(dt, 2), |
| "passages": [{"n": i, "score": round(s, 4), "chunk_id": m["chunk_id"], "speech_id": m["speech_id"], "citation": citation_of(m), "text": m["text"]} |
| for i, (s, m) in enumerate(hits, 1)]} |
|
|
| def show(res): |
| print(f"\nQ: {res['question']}\n\nA: {res['answer']}\n") |
| for p in res["passages"]: print(f"[{p['n']}] ({p['score']:.3f}) {p['citation']}") |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("question", nargs="?") |
| ap.add_argument("--batch", help="jsonl with a 'question' field per line") |
| ap.add_argument("--out", default=None) |
| ap.add_argument("--k", type=int, default=K) |
| ap.add_argument("--provider", default=DEFAULT_PROVIDER, choices=list(PROVIDERS)) |
| ap.add_argument("--model", default=DEFAULT_MODEL, help="model name for the provider (default: provider's default)") |
| ap.add_argument("--no-retrieval", action="store_true") |
| ap.add_argument("--resume", action="store_true", help="continue a batch: keep answers already in --out, skip their ids") |
| ap.add_argument("--no-fallback", action="store_true", help="never switch model; stop the batch instead (use for the SQ2 runs)") |
| ap.add_argument("--embed", default=EMBED, help="embedding model key (embed_models.py)") |
| ap.add_argument("--mode", default=MODE, choices=["dense", "bm25", "hybrid"]) |
| ap.add_argument("--per-speech", type=int, default=PER_SPEECH, help="max passages per speech in the top-k (0 = no cap)") |
| ap.add_argument("--list-models", action="store_true", help="show the models this API key can call, then exit") |
| ap.add_argument("--probe", action="store_true", help="try each candidate model once and report which respond, then exit") |
| args = ap.parse_args() |
| gen = Generator(args.provider, args.model, args.no_fallback) |
| if args.list_models: gen.list_models(); return |
| if args.probe: |
| extra = ["gemini-3.7-flash", "gemini-3.5-flash"] if args.provider == "gemini" else [] |
| gen.probe([gen.model] + gen.fallbacks + extra); return |
| idx = None if args.no_retrieval else Retriever(model=args.embed, mode=args.mode) |
| ps = args.per_speech or None |
| if args.batch: |
| qs = [json.loads(l) for l in open(args.batch, encoding="utf-8") if l.strip()] |
| out_path = args.out or args.batch.replace(".jsonl", "_answers.jsonl") |
| done = set() |
| if args.resume and os.path.exists(out_path): |
| done = {json.loads(l)["id"] for l in open(out_path, encoding="utf-8") if l.strip() and "[generation failed]" not in l} |
| print(f"resuming: {len(done)} answers already in {out_path}") |
| out = open(out_path, "a" if args.resume else "w", encoding="utf-8") |
| for i, q in enumerate(qs, 1): |
| qid = q.get("id", i) |
| if qid in done: continue |
| res = answer(q["question"], idx, gen, args.k, args.no_retrieval, ps); res["id"] = qid |
| if res["answer"] == "[generation failed]": |
| print(f"{i}/{len(qs)} {qid}: generation failed — stopping; rerun with --resume later", flush=True); break |
| out.write(json.dumps(res, ensure_ascii=False) + "\n"); out.flush() |
| print(f"{i}/{len(qs)} done ({res['gen_seconds']}s)", flush=True) |
| print(f"→ {out_path}") |
| elif args.question: |
| show(answer(args.question, idx, gen, args.k, args.no_retrieval, ps)) |
| else: |
| print("Ask a question (empty line to quit)") |
| while True: |
| try: q = input("\n> ").strip() |
| except (EOFError, KeyboardInterrupt): break |
| if not q: break |
| show(answer(q, idx, gen, args.k, args.no_retrieval, ps)) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
