47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
from dataset.search_shortqa_dataset import build_sft_conversation
|
|
|
|
|
|
def iter_jsonl(path):
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
yield json.loads(line)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Build SFT data for SearchShortQA")
|
|
parser.add_argument("--input", required=True, help="Raw JSONL with question/answer/contexts fields")
|
|
parser.add_argument("--output", required=True, help="Output SFT JSONL path")
|
|
parser.add_argument("--max_contexts", type=int, default=6)
|
|
parser.add_argument("--max_chars_per_context", type=int, default=700)
|
|
parser.add_argument("--limit", type=int, default=0)
|
|
args = parser.parse_args()
|
|
|
|
os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
|
|
count = 0
|
|
with open(args.output, "w", encoding="utf-8") as out:
|
|
for sample in iter_jsonl(args.input):
|
|
if args.limit and count >= args.limit:
|
|
break
|
|
conversations = build_sft_conversation(
|
|
sample,
|
|
max_contexts=args.max_contexts,
|
|
max_chars_per_context=args.max_chars_per_context,
|
|
)
|
|
out.write(json.dumps({"conversations": conversations}, ensure_ascii=False) + "\n")
|
|
count += 1
|
|
|
|
print(f"Wrote {count} SFT samples to {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|