Upload 5 files

context_retreiver.py

@@ -0,0 +1,75 @@
+# context_retriever.py
+import os, re, json, pickle, logging, numpy as np, faiss
+from tqdm.notebook import tqdm
+from sentence_transformers import SentenceTransformer
+from langchain_community.retrievers import BM25Retriever
+from langchain.docstore.document import Document
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+WORK = "context"
+JSONL = f"{WORK}/rag_documents.jsonl"
+FAISS_INDEX = f"{WORK}/faiss_ivf.index"
+BM25_PICKLE = f"{WORK}/bm25_retriever.pkl"
+
+logger.info("Loading all RAG documents...")
+with open(JSONL, encoding='utf-8') as f:
+    ALL_DOCS = [json.loads(line) for line in f]
+
+LINE_TO_TEXT = {i: doc["text"] for i, doc in enumerate(ALL_DOCS)}
+LINE_TO_META = {i: doc["metadata"] for i, doc in enumerate(ALL_DOCS)}
+
+class HybridRetriever:
+    def __init__(self):
+        # FAISS CPU
+        self.faiss_index = faiss.read_index(FAISS_INDEX)
+        logger.info(f"FAISS loaded ({self.faiss_index.ntotal:,} vectors)")
+
+        # SentenceTransformer (GPU if available)
+        self.model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2",
+                                         device="cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu")
+
+        # BM25
+        if os.path.exists(BM25_PICKLE):
+            self.bm25 = pickle.load(open(BM25_PICKLE, "rb"))
+            logger.info("BM25 loaded")
+        else:
+            logger.info("Building BM25...")
+            docs = [Document(page_content=re.sub(r"^Filename:.*\nFullPath:.*\n\n", "",
+                                                 doc["text"], flags=re.M),
+                             metadata=doc["metadata"]) for doc in ALL_DOCS]
+            self.bm25 = BM25Retriever.from_documents(docs)
+            self.bm25.k = 30
+            pickle.dump(self.bm25, open(BM25_PICKLE, "wb"))
+            logger.info("BM25 built and saved")
+
+    def batch_retrieve(self, queries, top_k=3, faiss_k=10, bm25_k=3):
+        qvecs = self.model.encode(queries, show_progress_bar=False, normalize_embeddings=True).astype("float32")
+        D, I = self.faiss_index.search(qvecs, faiss_k)
+
+        batch_results = []
+        for qi, (scores, indices) in enumerate(zip(D, I)):
+            results = []
+            seen = set()
+            for score, idx in zip(scores, indices):
+                if idx == -1 or idx in seen: continue
+                results.append({"score": float(score), "text": LINE_TO_TEXT[idx],
+                                "metadata": LINE_TO_META[idx], "source": "FAISS"})
+                seen.add(idx)
+                if len(results) >= top_k: break
+
+            # BM25
+            bm25_docs = self.bm25.invoke(queries[qi])
+            for doc in bm25_docs[:bm25_k]:
+                ln = doc.metadata.get("line_no")
+                if ln in seen: continue
+                results.append({"score": 0.0, "text": LINE_TO_TEXT.get(ln, ""),
+                                "metadata": LINE_TO_META.get(ln, doc.metadata), "source": "BM25"})
+                seen.add(ln)
+                if len(results) >= top_k: break
+            batch_results.append(results)
+        return batch_results
+
+# Singleton retriever
+retriever = HybridRetriever()

full_rag.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9c92278e3df812534acaa211928b76a888453c81cfbe6b70bdea5d5cb330c61
+size 1597083267

