File size: 4,959 Bytes
4c67792 2d61d41 4c67792 8b2d098 4c67792 2d61d41 4c67792 | 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 86 87 88 89 90 91 92 93 94 95 | """
Configuration dataclass for the passage-entity KG pipeline in QAFD-RAG.
Combines passage-entity BaseConfig fields with QAFD algorithm parameters.
"""
from dataclasses import dataclass, field
from typing import Optional, Literal
@dataclass
class PassageEntityConfig:
"""Unified configuration for indexing, retrieval, and QAFD."""
# ── LLM ────────────────────────────────────────────────────────────────
llm_model: str = "gpt-4o-mini"
llm_base_url: str = "https://api.openai.com/v1"
llm_api_key: str = "" # falls back to OPENAI_API_KEY env
max_new_tokens: Optional[int] = 2048
temperature: float = 0.0
# ── Embedding ──────────────────────────────────────────────────────────
embedding_model_key: str = "nvidia-nv-embed-v2" # key in QAFD-RAG registry
embedding_batch_size: int = 16
# ── Dataset / paths ────────────────────────────────────────────────────
dataset: Optional[str] = None # musique, hotpotqa, 2wikimultihopqa
save_dir: str = "outputs"
force_index_from_scratch: bool = False
force_openie_from_scratch: bool = False
save_openie: bool = True
# ── Graph construction ─────────────────────────────────────────────────
is_directed_graph: bool = False
synonymy_edge_topk: int = 2047
synonymy_edge_query_batch_size: int = 1000
synonymy_edge_key_batch_size: int = 10000
synonymy_edge_sim_threshold: float = 0.8
# ── Retrieval ──────────────────────────────────────────────────────────
linking_top_k: int = 10
retrieval_top_k: int = 200
passage_node_weight: float = 0.05
damping: float = 0.5
# ── QA ─────────────────────────────────────────────────────────────────
qa_top_k: int = 5
# ── QAFD algorithm parameters ──────────────────────────────────────────
use_qafd: bool = True
qafd_alpha: float = 1.5
qafd_epsilon: float = 0.01
qafd_max_iterations: int = 500
qafd_weight_scheme: str = "multiply" # "multiply", "add", "original"
qafd_use_node_degree: bool = True
qafd_step_size: float = 0.2
qafd_random_seed: int = 42
# ── Query-aware enhancements (all default = original behaviour) ────────
sim_mode: str = "normalized" # Similarity contrast: "normalized", "relu", "relu_sq"
qa_sink_gamma: float = 0.0 # Query-aware sink capacity (0=off)
qa_warm_delta: float = 0.0 # Query-aware seed bias (0=off)
qa_warm_walk: bool = False # Use QA edge weights in warm-start walk
qa_warm_steps: int = 2 # Number of warm-start steps (default 2)
qa_accum_gamma: float = 0.0 # Query-aware x accumulation boost (0=off)
qa_post_lambda: float = 0.0 # Post-diffusion reranking (0=off)
batch_push: bool = False # Batch push-relabel (process all excess nodes per iter)
# ── Reranker ───────────────────────────────────────────────────────────
rerank_dspy_file_path: Optional[str] = None # path to DSPy JSON; None → built-in prompt
def __post_init__(self):
if self.save_dir == "outputs" and self.dataset:
self.save_dir = f"outputs/{self.dataset}"
@property
def working_dir(self) -> str:
"""Model-specific sub-directory under save_dir.
Also checks kg/multihop/ for pre-downloaded KGs from HuggingFace.
If found there, uses that path instead of outputs/.
"""
import os
llm_label = self.llm_model.replace("/", "_")
emb_label = self.embedding_model_key.replace("/", "_")
# Check HuggingFace download location (kg/multihop/{llm}_{emb}_{dataset}/)
if self.dataset:
for task_dir in ["multihop", "ultradomain"]:
hf_path = os.path.join("kg", task_dir, f"{llm_label}_{emb_label}_{self.dataset}")
if os.path.isdir(hf_path) and os.path.exists(os.path.join(hf_path, "graph.pickle")):
return hf_path
return f"{self.save_dir}/{llm_label}_{emb_label}"
|