【案例】新媒体评论分类与信息抽取 · 数据工程

{"context": "Instruction: 现在你是一个非常厉害的SPO抽取器。\nInput: 下面这句中包含了哪些三元组，用json列表的形式回答，不要输出除json外的其他答案。\n\n73获奖记录人物评价：黄磊是一个特别幸运的演员，拍第一部戏就碰到了导演陈凯歌，而且在他的下一部电影《夜半歌声》中演对手戏的张国荣、吴倩莲、黎明等都是著名的港台演员。\nAnswer: ", "target": "```json\n[{\"subject\": \"夜半歌声\",\"subject_type\": \"影视作品\",\"predicate\": \"主演\",\"object\": \"黄磊\",\"object_type\": \"人物\"}]\n```"}
{"context": "Instruction: 现在你是一个非常厉害的SPO抽取器。\nInput: 抽取下面这句话中的SPO，只返回json的答案。\n\n马红宝 男 汉族，1949年8月生，浙江省长兴县人，1978年8月加入中国共产党，1967年9月参加工作，大普文化\nAnswer: ", "target": "```json\n[{\"predicate\": \"民族\", \"object_type\": \"Text\", \"subject_type\": \"人物\", \"object\": \"汉族\", \"subject\": \"马红宝\"}, {\"predicate\": \"国籍\", \"object_type\": \"国家\", \"subject_type\": \"人物\", \"object\": \"中国\", \"subject\": \"马红宝\"}, {\"predicate\": \"出生地\", \"object_type\": \"地点\", \"subject_type\": \"人物\", \"object\": \"浙江省长兴\", \"subject\": \"马红宝\"}, {\"predicate\": \"出生日期\", \"object_type\": \"Date\", \"subject_type\": \"人物\", \"object\": \"1949年8月\", \"subject\": \"马红宝\"}]\n```"}
{"context": "Instruction: 现在你是一个非常厉害的SPO抽取器。\nInput: 前秦世祖宣昭皇帝苻坚（338年－385年10月16日），字永固，又字文玉，小名坚头，氐族，略阳临渭（今甘肃秦安）人，十六国时期前秦的君主，公元357－385年在位。\n\n上面这句话中包含了哪些三元组信息？。\nAnswer: ", "target": "这句话中包含了前秦世祖宣昭皇帝苻坚的生日、字、小名、出生地、民族等信息，以上信息用json可以表示为：\n```json\n[{\"predicate\": \"出生日期\", \"object_type\": \"Date\", \"subject_type\": \"人物\", \"object\": \"338年\", \"subject\": \"苻坚\"}, {\"predicate\": \"字\", \"object_type\": \"Text\", \"subject_type\": \"历史人物\", \"object\": \"文玉\", \"subject\": \"苻坚\"}, {\"predicate\": \"出生地\", \"object_type\": \"地点\", \"subject_type\": \"人物\", \"object\": \"略阳临渭\", \"subject\": \"苻坚\"}, {\"predicate\": \"民族\", \"object_type\": \"Text\", \"subject_type\": \"人物\", \"object\": \"氐族\", \"subject\": \"苻坚\"}]\n```"}
{"context": "Instruction: 现在你是一个非常厉害的SPO抽取器。\nInput: “《大明按察使》该剧由黄海涛，申积军，阙云霞，沈剑波，王江红担任策划、黄克敏执导，姚橹，丁勇岱，杨旸，高鑫，李芯逸等主演的悬疑断案电视剧。”中包含了哪些SPO信息？用json的格式回答，不要说多余的话。\nAnswer: ", "target": "```json\n[{\"predicate\": \"导演\", \"object_type\": \"人物\", \"subject_type\": \"影视作品\", \"object\": \"黄克敏\", \"subject\": \"大明按察使\"}, {\"predicate\": \"主演\", \"object_type\": \"人物\", \"subject_type\": \"影视作品\", \"object\": \"丁勇岱\", \"subject\": \"大明按察使\"}, {\"predicate\": \"主演\", \"object_type\": \"人物\", \"subject_type\": \"影视作品\", \"object\": \"姚橹\", \"subject\": \"大明按察使\"}, {\"predicate\": \"主演\", \"object_type\": \"人物\", \"subject_type\": \"影视作品\", \"object\": \"杨旸\", \"subject\": \"大明按察使\"}]\n```"}
{"context": "Instruction: 现在你是一个非常厉害的SPO抽取器。\nInput: 《中国姓氏大辞典》是在国家自然科学基金的支持下，于2010年由江西人民出版社正式编纂出版的图书，作者是袁义达、邱家儒。\n\n抽取出这句话中包含了哪些关系，并用json的形式回答，不要输出多余信息。\nAnswer: ", "target": "```json\n[{\"predicate\": \"出版社\", \"object_type\": \"出版社\", \"subject_type\": \"书籍\", \"object\": \"江西人民出版社\", \"subject\": \"中国姓氏大辞典\"}, {\"predicate\": \"作者\", \"object_type\": \"人物\", \"subject_type\": \"图书作品\", \"object\": \"袁义达\", \"subject\": \"中国姓氏大辞典\"}, {\"predicate\": \"作者\", \"object_type\": \"人物\", \"subject_type\": \"图书作品\", \"object\": \"邱家儒\", \"subject\": \"中国姓氏大辞典\"}]\n```"}

