File size: 2,269 Bytes
cd6775d | 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 | """① 监督式查询分解（本地 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]
|