139 lines
5.3 KiB
Python
139 lines
5.3 KiB
Python
import json
|
||
import os
|
||
from typing import Any, Dict, List
|
||
|
||
from torch.utils.data import Dataset
|
||
try:
|
||
from datasets import load_dataset
|
||
except ImportError:
|
||
load_dataset = None
|
||
|
||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
|
||
|
||
SEARCH_SHORTQA_SYSTEM = (
|
||
"你是一个搜索增强短问答助手。你只能根据给定检索片段回答问题；"
|
||
"如果证据不足，回答“不确定”。输出必须是JSON，格式为"
|
||
"{\"answer\":\"...\", \"citations\":[\"D1\"]}。answer必须简短，citations只能引用给定片段ID。"
|
||
)
|
||
|
||
|
||
def _as_list(value):
|
||
if value is None:
|
||
return []
|
||
if isinstance(value, list):
|
||
return value
|
||
return [value]
|
||
|
||
|
||
def normalize_contexts(sample: Dict[str, Any]) -> List[Dict[str, str]]:
|
||
contexts = (
|
||
sample.get("contexts")
|
||
or sample.get("documents")
|
||
or sample.get("ctxs")
|
||
or sample.get("passages")
|
||
or sample.get("snippets")
|
||
or sample.get("search_results")
|
||
or []
|
||
)
|
||
normalized = []
|
||
for i, item in enumerate(contexts):
|
||
if isinstance(item, str):
|
||
normalized.append({"id": f"D{i + 1}", "title": "", "text": item})
|
||
continue
|
||
doc_id = str(item.get("id") or item.get("doc_id") or item.get("source") or f"D{i + 1}")
|
||
title = str(item.get("title") or item.get("name") or "")
|
||
text = str(item.get("text") or item.get("content") or item.get("snippet") or "")
|
||
normalized.append({"id": doc_id, "title": title, "text": text})
|
||
return normalized
|
||
|
||
|
||
def normalize_answers(sample: Dict[str, Any]) -> List[str]:
|
||
answers = sample.get("answers")
|
||
if answers is None:
|
||
answers = sample.get("answer")
|
||
if answers is None:
|
||
answers = sample.get("short_answers")
|
||
if answers is None:
|
||
answers = sample.get("output")
|
||
return [str(a) for a in _as_list(answers) if str(a).strip()]
|
||
|
||
|
||
def normalize_gold_doc_ids(sample: Dict[str, Any]) -> List[str]:
|
||
doc_ids = sample.get("gold_doc_ids") or sample.get("citations") or sample.get("supporting_doc_ids")
|
||
return [str(x) for x in _as_list(doc_ids)]
|
||
|
||
|
||
def select_contexts(sample: Dict[str, Any], max_contexts=6) -> List[Dict[str, str]]:
|
||
contexts = normalize_contexts(sample)
|
||
if max_contexts <= 0 or len(contexts) <= max_contexts:
|
||
return contexts
|
||
gold_ids = set(normalize_gold_doc_ids(sample))
|
||
gold_contexts = [doc for doc in contexts if doc["id"] in gold_ids]
|
||
other_contexts = [doc for doc in contexts if doc["id"] not in gold_ids]
|
||
return (gold_contexts + other_contexts)[:max_contexts]
|
||
|
||
|
||
def format_search_prompt(sample: Dict[str, Any], max_contexts=6, max_chars_per_context=700) -> str:
|
||
question = str(sample.get("question") or sample.get("query") or sample.get("prompt") or "")
|
||
contexts = select_contexts(sample, max_contexts)
|
||
lines = ["请基于以下检索片段回答问题。", "", f"问题：{question}", "", "检索片段："]
|
||
for doc in contexts:
|
||
title = f"｜{doc['title']}" if doc["title"] else ""
|
||
text = doc["text"].replace("\n", " ").strip()[:max_chars_per_context]
|
||
lines.append(f"[{doc['id']}]{title} {text}")
|
||
lines.extend(["", "要求：输出JSON；answer简短；citations只填写真正支撑答案的片段ID。"])
|
||
return "\n".join(lines)
|
||
|
||
|
||
def build_search_shortqa_messages(sample: Dict[str, Any], max_contexts=6, max_chars_per_context=700):
|
||
return [
|
||
{"role": "system", "content": SEARCH_SHORTQA_SYSTEM},
|
||
{"role": "user", "content": format_search_prompt(sample, max_contexts, max_chars_per_context)},
|
||
]
|
||
|
||
|
||
def build_sft_conversation(sample: Dict[str, Any], max_contexts=6, max_chars_per_context=700):
|
||
answers = normalize_answers(sample)
|
||
gold_doc_ids = normalize_gold_doc_ids(sample)
|
||
answer = answers[0] if answers else "不确定"
|
||
assistant = {
|
||
"answer": answer,
|
||
"citations": gold_doc_ids[:3],
|
||
}
|
||
return build_search_shortqa_messages(sample, max_contexts, max_chars_per_context) + [
|
||
{"role": "assistant", "content": json.dumps(assistant, ensure_ascii=False)}
|
||
]
|
||
|
||
|
||
class SearchShortQARLDataset(Dataset):
|
||
def __init__(self, jsonl_path, tokenizer, max_length=1024, max_contexts=6, max_chars_per_context=700):
|
||
super().__init__()
|
||
if load_dataset is None:
|
||
raise ImportError("SearchShortQARLDataset requires the 'datasets' package. Install project requirements before training.")
|
||
self.tokenizer = tokenizer
|
||
self.max_length = max_length
|
||
self.max_contexts = max_contexts
|
||
self.max_chars_per_context = max_chars_per_context
|
||
self.samples = load_dataset("json", data_files=jsonl_path, split="train")
|
||
|
||
def __len__(self):
|
||
return len(self.samples)
|
||
|
||
def __getitem__(self, index):
|
||
sample = dict(self.samples[index])
|
||
messages = build_search_shortqa_messages(sample, self.max_contexts, self.max_chars_per_context)
|
||
prompt = self.tokenizer.apply_chat_template(
|
||
messages,
|
||
tokenize=False,
|
||
add_generation_prompt=True,
|
||
open_thinking=False,
|
||
)
|
||
return {
|
||
"prompt": prompt,
|
||
"question": str(sample.get("question") or sample.get("query") or ""),
|
||
"answers": normalize_answers(sample),
|
||
"gold_doc_ids": normalize_gold_doc_ids(sample),
|
||
"contexts": select_contexts(sample, self.max_contexts),
|
||
}
|