# -*- coding:utf-8 -*-
"""数据长什么样：一条 jsonl 拆开看，纯标准库，可直接运行。

两个任务（文本分类、SPO 抽取）共用同一套字段：context 与 target。
context 里塞的是「Instruction + Input + Answer:」三段式，
Answer: 后面故意留空——那是留给模型写的位置。
"""
import json

CLS_SAMPLE = {
    "context": "Instruction: 你现在是一个很厉害的阅读理解器，严格按照人类指令进行回答。\n"
               "Input: 下面句子中的主语是什么类别，输出成列表形式。\n\n"
               "朋友说京东上买的水果很新鲜，我就也下单了一箱\nAnswer: ",
    "target": "[\"水果\"]"
}

SPO_SAMPLE = {
    "context": "Instruction: 你现在是一个很厉害的阅读理解器，严格按照人类指令进行回答。\n"
               "Input: 找到句子中的三元组信息并输出成json给我：\n\n"
               "《霸王别姬》是由张国荣主演的一部电影\nAnswer: ",
    "target": "```json\n[{\"predicate\": \"主演\", \"object_type\": \"人物\", "
              "\"subject_type\": \"影视作品\", \"object\": \"张国荣\", \"subject\": \"霸王别姬\"}]\n```"
}


def show(name, sample):
    print('=' * 60)
    print(name)
    print('=' * 60)
    context, target = sample['context'], sample['target']
    for i, seg in enumerate(context.split('\n')):
        print('  ctx[%d] %s' % (i, seg if seg else '(空行)'))
    print('  target: %s' % target)
    print('  context %d 字 / target %d 字' % (len(context), len(target)))


if __name__ == '__main__':
    show('任务一：文本分类', CLS_SAMPLE)
    show('任务二：SPO 信息抽取', SPO_SAMPLE)

    # 两条样本的字段名完全一样——这是「换任务不改代码」的前提
    assert set(CLS_SAMPLE) == set(SPO_SAMPLE) == {'context', 'target'}
    print('\n两个任务的字段结构完全一致:', sorted(CLS_SAMPLE))

    # context 一定以 'Answer: ' 收尾，推理时按这个标记切结果
    for s in (CLS_SAMPLE, SPO_SAMPLE):
        assert s['context'].rstrip(' ').endswith('Answer:'), 'Answer: 是切分锚点，不能省'
    print("context 都以 'Answer: ' 结尾，推理时按它切出模型写的部分")

    # target 必须是合法 JSON（SPO 的要先剥掉代码围栏）
    print('\n分类 target 解析:', json.loads(CLS_SAMPLE['target']))
    spo = SPO_SAMPLE['target'].strip().removeprefix('```json').removesuffix('```').strip()
    parsed = json.loads(spo)
    print('抽取 target 解析:', parsed[0]['subject'], '-', parsed[0]['predicate'], '-', parsed[0]['object'])
    assert parsed[0]['subject'] == '霸王别姬'

    # 写回一行 jsonl 看看真实存储形态
    line = json.dumps(CLS_SAMPLE, ensure_ascii=False)
    print('\n存盘时就是这么一行:')
    print(' ', line[:100] + ' ...')
    assert '\\n' in line, 'jsonl 里换行是转义的 \\n，一条样本永远占一行'
    print('\n一条样本 = 一行 JSON，换行被转义，所以可以按行流式读')

# -*- coding:utf-8 -*-
"""Instruction 是怎么拼出来的：把三段式做成函数，纯标准库，可直接运行。

手写 context 字符串很容易在空格、换行、冒号上出入，
而训练与推理只要有一个字不一致，模型就像换了一场考试。
所以把它固化成一个函数，两边都调它。
"""
import json

SYSTEM = '你现在是一个很厉害的阅读理解器，严格按照人类指令进行回答。'

TASK_PROMPT = {
    'classify': '下面句子中的主语是什么类别，输出成列表形式。',
    'spo': '找到句子中的三元组信息并输出成json给我：',
    'ner': '找到文章中所有【{entity_type}】类型的实体，并按照JSON格式输出。',
}


