# -*- coding: utf8 -*-
#

from transformers import AutoTokenizer, BertTokenizerFast

tokenizer = AutoTokenizer.from_pretrained('hfl/chinese-electra-180g-small-discriminator', use_fast=True)
question = '南京天气怎么样'  # 7
context = '我今天早上站在阳台看天空，今天南京天气很好！'  # 22

tokenized_examples = tokenizer(
    question,  # 问题文本
    context,  # 篇章文本
    truncation="only_second",  # 截断只发生在第二部分，即篇章
    max_length=20,  # 设定最大长度为384
    # stride=5,  # 设定篇章切片步长为128
    return_overflowing_tokens=True,  # 返回超出最大长度的标记，将篇章切成多片
    return_offsets_mapping=True,  # 返回偏置信息，用于对齐答案位置
    padding="max_length",  # 按最大长度进行补齐
)

print(tokenized_examples)
input_ids = tokenized_examples['input_ids']
token_type_ids = tokenized_examples['token_type_ids']
attention_masks = tokenized_examples['attention_mask']
offset_mappings = tokenized_examples['offset_mapping']
overflow_to_sample_mapping = tokenized_examples['overflow_to_sample_mapping']
for index, _input_ids in enumerate(input_ids):
    print('input_ids -> ', tokenizer.convert_ids_to_tokens(_input_ids))
    print('token_type_ids -> ', token_type_ids[index])
    print('attention_masks -> ', attention_masks[index])
    print('offset_mappings -> ', offset_mappings[index])
    print('overflow_to_sample_mapping -> ', overflow_to_sample_mapping[index])