qd_baselines/_icl.py

"""③ ICL 查询分解（Socratic Questioning，进程内）。

复用 methods/icl_socratic/.../code/socratic_tree.py 的 SocraticTree，不再启动 Flask 服务。
默认用 gpt-4o-mini（需 CHATANYWHERE_API_KEY）；也可传入 llm_fn 使用本地/自定义 LLM。
注意：Socratic 内部会调用检索（原实现请求 localhost:8895 的 ColBERT）；进程内版默认注入
空检索桩（不发任何 HTTP），纯查询分解无需真实检索。需要时可传 retriever_fn 接你的检索器。
"""

import os
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
_PKG_ROOT = os.path.dirname(_HERE)
_SOC = os.path.join(_PKG_ROOT, "methods", "icl_socratic", "socratic_questioning_textOnly")
_CODE_DIR = os.path.join(_SOC, "code")
_DATA_DIR = os.path.join(_SOC, "data")
_RESULTS_DIR = os.path.join(_SOC, "results", "functional")

_tree = None
_stmod = None


def _build():
    global _tree, _stmod
    if _tree is not None:
        return
    import json
    if _CODE_DIR not in sys.path:
        sys.path.insert(0, _CODE_DIR)
    import socratic_tree as stmod
    prompt_map = json.load(open(os.path.join(_DATA_DIR, "prompt_map.json"), "r", encoding="utf-8"))
    os.makedirs(_RESULTS_DIR, exist_ok=True)
    _stmod = stmod
    _tree = stmod.SocraticTree(
        backbone="gpt",
        api=os.environ.get("CHATANYWHERE_API_KEY", ""),
        prompt_map=prompt_map,
        dataset="lqa",
        num_question=2,
        max_turn=0,
        max_depth=0,
        save_dir=_RESULTS_DIR,
    )


def icl_decompose(query, llm_fn=None, retriever_fn=None):
    """把一个 query 交给 Socratic 分解器，返回子问题 list[str]。

    llm_fn: 可选 callable(system_define:str, prompt:str)->str，替代默认 gpt-4o-mini
            （例如包一个本地 chat 模型），实现「完全不连外部 LLM」。
    retriever_fn: 可选 callable(query:str)->str（top-1 passage 文本）。Socratic 内部会检索，
            默认注入空检索桩（返回 ""），从而**完全不发 HTTP**；纯查询分解无需真实检索。
            若想要完整苏格拉底流程，可传入你自己的检索器。
    """
    _build()
    if llm_fn is not None:
        _stmod.set_llm(llm_fn)
    # 默认用空检索桩，避免任何 HTTP（原实现会请求 localhost:8895 的 ColBERT）。
    _stmod.set_retriever(retriever_fn if retriever_fn is not None else (lambda q: ""))
    subs = _tree.start(1, query, options=None, context=None)
    # 去重并保序
    subs = list(dict.fromkeys([s.strip() for s in subs if s and s.strip()]))
    return subs or [query]