def build_context(task: str, sentence: str, **kwargs) -> str:
    """拼出 context。注意三个细节，缺一个就和训练数据对不上：
    1) Input 与正文之间是两个换行；
    2) Answer 后面有一个空格；
    3) 结尾不加换行。
    """
    if task not in TASK_PROMPT:
        raise ValueError('未知任务 %s，可选: %s' % (task, list(TASK_PROMPT)))
    instruction = TASK_PROMPT[task].format(**kwargs) if kwargs else TASK_PROMPT[task]
    return ('Instruction: %s\nInput: %s\n\n%s\nAnswer: '
            % (SYSTEM, instruction, sentence))


def build_sample(task: str, sentence: str, target, **kwargs) -> dict:
    """训练样本 = context + target。target 是对象时序列化成紧凑 JSON。"""
    if not isinstance(target, str):
        target = json.dumps(target, ensure_ascii=False)
    return {'context': build_context(task, sentence, **kwargs), 'target': target}


if __name__ == '__main__':
    ctx = build_context('classify', '朋友说京东上买的水果很新鲜，我就也下单了一箱')
    print(repr(ctx))

    # 训练与推理必须走同一个函数，结果逐字相同
    assert build_context('classify', 'x') == build_context('classify', 'x')
    assert ctx.endswith('Answer: '), 'Answer: 后面那个空格不能丢'
    assert '\n\n' in ctx, 'Input 与正文之间是两个换行'

    print('\n三个任务各生成一条:')
    samples = [
        build_sample('classify', '这款平板续航很顶', ['平板']),
        build_sample('spo', '《霸王别姬》是由张国荣主演的一部电影',
                     [{'predicate': '主演', 'subject': '霸王别姬', 'object': '张国荣'}]),
        build_sample('ner', '北京冬奥会在张家口设了赛区', ['北京', '张家口'], entity_type='地点'),
    ]
    for s in samples:
        print(' ', json.dumps(s, ensure_ascii=False)[:96], '...')

    # 换任务只换 instruction 文本，格式骨架一模一样
    heads = [s['context'].split('\n')[0] for s in samples]
    assert len(set(heads)) == 1, '三个任务的第一行完全相同'
    print('\n三条样本的 Instruction 行完全相同，差别只在 Input 行')

    # 带占位符的任务能格式化进去
    assert '【地点】' in samples[2]['context']
    print('ner 任务的实体类型被填进了模板:',
          samples[2]['context'].split('\n')[1])

# -*- coding:utf-8 -*-
"""labels 是怎么涂成 -100 的：纯标准库复刻，可直接运行。

这一段是整个生成式微调最容易写错的地方。
写错了不会报错，只会让模型把「提示词」也当成要背下来的内容，
表现是：推理时它把你的 Instruction 一字不差地复读一遍。
"""

GMASK, SOP, EOP, PAD = 150001, 150004, 150005, 3
IGNORE = -100


def build_inputs(prompt_ids, target_ids):
    """复刻 ChatGLM 的 build_inputs_with_special_tokens。"""
    return prompt_ids + [GMASK, SOP] + target_ids + [EOP]


def make_labels(input_ids, bos_token_id=SOP):
    """提示部分涂 -100，答案部分照抄 input_ids。"""
    context_length = input_ids.index(bos_token_id)
    mask_position = context_length - 1
    return [IGNORE] * context_length + input_ids[mask_position + 1:]


def pad_to(input_ids, labels, max_len):
    pad_len = max_len - len(input_ids)
    return (input_ids + [PAD] * pad_len,
            labels + [IGNORE] * pad_len)   # 标签补 -100，不是补 PAD


