examples/demo_decompose.py

"""示例：把你自己的一批 query 逐条交给某个 baseline 分解，结果写成 jsonl。

①②③ 为进程内函数调用，**无需启动服务**；填好 .env（模型/密钥）后直接运行：
  python examples/demo_decompose.py --method icl --input examples/sample_queries.txt
  python examples/demo_decompose.py --method unsupervised --input 你的问题.txt --output out.jsonl
（首次调用会惰性加载对应模型。）

--input：纯文本文件，每行一个 query（换成你自己的数据即可）。
输出：每行一个 {"query": ..., "sub_queries": [...]} 的 jsonl。
"""

import argparse
import json
import os
import sys

# 让 `import decompose` 找到上级目录的 decompose.py
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from decompose import METHODS  # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--method", choices=list(METHODS), default="icl",
                    help="用哪个 baseline 分解：supervised / unsupervised / icl")
    ap.add_argument("--input", required=True, help="每行一个 query 的文本文件")
    ap.add_argument("--output", default=None, help="输出 jsonl（默认打印到屏幕）")
    args = ap.parse_args()

    decompose = METHODS[args.method]
    with open(args.input, "r", encoding="utf-8") as f:
        queries = [ln.strip() for ln in f if ln.strip()]

    out = open(args.output, "w", encoding="utf-8") if args.output else None
    for q in queries:
        try:
            subs = decompose(q)
        except Exception as e:  # 服务没起/网络问题时给出清晰提示
            print(f"[warn] 分解失败（服务是否已启动？）: {q}\n       {e}", file=sys.stderr)
            continue
        rec = {"method": args.method, "query": q, "sub_queries": subs}
        line = json.dumps(rec, ensure_ascii=False)
        if out:
            out.write(line + "\n")
        else:
            print(line)
    if out:
        out.close()
        print(f"已写出 -> {args.output}")


if __name__ == "__main__":
    main()