| """Build train JSON with top-3 AfriE5-similar same-subset (Q,A) demos prepended. |
| Source: msrh_rag_train_afrie5_TV_k3.json (36501 rows = Train+Val). |
| Per row: find top-3 same-subset OTHER (Q,A) pairs, exclude self by ID. |
| """ |
| import os, json, time, pathlib |
| os.environ["HF_HOME"] = "/mnt/msrh/Magic_submission/hf_cache" |
| os.environ.setdefault("CUDA_VISIBLE_DEVICES", "4") |
| import numpy as np |
| import pandas as pd |
| from sentence_transformers import SentenceTransformer |
|
|
| WORK = pathlib.Path("/mnt/msrh/Magic_submission") |
| TRAIN_CSV = WORK/"data/Train.csv" |
| VAL_CSV = WORK/"data/Val.csv" |
| SRC_JSON = WORK/"LF/data/msrh_rag_train_afrie5_TV_k3.json" |
| OUT_JSON = WORK/"LF/data/msrh_rag_train_afrie5_TV_k3_fewshot.json" |
|
|
| train = pd.read_csv(TRAIN_CSV).reset_index(drop=True) |
| val = pd.read_csv(VAL_CSV).reset_index(drop=True) |
| tv_pool = pd.concat([train, val], ignore_index=True) |
| print(f"TV pool size: {len(tv_pool)}") |
|
|
| MODEL = "/mnt/msrh/Magic_submission/hub/AfriE5-Large-instruct" |
| print(f"[{time.strftime('%H:%M:%S')}] loading {MODEL}") |
| model = SentenceTransformer(MODEL, device="cuda", trust_remote_code=True) |
| QUERY_INST = "Instruct: Given a maternal/sexual/reproductive health question, retrieve relevant passages that answer the question\nQuery: " |
|
|
| K = 3 |
| demos_by_id = {} |
|
|
| for subset in sorted(tv_pool["subset"].unique()): |
| pool_sub = tv_pool[tv_pool["subset"] == subset].reset_index(drop=True) |
| print(f"[{time.strftime('%H:%M:%S')}] subset={subset} n={len(pool_sub)}") |
| pool_emb = model.encode([QUERY_INST + str(q) for q in pool_sub["input"]], |
| batch_size=32, normalize_embeddings=True, |
| show_progress_bar=False, convert_to_numpy=True) |
| sims = pool_emb @ pool_emb.T |
| np.fill_diagonal(sims, -1.0) |
| top_idx = np.argsort(-sims, axis=1)[:, :K] |
| for i, row in pool_sub.iterrows(): |
| demos_by_id[row["ID"]] = [ |
| (str(pool_sub.iloc[int(j)]["input"]), str(pool_sub.iloc[int(j)]["output"])) |
| for j in top_idx[i] |
| ] |
| print(f"[{time.strftime('%H:%M:%S')}] demos built for {len(demos_by_id)} rows") |
|
|
| |
| n_inj = 0 |
| n_skip = 0 |
| with open(SRC_JSON) as fin, open(OUT_JSON, "w") as fout: |
| for i, ln in enumerate(fin): |
| r = json.loads(ln) |
| pool_row = tv_pool.iloc[i] |
| dem = demos_by_id.get(pool_row["ID"], []) |
| if dem: |
| content = r["messages"][0]["content"] |
| idx = content.find("Context 1:") |
| if idx > 0: |
| before = content[:idx].rstrip() |
| after = content[idx:] |
| ex_block = "\n\n".join([f"Example {k+1}:\nQ: {q}\nA: {a}" for k, (q, a) in enumerate(dem)]) |
| new_content = f"{before}\n\nFew-shot examples (similar questions from training):\n{ex_block}\n\n{after}" |
| r["messages"][0]["content"] = new_content |
| n_inj += 1 |
| else: |
| n_skip += 1 |
| else: |
| n_skip += 1 |
| fout.write(json.dumps(r, ensure_ascii=False) + "\n") |
| print(f"\nInjected demos in {n_inj} rows, skipped {n_skip}") |
| print(f"wrote {OUT_JSON}") |
|
|