if __name__ == '__main__':
    prompt_ids = [5001, 5002, 5003, 5004]      # "Instruction: ... Answer: "
    target_ids = [7001, 7002]                  # '["水果"]'

    input_ids = build_inputs(prompt_ids, target_ids)
    print('input_ids:', input_ids)
    assert input_ids == [5001, 5002, 5003, 5004, GMASK, SOP, 7001, 7002, EOP]

    labels = make_labels(input_ids)
    print('labels   :', labels)
    assert len(labels) == len(input_ids), 'labels 必须和 input_ids 等长'

    # 关键性质一：<sop> 及之前全部是 -100
    context_length = input_ids.index(SOP)
    print('\ncontext_length（<sop> 的下标）:', context_length)
    assert labels[:context_length] == [IGNORE] * context_length
    print('前 %d 个位置不计损失' % context_length)

    # 关键性质二：答案段与 input_ids 逐元素相同
    assert labels[context_length:] == input_ids[context_length:]
    print('从下标 %d 起，labels 与 input_ids 完全一致:' % context_length,
          labels[context_length:])

    # 关键性质三：<eop> 要参与损失，模型才学得会「写完就停」
    assert labels[-1] == EOP, '结束符必须计损失，否则模型不知道何时停下'
    print('结束符 <eop> 参与损失 -> 模型学得会停笔')

    # padding 之后再验一次
    ids_p, labels_p = pad_to(input_ids, labels, 16)
    print('\npad 到 16:')
    print('  input_ids:', ids_p)
    print('  labels   :', labels_p)
    assert labels_p[len(input_ids):] == [IGNORE] * 7
    assert IGNORE not in ids_p, 'input_ids 里不能出现 -100'

    # 对照：labels 补 PAD 的错误写法
    wrong = labels + [PAD] * (16 - len(labels))
    diff = sum(1 for a, b in zip(labels_p, wrong) if a != b)
    print('\n标签补 PAD（错误）会让 %d 个位置变成「要模型学着输出 padding」' % diff)
    assert diff == 7

    # 统计参与损失的 token 占比 —— 数据质量的一个硬指标
    valid = sum(1 for t in labels_p if t != IGNORE)
    print('\n参与损失的 token: %d / %d = %.1f%%' % (valid, len(labels_p), 100 * valid / len(labels_p)))
    print('这个比例太低（比如 5%）说明提示太长、答案太短，训练信号被稀释')

# -*- coding:utf-8 -*-
"""labels 是怎么涂成 -100 的：纯标准库复刻，可直接运行。

这一段是整个生成式微调最容易写错的地方。
写错了不会报错，只会让模型把「提示词」也当成要背下来的内容，
表现是：推理时它把你的 Instruction 一字不差地复读一遍。
"""

GMASK, SOP, EOP, PAD = 150001, 150004, 150005, 3
IGNORE = -100


def build_inputs(prompt_ids, target_ids):
    """复刻 ChatGLM 的 build_inputs_with_special_tokens。"""
    return prompt_ids + [GMASK, SOP] + target_ids + [EOP]


def make_labels(input_ids, bos_token_id=SOP):
    """提示部分涂 -100，答案部分照抄 input_ids。"""
    context_length = input_ids.index(bos_token_id)
    mask_position = context_length - 1
    return [IGNORE] * context_length + input_ids[mask_position + 1:]


def pad_to(input_ids, labels, max_len):
    pad_len = max_len - len(input_ids)
    return (input_ids + [PAD] * pad_len,
            labels + [IGNORE] * pad_len)   # 标签补 -100，不是补 PAD


if __name__ == '__main__':
    prompt_ids = [5001, 5002, 5003, 5004]      # "Instruction: ... Answer: "
    target_ids = [7001, 7002]                  # '["水果"]'

    input_ids = build_inputs(prompt_ids, target_ids)
    print('input_ids:', input_ids)
    assert input_ids == [5001, 5002, 5003, 5004, GMASK, SOP, 7001, 7002, EOP]

    labels = make_labels(input_ids)
    print('labels   :', labels)
    assert len(labels) == len(input_ids), 'labels 必须和 input_ids 等长'

    # 关键性质一：<sop> 及之前全部是 -100
    context_length = input_ids.index(SOP)
    print('\ncontext_length（<sop> 的下标）:', context_length)
    assert labels[:context_length] == [IGNORE] * context_length
    print('前 %d 个位置不计损失' % context_length)

    # 关键性质二：答案段与 input_ids 逐元素相同
    assert labels[context_length:] == input_ids[context_length:]
    print('从下标 %d 起，labels 与 input_ids 完全一致:' % context_length,
          labels[context_length:])

    # 关键性质三：<eop> 要参与损失，模型才学得会「写完就停」
    assert labels[-1] == EOP, '结束符必须计损失，否则模型不知道何时停下'
    print('结束符 <eop> 参与损失 -> 模型学得会停笔')

    # padding 之后再验一次
    ids_p, labels_p = pad_to(input_ids, labels, 16)
    print('\npad 到 16:')
    print('  input_ids:', ids_p)
    print('  labels   :', labels_p)
    assert labels_p[len(input_ids):] == [IGNORE] * 7
    assert IGNORE not in ids_p, 'input_ids 里不能出现 -100'

    # 对照：labels 补 PAD 的错误写法
    wrong = labels + [PAD] * (16 - len(labels))
    diff = sum(1 for a, b in zip(labels_p, wrong) if a != b)
    print('\n标签补 PAD（错误）会让 %d 个位置变成「要模型学着输出 padding」' % diff)
    assert diff == 7

    # 统计参与损失的 token 占比 —— 数据质量的一个硬指标
    valid = sum(1 for t in labels_p if t != IGNORE)
    print('\n参与损失的 token: %d / %d = %.1f%%' % (valid, len(labels_p), 100 * valid / len(labels_p)))
    print('这个比例太低（比如 5%）说明提示太长、答案太短，训练信号被稀释')

