"""训练和评估共用的部分：数据读取、提示格式。训练和推理必须用同一种提示格式。"""
import json
from pathlib import Path

HERE = Path(__file__).resolve().parent
INSTRUCTION = '把下面这段改写成我的写法，意思不变，只输出改写后的这一段：\n\n'


def read_pairs(split, data_dir=None):
    """读 data/pairs_<split>.jsonl。每行至少有 ai（默认输出）和 original（你改过的版本）两个字段。"""
    path = Path(data_dir or HERE / 'data') / f'pairs_{split}.jsonl'
    return [json.loads(l) for l in path.read_text(encoding='utf-8').splitlines() if l.strip()]


def user_message(ai_text):
    return INSTRUCTION + ai_text


def to_chat(pair):
    return {'messages': [{'role': 'user', 'content': user_message(pair['ai'])},
                         {'role': 'assistant', 'content': pair['original']}]}


def fewshot_messages(ai_text, examples):
    """基线：不训练，把几组“默认输出 → 改过的版本”直接写进上下文。"""
    msgs = []
    for ex in examples:
        msgs += [{'role': 'user', 'content': user_message(ex['ai'])}, {'role': 'assistant', 'content': ex['original']}]
    return msgs + [{'role': 'user', 'content': user_message(ai_text)}]