"""① 监督式查询分解（本地 Llama + LoRA，transformers + peft，进程内）。

首次调用 supervised_decompose() 时惰性加载模型。需要环境变量：
  RAGQA_LLAMA_MODEL            Llama-3.1-8B-Instruct 基座目录
  RAGQA_LLAMA_LORA_SUPERVISED  监督式分解 LoRA 适配器目录（可选；不设则用基座本身）
"""

import os

_HERE = os.path.dirname(os.path.abspath(__file__))
_PKG_ROOT = os.path.dirname(_HERE)  # .../query_decomposition_baselines

_model = None
_tokenizer = None


def _load():
    global _model, _tokenizer
    if _model is not None:
        return
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM

    base = os.environ.get("RAGQA_LLAMA_MODEL", "")
    lora = os.environ.get("RAGQA_LLAMA_LORA_SUPERVISED", "")
    # 未设 LoRA 路径时，自动回退到仓库内 weights/ 下的适配器（HF 仓库 clone 后即可用）。
    if not lora:
        cand = os.path.join(_PKG_ROOT, "weights", "decom_baseline_strategyqa_lora")
        if os.path.isdir(cand):
            lora = cand
    assert base, "[qd] 请设置环境变量 RAGQA_LLAMA_MODEL 指向 Llama-3.1-8B-Instruct 基座目录"

    _tokenizer = AutoTokenizer.from_pretrained(base)
    model = AutoModelForCausalLM.from_pretrained(
        base, torch_dtype=torch.float16, device_map="auto"
    )
    if lora:
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, lora)
    _model = model.eval()


def supervised_decompose(query, max_new_tokens=512):
    """把一个 query 交给监督式 LoRA 模型，返回子查询 list[str]（模型输出按 '|' 切分）。"""
    import torch
    _load()
    messages = [{"role": "user", "content": query}]
    input_ids = _tokenizer.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
    ).to(_model.device)
    with torch.no_grad():
        out = _model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=_tokenizer.eos_token_id,
        )
    text = _tokenizer.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True)
    subs = [s.strip() for s in text.split("|") if s.strip()]
    return subs or [query]