# -*- coding:utf-8 -*-
"""把一条 {context, target} 变成 ChatGLM 能吃的 input_ids / labels。

生成式微调与 BERT 分类最大的认知差：
输入和标签是**同一条序列**，只是标签把「提示部分」全部涂成 -100，
让损失只落在模型自己要写出来的那一段上。
"""
import numpy as np


def convert_example(examples: dict,
                    tokenizer,
                    max_source_seq_len: int,
                    max_target_seq_len: int):
    """
    参数:
        examples: {'text': ['{"context": "...", "target": "..."}', ...]}，每行一条 jsonl
        max_source_seq_len: 提示部分最长 token 数
        max_target_seq_len: 回答部分最长 token 数
    返回:
        {'input_ids': [...], 'labels': [...]}，两者等长
    """
    import json

    tokenized_output = {'input_ids': [], 'labels': []}
    max_seq_length = max_source_seq_len + max_target_seq_len

    for example in examples['text']:
        try:
            example = json.loads(example)
            context = example['context']
            target = example['target']

            # prompt 部分：不加特殊 token，后面手动拼
            prompts_ids = tokenizer.encode(text=context, add_special_tokens=False)
            target_ids = tokenizer.encode(text=target, add_special_tokens=False)

            # 超长就截断。截 prompt 会丢掉 Instruction 的尾部，
            # 截 target 会让模型学到「答案写一半就停」——两种都要在数据阶段先统计好
            if len(prompts_ids) >= max_source_seq_len:
                prompts_ids = prompts_ids[:max_source_seq_len - 1]
            if len(target_ids) >= max_target_seq_len - 1:
                target_ids = target_ids[:max_target_seq_len - 2]

            # build_inputs_with_special_tokens 会补上 [gMASK] 与 <sop>，
            # 并在结尾补 <eop>，这是 GLM 系列的生成格式，别自己手拼
            input_ids = tokenizer.build_inputs_with_special_tokens(prompts_ids, target_ids)

            # <sop> 的下标 = 提示部分的长度。labels 从这里开始才是真答案
            context_length = input_ids.index(tokenizer.bos_token_id)
            mask_position = context_length - 1

            labels = [-100] * context_length + input_ids[mask_position + 1:]

            pad_len = max_seq_length - len(input_ids)
            input_ids = input_ids + [tokenizer.pad_token_id] * pad_len
            # 标签这边补的是 -100 而不是 pad_token_id：
            # 补 pad_token_id 会让模型努力去学「输出一串 padding」
            labels = labels + [-100] * pad_len

            tokenized_output['input_ids'].append(input_ids)
            tokenized_output['labels'].append(labels)
        except Exception as e:
            print('[跳过] 解析失败:', e)
            continue

    for k, v in tokenized_output.items():
        tokenized_output[k] = np.array(v)
    return tokenized_output


if __name__ == '__main__':
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained('./ChatGLM-6B', trust_remote_code=True)
    out = convert_example(
        {'text': ['{"context": "Instruction: 你是阅读理解器\\nAnswer: ", "target": "[\\"水果\\"]"}']},
        tokenizer, max_source_seq_len=100, max_target_seq_len=100)
    print('input_ids 形状:', out['input_ids'].shape)
    print('labels    形状:', out['labels'].shape)
    print('labels 里 -100 的个数:', int((out['labels'] == -100).sum()))

# -*- coding:utf-8 -*-
"""ChatGLM + LoRA 的工程配置。

和 BERT 路线最大的差别：这里不再有「类别数」这种概念。
任务是什么，全靠 instruction 文本说清楚；换任务不改代码，只改数据。
"""
import os

import torch


class ProjectConfig(object):
    def __init__(self):
        self.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
        self.root = os.environ.get(
            'GLM_PROJECT_ROOT',
            os.path.dirname(os.path.abspath(__file__)),
        )
        # 底座：ChatGLM-6B，62 亿参数，GLM 架构，28 层，hidden_size 4096
        self.pre_model = os.path.join(self.root, 'ChatGLM-6B')
        self.train_path = os.path.join(self.root, 'data', 'mixed_train_dataset.jsonl')
        self.dev_path = os.path.join(self.root, 'data', 'mixed_dev_dataset.jsonl')

        # LoRA 三件套
        self.use_lora = True
        self.lora_rank = 8          # 秩：A 是 (hidden, 8)，B 是 (8, out)
        self.lora_alpha = 32        # 缩放系数，实际缩放 = alpha / rank = 4
        self.lora_dropout = 0.1

        self.use_ptuning = False    # 与 LoRA 二选一，不要同时开
        self.pre_seq_len = 128
        self.prefix_projection = False

        # 显存敏感项：batch 只能给到 2，靠梯度累积把等效 batch 撑起来
        self.batch_size = 2
        self.gradient_accumulation_steps = 8
        self.epochs = 2
        self.learning_rate = 1e-4   # 比全参微调的 5e-5 大一个量级：只训 0.06% 的参数，学慢了学不动
        self.weight_decay = 0
        self.warmup_ratio = 0.06

        # 输入输出各留 100 个 token。超了直接截断，不会报错
        self.max_source_seq_len = 100
        self.max_target_seq_len = 100

        self.logging_steps = 10
        self.save_freq = 1000
        self.save_dir = os.path.join(self.root, 'checkpoints', 'ptune')