prompter.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""
+generate_prompts_v8_batch_fixed.py
+
+- Uses batch retrieval for Context, QA, and Relationships
+- Saves in batches with checkpointing
+- Pads contexts and QA to fixed sizes
+- Appends metadata at the end
+"""
+
+import os, json, torch, numpy as np
+from pathlib import Path
+from tqdm import tqdm
+from sentence_transformers import SentenceTransformer
+from concurrent.futures import ThreadPoolExecutor
+
+from context_retreiver import retriever as context_retriever
+from qa_retreiver import search_topk as qa_retreiver
+from relationships_retreiver import batch_relationships
+
+QA_FILE = Path("got_all_qa_final.json")
+OUT_DIR = Path("prompts_out")
+CHECKPOINT_FILE = OUT_DIR / "checkpoint.json"
+SAVE_BATCH_SIZE = 512
+EMBED_BATCH_SIZE = 32  # GPU batch size
+
+DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
+print(f"[INFO] Using device: {DEVICE}")
+
+EMBED_MODEL = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device=DEVICE)
+
+STRUCTURAL_TOKENS = [
+    "<|CTX_QA|>", "<|/CTX_QA|>",
+    "<|CTX_REL|>", "<|/CTX_REL|>",
+    "<|INSTR|>", "<|/INSTR|>",
+    "<|QUESTION|>", "<|/QUESTION|>",
+    "<|ANSWER|>", "<|/ANSWER|>",
+    "<|QA_SIM_1|>", "<|/QA_SIM_1|>",
+    "<|QA_SIM_2|>", "<|/QA_SIM_2|>",
+    "<|QA_SIM_3|>", "<|/QA_SIM_3|>",
+    "<|QA_SIM_4|>", "<|/QA_SIM_4|>",
+    "<|QA_SIM_5|>", "<|/QA_SIM_5|>"
+]
+
+def read_checkpoint(): 
+    if CHECKPOINT_FILE.exists():
+        try: 
+            return int(json.loads(CHECKPOINT_FILE.read_text())["next_index"])
+        except: 
+            return 0
+    return 0
+
+def write_checkpoint(idx): 
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    CHECKPOINT_FILE.write_text(json.dumps({"next_index": idx}))
+
+def metadata_to_str(meta):
+    if not meta: return ""
+    return "; ".join(f"{k}={v}" for k,v in meta.items() if isinstance(v,(str,int,float,bool)))
+
+def append_metadata_at_end(answer, context1_text, context1_meta):
+    parts=[]
+    if answer: parts.append(answer.strip())
+    if context1_text: parts.append(f"[Context1: {context1_text.strip()}]")
+    meta_str = metadata_to_str(context1_meta)
+    if meta_str: parts.append(f"(meta: {meta_str})")
+    return " ".join(parts)
+
+def build_prompt(ctx_texts, rel_text, sim_qas, question):
+    parts=[]
+    # ctx_texts = [ctx2, ctx3]
+    for ctx in ctx_texts:
+        if ctx: parts.append(f"<|CTX_QA|> {ctx} <|/CTX_QA|>")
+    if rel_text: parts.append(f"<|CTX_REL|> {rel_text} <|/CTX_REL|>")
+    for i in range(5):
+        if i < len(sim_qas):
+            qa = sim_qas[i]
+            parts.append(f"<|QA_SIM_{i+1}|> Q: {qa['question']} A: {qa['answer']} <|/QA_SIM_{i+1}|>")
+        else:
+            parts.append(f"<|QA_SIM_{i+1}|> <|/QA_SIM_{i+1}|>")
+    parts.append("<|INSTR|> Use above contexts to answer concisely. <|/INSTR|>")
+    parts.append(f"<|QUESTION|> {question} <|/QUESTION|>")
+    parts.append("<|ANSWER|>")
+    return "\n\n".join(parts)
+
+def retrieve_contexts(questions, top_k=3):
+    """Batch retrieve context texts + metadata"""
+    batch_res = context_retriever.batch_retrieve(questions, top_k=top_k)
+    contexts=[]
+    for res_list in batch_res:
+        ctx_texts = [r["text"] for r in res_list[:top_k]]
+        ctx_metas = [r["metadata"] for r in res_list[:top_k]]
+        # pad to top_k
+        while len(ctx_texts)<top_k: ctx_texts.append(""); ctx_metas.append({})
+        contexts.append((ctx_texts, ctx_metas))
+    return contexts
+
+def retrieve_qas_and_rels(questions, max_workers=20):
+    """Threaded retrieval of QA and relationships"""
+    sim_qas_list=[]
+    rel_list=[]
+    with ThreadPoolExecutor(max_workers=max_workers) as ex:
+        sim_qas_list = list(ex.map(lambda q: qa_retreiver([q], k=5), questions))
+        rel_list = list(ex.map(lambda q: batch_relationships([q], top_k=1)[0], questions))
+    return sim_qas_list, rel_list
+
+def main():
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    with open(QA_FILE,'r',encoding='utf-8') as f:
+        qas = json.load(f)
+    total = len(qas)
+    start_idx = read_checkpoint()
+    if start_idx >= total:
+        print("[INFO] Checkpoint beyond dataset length.")
+        return
+
+    prompts_accum=[]
+    batch_count=start_idx//SAVE_BATCH_SIZE
+
+    for batch_start in tqdm(range(start_idx, total, EMBED_BATCH_SIZE)):
+        batch_end = min(batch_start + EMBED_BATCH_SIZE, total)
+        batch_items = qas[batch_start:batch_end]
+        questions = [it.get("question") or it.get("q") or it.get("Question") for it in batch_items]
+        orig_answers = [it.get("answer") or it.get("a") or it.get("Answer","") for it in batch_items]
+
+        # --- retrieve contexts ---
+        contexts = retrieve_contexts(questions, top_k=3)
+        # --- QA & relationships ---
+        sim_qas_list, rel_list = retrieve_qas_and_rels(questions)
+
+        for i,q in enumerate(questions):
+            if not q:
+                write_checkpoint(batch_start+i+1)
+                continue
+            ctx_texts, ctx_metas = contexts[i]
+            context1, context2, context3 = ctx_texts
+            meta1 = ctx_metas[0]
+            prompt_text = build_prompt([context2, context3], rel_list[i], sim_qas_list[i], q)
+            gold = append_metadata_at_end(orig_answers[i], context1, meta1)
+
+            obj={
+                "id": batch_start+i,
+                "question": q,
+                "prompt": prompt_text,
+                "gold_answer": gold,
+                "context1": context1,
+                "retrieved_qas": sim_qas_list[i],
+                "relation_text": rel_list[i]
+            }
+            prompts_accum.append(obj)
+
+            # --- Save batch ---
+            if len(prompts_accum)>=SAVE_BATCH_SIZE:
+                out_path = OUT_DIR/f"prompts_batch_{batch_count:03d}.json"
+                out_path.write_text(json.dumps(prompts_accum, ensure_ascii=False, indent=2),encoding='utf-8')
+                batch_count+=1
+                prompts_accum=[]
+
+            write_checkpoint(batch_start+i+1)
+
+    # save remaining
+    if prompts_accum:
+        out_path = OUT_DIR/f"prompts_batch_{batch_count:03d}.json"
+        out_path.write_text(json.dumps(prompts_accum, ensure_ascii=False, indent=2))
+
+    OUT_DIR.joinpath("special_tokens_used.txt").write_text("\n".join(STRUCTURAL_TOKENS))
+    print("[DONE] All prompts processed.")
+
+if __name__=="__main__":
+    main()

qa_retreiver.py

@@ -0,0 +1,68 @@
+# qa_retriever.py
+import os, pickle, faiss
+from sentence_transformers import SentenceTransformer
+from typing import List, Dict, Any, Optional
+
+MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+CLEAN_JSON = "qa_pairs/asoiaf_qa_clean.json"
+INDEX_FILE = "qa_pairs/faiss_index.index"
+QA_DATA_FILE = "qa_pairs/qa_data.pkl"
+
+EMBED_MODEL: Optional[SentenceTransformer] = None
+INDEX = None
+QA_PAIRS: List[Dict[str, Any]] = []
+
+def _load_embed_model():
+    global EMBED_MODEL
+    if EMBED_MODEL is None:
+        EMBED_MODEL = SentenceTransformer(MODEL_NAME,
+                                          device="cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu")
+    return EMBED_MODEL
+
+def build_or_load_index():
+    global INDEX, QA_PAIRS
+    if INDEX and QA_PAIRS: return INDEX, QA_PAIRS, EMBED_MODEL
+
+    INDEX = faiss.read_index(INDEX_FILE)
+    with open(QA_DATA_FILE, "rb") as f:
+        QA_PAIRS = pickle.load(f)
+    _load_embed_model()
+    return INDEX, QA_PAIRS, EMBED_MODEL
+def search_topk(query: str, index=None, qa_pairs=None, model=None, k: int = 5):
+    """
+    Returns up to `k` similar Q&A entries as list of dicts.
+    Handles single query string input safely.
+    """
+    if not isinstance(query, list):
+        query_list = [query]
+    else:
+        query_list = query
+
+    if model is None:
+        model = _load_embed_model()
+    if index is None or qa_pairs is None:
+        index, qa_pairs, model = build_or_load_index()
+
+    q_vecs = model.encode(query_list, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False).astype("float32")
+
+    results = []
+    for q_vec in q_vecs:
+        scores, indices = index.search(q_vec[None, :], k*3)
+        seen = set()
+        q_results = []
+        for score, idx in zip(scores[0], indices[0]):
+            if len(q_results) >= k: break
+            if idx < 0 or idx >= len(qa_pairs): continue
+            q_text = qa_pairs[idx].get("question", "")
+            if q_text in seen: continue
+            seen.add(q_text)
+            raw_ans = qa_pairs[idx].get("answer", "")
+            clean_ans = raw_ans.split("\n\nReference:")[0].strip()
+            q_results.append({
+                "similarity": float(score),
+                "question": q_text,
+                "answer": clean_ans
+            })
+        results.append(q_results)
+
+    return results[0] if len(results) == 1 else results

relationships_retreiver.py

@@ -0,0 +1,49 @@
+# relationship_retriever.py
+import os, pickle, logging
+import faiss
+from sentence_transformers import SentenceTransformer
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+RELATIONS = "relations"
+REL_INDEX = f"{RELATIONS}/got_rels.faiss"
+REL_DATA = f"{RELATIONS}/got_rels_meta.pkl"
+
+logger.info("Loading relationship FAISS index...")
+rel_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2",
+                                device="cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu")
+rel_index = faiss.read_index(REL_INDEX)
+with open(REL_DATA, "rb") as f:
+    rel_data = pickle.load(f)
+name_map = rel_data["name_map"]
+
+def batch_relationships(questions, top_k=3):
+    batch_results = []
+    for q in questions:
+        q_upper = q.upper()
+        candidates = []
+        for variant in name_map.keys():
+            if len(variant) < 3: continue
+            if variant in q_upper or variant.replace(" ","") in q_upper.replace(" ",""):
+                candidates.append(name_map[variant])
+        candidates = list(dict.fromkeys(candidates))[:2]
+        if not candidates:
+            batch_results.append(["No known character relationships found"])
+            continue
+
+        query = f"Relationships of {' and '.join(candidates)} in Game of Thrones books"
+        q_vec = rel_model.encode([query], normalize_embeddings=True, show_progress_bar=False).astype("float32")
+        D, I = rel_index.search(q_vec, top_k*2)
+        results = []
+        seen = set()
+        for idx in I[0]:
+            if idx == -1: continue
+            sent = rel_data["sentences"][idx]
+            char = rel_data["metadata"][idx]["display_name"]
+            if char not in seen:
+                results.append(sent)
+                seen.add(char)
+            if len(results) >= top_k: break
+        batch_results.append(results if results else ["No confirmed relationships found"])
+    return batch_results