if __name__ == '__main__':
    pc = ProjectConfig()
    print('设备          :', pc.device)
    print('等效 batch    :', pc.batch_size * pc.gradient_accumulation_steps)
    print('LoRA 缩放系数 :', pc.lora_alpha / pc.lora_rank)
    print('单条最长 token:', pc.max_source_seq_len + pc.max_target_seq_len)

# -*- coding:utf-8 -*-
"""开训前必做：统计有多少条样本会被 max_source/target_seq_len 截断。

max_source_seq_len=100 是配置里写死的默认值，不是普适真理。
新媒体评论普遍短，但 SPO 抽取任务的 target 是一串 JSON，
很容易超过 100 个 token —— 被截断的样本会教模型「写一半就交卷」。
"""
import json

MAX_SOURCE, MAX_TARGET = 100, 100


def rough_token_len(text: str) -> int:
    """粗估 token 数：中文按字算，ASCII 连续段按 4 个字符 1 token 估。

    真实统计请换成 tokenizer.encode(text, add_special_tokens=False) 的长度；
    这里用粗估是为了在没有模型权重的机器上也能先跑一遍体检。
    """
    n, ascii_run = 0, 0
    for ch in text:
        if ord(ch) < 128:
            ascii_run += 1
        else:
            n += 1 + (ascii_run + 3) // 4
            ascii_run = 0
    return n + (ascii_run + 3) // 4


def audit(samples):
    over_src = over_tgt = 0
    src_lens, tgt_lens = [], []
    for s in samples:
        sl, tl = rough_token_len(s['context']), rough_token_len(s['target'])
        src_lens.append(sl)
        tgt_lens.append(tl)
        over_src += sl >= MAX_SOURCE
        over_tgt += tl >= MAX_TARGET - 1
    return {
        'total': len(samples),
        'src_max': max(src_lens), 'src_avg': sum(src_lens) / len(src_lens),
        'tgt_max': max(tgt_lens), 'tgt_avg': sum(tgt_lens) / len(tgt_lens),
        'over_src': over_src, 'over_tgt': over_tgt,
    }


if __name__ == '__main__':
    head = ('Instruction: 你现在是一个很厉害的阅读理解器，严格按照人类指令进行回答。\n'
            'Input: 找到句子中的三元组信息并输出成json给我：\n\n')
    samples = [
        {'context': head + '这款平板续航很顶\nAnswer: ', 'target': '["平板"]'},
        {'context': head + '《霸王别姬》是由张国荣主演的一部电影\nAnswer: ',
         'target': '```json\n[{"predicate": "主演", "object_type": "人物", '
                   '"subject_type": "影视作品", "object": "张国荣", "subject": "霸王别姬"}]\n```'},
        {'context': head + '《活着》是余华在1992年发表的长篇小说，讲述了福贵一生的故事\nAnswer: ',
         'target': '```json\n[{"predicate": "作者", "object_type": "人物", '
                   '"subject_type": "图书作品", "object": "余华", "subject": "活着"}, '
                   '{"predicate": "出版时间", "object_type": "时间", '
                   '"subject_type": "图书作品", "object": "1992年", "subject": "活着"}]\n```'},
    ]

    stat = audit(samples)
    print('样本数        :', stat['total'])
    print('source 平均/最长: %.1f / %d  (上限 %d)' % (stat['src_avg'], stat['src_max'], MAX_SOURCE))
    print('target 平均/最长: %.1f / %d  (上限 %d)' % (stat['tgt_avg'], stat['tgt_max'], MAX_TARGET))
    print('会被截断的 source: %d 条' % stat['over_src'])
    print('会被截断的 target: %d 条' % stat['over_tgt'])

    # 光是这条固定的 Instruction 头就吃掉多少预算
    head_len = rough_token_len(head)
    print('\n固定 Instruction 头占 %d 个 token，正文只剩 %d 个'
          % (head_len, MAX_SOURCE - head_len))
    assert head_len > 40, 'Instruction 头本身就很长，这是新手最容易忽略的一笔开销'

    # 两条 SPO 的样本，target 比分类长一个数量级
    cls_t = rough_token_len(samples[0]['target'])
    spo_t = rough_token_len(samples[2]['target'])
    print('分类 target %d token vs 双三元组 target %d token，差 %.1f 倍'
          % (cls_t, spo_t, spo_t / cls_t))
    assert spo_t > cls_t * 5
    print('\n结论：同一份 max_target_seq_len 同时服务两个任务时，按长的那个任务定')

# -*- coding:utf-8 -*-
"""生成式微调的 DataLoader。

比 BERT 路线简单：没有 token_type_ids，没有 attention_mask —— 
ChatGLM 的 forward 会根据 input_ids 里的 <sop> 位置自己构造注意力掩码。
我们只负责给对 input_ids 与 labels。
"""
from functools import partial

from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, default_data_collator

from data_handle import convert_example
from glm_config import ProjectConfig

pc = ProjectConfig()
# trust_remote_code=True：ChatGLM 的分词器实现放在模型仓库里，不在 transformers 本体
tokenizer = AutoTokenizer.from_pretrained(pc.pre_model, trust_remote_code=True)


def get_data():
    dataset = load_dataset('text', data_files={'train': pc.train_path,
                                               'dev': pc.dev_path})
    new_func = partial(convert_example,
                       tokenizer=tokenizer,
                       max_source_seq_len=pc.max_source_seq_len,
                       max_target_seq_len=pc.max_target_seq_len)

    # batched=True：一次处理一批，比逐条快很多。
    # 注意 convert_example 内部是按条循环的，这里的 batch 只是为了减少函数调用开销
    dataset = dataset.map(new_func, batched=True)

    train_dataloader = DataLoader(dataset['train'],
                                  shuffle=True,
                                  collate_fn=default_data_collator,
                                  batch_size=pc.batch_size)
    dev_dataloader = DataLoader(dataset['dev'],
                                collate_fn=default_data_collator,
                                batch_size=pc.batch_size)
    return train_dataloader, dev_dataloader


if __name__ == '__main__':
    train_dataloader, dev_dataloader = get_data()
    print('train batch 数:', len(train_dataloader))
    print('dev   batch 数:', len(dev_dataloader))
    for batch in train_dataloader:
        print('字段:', list(batch.keys()))
        print('形状:', {k: tuple(v.shape) for k, v in batch.items()})
        # 训练开始前最值得看的一眼：这一批里有多少 token 真的参与了损失
        valid = int((batch['labels'] != -100).sum())
        total = batch['labels'].numel()
        print('参与损失的 token: %d / %d = %.1f%%' % (valid, total, 100 * valid / total))
        break

# -*- coding:utf-8 -*-
"""骨架模板：把你自己的标注数据转成这套 jsonl。

拿到一份业务数据（Excel 导出、爬虫结果、标注平台导出），
到能开训之间就差这一个脚本。四个 TODO 填完即可。
"""
import json
import random

# TODO-1: 系统角色描述。定了之后训练与推理都不能再改
SYSTEM = '你现在是一个很厉害的阅读理解器，严格按照人类指令进行回答。'
# TODO-2: 每个任务一句话指令
INSTRUCTIONS = {
    'classify': '下面句子中的主语是什么类别，输出成列表形式。',
    'spo': '找到句子中的三元组信息并输出成json给我：',
}
# TODO-3: 验证集比例
DEV_RATIO = 0.1
SEED = 42


def build_context(task: str, sentence: str) -> str:
    return ('Instruction: %s\nInput: %s\n\n%s\nAnswer: '
            % (SYSTEM, INSTRUCTIONS[task], sentence))


def build_target(task: str, answer) -> str:
    """分类返回裸 JSON 列表；抽取包一层 ```json 围栏，与原项目保持一致。"""
    body = json.dumps(answer, ensure_ascii=False)
    if task == 'spo':
        return '```json\n%s\n```' % body
    return body


def iter_raw():
    """TODO-4: 换成你自己的数据源，yield (task, sentence, answer)。"""
    yield 'classify', '朋友说京东上买的水果很新鲜，我就也下单了一箱', ['水果']
    yield 'classify', '这款平板的续航是真顶，出差两天没充电', ['平板']
    yield 'spo', '《霸王别姬》是由张国荣主演的一部电影', [
        {'predicate': '主演', 'subject': '霸王别姬', 'object': '张国荣'}]


def main(train_path='mixed_train_dataset.jsonl', dev_path='mixed_dev_dataset.jsonl'):
    samples = []
    for task, sentence, answer in iter_raw():
        if task not in INSTRUCTIONS:
            print('[跳过] 未知任务', task)
            continue
        samples.append({'context': build_context(task, sentence),
                        'target': build_target(task, answer)})

    # 固定随机种子再打散：两个任务混在一起训练时，顺序必须是乱的，
    # 否则前半程全是分类、后半程全是抽取，模型会在后半程「忘掉」前一个任务
    random.Random(SEED).shuffle(samples)

    n_dev = max(1, int(len(samples) * DEV_RATIO))
    dev, train = samples[:n_dev], samples[n_dev:]

    for path, rows in ((train_path, train), (dev_path, dev)):
        with open(path, 'w', encoding='utf8') as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + '\n')
        print('写出 %s: %d 条' % (path, len(rows)))
    return len(train), len(dev)


if __name__ == '__main__':
    import os
    import tempfile

    tmp = tempfile.mkdtemp()
    n_train, n_dev = main(os.path.join(tmp, 'train.jsonl'), os.path.join(tmp, 'dev.jsonl'))
    assert n_train + n_dev == 3

    # 回读校验：每一行都必须是合法 JSON 且字段齐全
    with open(os.path.join(tmp, 'train.jsonl'), encoding='utf8') as f:
        for line in f:
            row = json.loads(line)
            assert set(row) == {'context', 'target'}
            assert row['context'].endswith('Answer: ')
    print('回读校验通过：每行都是合法 JSON，字段齐全，context 以 Answer: 收尾')

# -*- coding:utf-8 -*-
"""骨架模板：把你自己的标注数据转成这套 jsonl。

拿到一份业务数据（Excel 导出、爬虫结果、标注平台导出），
到能开训之间就差这一个脚本。四个 TODO 填完即可。
"""
import json
import random

# TODO-1: 系统角色描述。定了之后训练与推理都不能再改
SYSTEM = '你现在是一个很厉害的阅读理解器，严格按照人类指令进行回答。'
# TODO-2: 每个任务一句话指令
INSTRUCTIONS = {
    'classify': '下面句子中的主语是什么类别，输出成列表形式。',
    'spo': '找到句子中的三元组信息并输出成json给我：',
}
# TODO-3: 验证集比例
DEV_RATIO = 0.1
SEED = 42


def build_context(task: str, sentence: str) -> str:
    return ('Instruction: %s\nInput: %s\n\n%s\nAnswer: '
            % (SYSTEM, INSTRUCTIONS[task], sentence))


def build_target(task: str, answer) -> str:
    """分类返回裸 JSON 列表；抽取包一层 ```json 围栏，与原项目保持一致。"""
    body = json.dumps(answer, ensure_ascii=False)
    if task == 'spo':
        return '```json\n%s\n```' % body
    return body


def iter_raw():
    """TODO-4: 换成你自己的数据源，yield (task, sentence, answer)。"""
    yield 'classify', '朋友说京东上买的水果很新鲜，我就也下单了一箱', ['水果']
    yield 'classify', '这款平板的续航是真顶，出差两天没充电', ['平板']
    yield 'spo', '《霸王别姬》是由张国荣主演的一部电影', [
        {'predicate': '主演', 'subject': '霸王别姬', 'object': '张国荣'}]


def main(train_path='mixed_train_dataset.jsonl', dev_path='mixed_dev_dataset.jsonl'):
    samples = []
    for task, sentence, answer in iter_raw():
        if task not in INSTRUCTIONS:
            print('[跳过] 未知任务', task)
            continue
        samples.append({'context': build_context(task, sentence),
                        'target': build_target(task, answer)})

    # 固定随机种子再打散：两个任务混在一起训练时，顺序必须是乱的，
    # 否则前半程全是分类、后半程全是抽取，模型会在后半程「忘掉」前一个任务
    random.Random(SEED).shuffle(samples)

    n_dev = max(1, int(len(samples) * DEV_RATIO))
    dev, train = samples[:n_dev], samples[n_dev:]

    for path, rows in ((train_path, train), (dev_path, dev)):
        with open(path, 'w', encoding='utf8') as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + '\n')
        print('写出 %s: %d 条' % (path, len(rows)))
    return len(train), len(dev)


if __name__ == '__main__':
    import os
    import tempfile

    tmp = tempfile.mkdtemp()
    n_train, n_dev = main(os.path.join(tmp, 'train.jsonl'), os.path.join(tmp, 'dev.jsonl'))
    assert n_train + n_dev == 3

    # 回读校验：每一行都必须是合法 JSON 且字段齐全
    with open(os.path.join(tmp, 'train.jsonl'), encoding='utf8') as f:
        for line in f:
            row = json.loads(line)
            assert set(row) == {'context', 'target'}
            assert row['context'].endswith('Answer: ')
    print('回读校验通过：每行都是合法 JSON，字段齐全，context 以 Answer: 收尾')