paddlepaddle--paddlenlp
252 行
7.9 KiB
Python
252 行
7.9 KiB
Python
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import collections
|
|
import json
|
|
import random
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
import paddle
|
|
from paddle.io import Dataset
|
|
|
|
from paddlenlp.transformers.bert.tokenizer import BertTokenizer
|
|
|
|
BiEncoderPassage = collections.namedtuple("BiEncoderPassage", ["text", "title"])
|
|
|
|
BiENcoderBatch = collections.namedtuple(
|
|
"BiEncoderInput",
|
|
[
|
|
"questions_ids",
|
|
"question_segments",
|
|
"context_ids",
|
|
"ctx_segments",
|
|
"is_positive",
|
|
"hard_negatives",
|
|
"encoder_type",
|
|
],
|
|
)
|
|
|
|
|
|
def normalize_question(question: str) -> str:
|
|
question = question.replace("’", "'")
|
|
return question
|
|
|
|
|
|
def normalize_passage(ctx_text: str):
|
|
ctx_text = ctx_text.replace("\n", " ").replace("’", "'")
|
|
if ctx_text.startswith('"'):
|
|
ctx_text = ctx_text[1:]
|
|
if ctx_text.endswith('"'):
|
|
ctx_text = ctx_text[:-1]
|
|
return ctx_text
|
|
|
|
|
|
class BiEncoderSample(object):
|
|
query: str
|
|
positive_passages: List[BiEncoderPassage]
|
|
negative_passages: List[BiEncoderPassage]
|
|
hard_negative_passages: List[BiEncoderPassage]
|
|
|
|
|
|
class NQdataSetForDPR(Dataset):
|
|
"""
|
|
class for managing dataset
|
|
"""
|
|
|
|
def __init__(self, dataPath, query_special_suffix=None):
|
|
super(NQdataSetForDPR, self).__init__()
|
|
self.data = self._read_json_data(dataPath)
|
|
self.tokenizer = BertTokenizer
|
|
self.query_special_suffix = query_special_suffix
|
|
self.new_data = []
|
|
for i in range(0, self.__len__()):
|
|
self.new_data.append(self.__getitem__(i))
|
|
|
|
def _read_json_data(self, dataPath):
|
|
results = []
|
|
with open(dataPath, "r", encoding="utf-8") as f:
|
|
print("Reading file %s" % dataPath)
|
|
data = json.load(f)
|
|
results.extend(data)
|
|
print("Aggregated data size: {}".format(len(results)))
|
|
return results
|
|
|
|
def __getitem__(self, index):
|
|
json_sample_data = self.data[index]
|
|
r = BiEncoderSample()
|
|
r.query = self._process_query(json_sample_data["question"])
|
|
|
|
positive_ctxs = json_sample_data["positive_ctxs"]
|
|
|
|
negative_ctxs = json_sample_data["negative_ctxs"] if "negative_ctxs" in json_sample_data else []
|
|
hard_negative_ctxs = json_sample_data["hard_negative_ctxs"] if "hard_negative_ctxs" in json_sample_data else []
|
|
|
|
for ctx in positive_ctxs + negative_ctxs + hard_negative_ctxs:
|
|
if "title" not in ctx:
|
|
ctx["title"] = None
|
|
|
|
def create_passage(ctx):
|
|
return BiEncoderPassage(normalize_passage(ctx["text"]), ctx["title"])
|
|
|
|
r.positive_passages = [create_passage(ctx) for ctx in positive_ctxs]
|
|
r.negative_passages = [create_passage(ctx) for ctx in negative_ctxs]
|
|
r.hard_negative_passages = [create_passage(ctx) for ctx in hard_negative_ctxs]
|
|
|
|
return r
|
|
|
|
def _process_query(self, query):
|
|
query = normalize_question(query)
|
|
|
|
if self.query_special_suffix and not query.endswith(self.query_special_suffix):
|
|
query += self.query_special_suffix
|
|
|
|
return query
|
|
|
|
def __len__(self):
|
|
return len(self.data)
|
|
|
|
|
|
class DataUtil:
|
|
"""
|
|
Class for working with datasets
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.tensorizer = BertTensorizer()
|
|
|
|
def create_biencoder_input(
|
|
self,
|
|
samples: List[BiEncoderSample],
|
|
inserted_title,
|
|
num_hard_negatives=0,
|
|
num_other_negatives=0,
|
|
shuffle=True,
|
|
shuffle_positives=False,
|
|
hard_neg_positives=False,
|
|
hard_neg_fallback=True,
|
|
query_token=None,
|
|
):
|
|
|
|
question_tensors = []
|
|
ctx_tensors = []
|
|
positive_ctx_indices = []
|
|
hard_neg_ctx_indices = []
|
|
|
|
for sample in samples:
|
|
|
|
if shuffle and shuffle_positives:
|
|
positive_ctxs = sample.positive_passages
|
|
positive_ctx = positive_ctxs[np.random.choice(len(positive_ctxs))]
|
|
else:
|
|
positive_ctx = sample.positive_passages[0]
|
|
|
|
neg_ctxs = sample.negative_passages
|
|
hard_neg_ctxs = sample.hard_negative_passages
|
|
question = sample.query
|
|
|
|
if shuffle:
|
|
random.shuffle(neg_ctxs)
|
|
random.shuffle(hard_neg_ctxs)
|
|
|
|
if hard_neg_fallback and len(hard_neg_ctxs) == 0:
|
|
hard_neg_ctxs = neg_ctxs[0:num_hard_negatives]
|
|
|
|
neg_ctxs = neg_ctxs[0:num_other_negatives]
|
|
hard_neg_ctxs = hard_neg_ctxs[0:num_hard_negatives]
|
|
|
|
all_ctxs = [positive_ctx] + neg_ctxs + hard_neg_ctxs
|
|
hard_negative_start_idx = 1
|
|
hard_negative_end_idx = 1 + len(hard_neg_ctxs)
|
|
|
|
current_ctxs_len = len(ctx_tensors)
|
|
|
|
sample_ctxs_tensors = [
|
|
self.tensorizer.text_to_tensor(ctx.text, title=ctx.title if (inserted_title and ctx.title) else None)
|
|
for ctx in all_ctxs
|
|
]
|
|
|
|
ctx_tensors.extend(sample_ctxs_tensors)
|
|
positive_ctx_indices.append(current_ctxs_len)
|
|
hard_neg_ctx_indices.append(
|
|
i
|
|
for i in range(
|
|
current_ctxs_len + hard_negative_start_idx,
|
|
current_ctxs_len + hard_negative_end_idx,
|
|
)
|
|
)
|
|
"""if query_token:
|
|
if query_token == "[START_END]":
|
|
query_span = _select_span
|
|
else:
|
|
question_tensors.append(self.tensorizer.text_to_tensor(" ".join([query_token, question])))
|
|
else:"""
|
|
|
|
question_tensors.append(self.tensorizer.text_to_tensor(question))
|
|
|
|
ctxs_tensor = paddle.concat([paddle.reshape(ctx, [1, -1]) for ctx in ctx_tensors], axis=0)
|
|
questions_tensor = paddle.concat([paddle.reshape(q, [1, -1]) for q in question_tensors], axis=0)
|
|
|
|
ctx_segments = paddle.zeros_like(ctxs_tensor)
|
|
question_segments = paddle.zeros_like(questions_tensor)
|
|
|
|
return BiENcoderBatch(
|
|
questions_tensor,
|
|
question_segments,
|
|
ctxs_tensor,
|
|
ctx_segments,
|
|
positive_ctx_indices,
|
|
hard_neg_ctx_indices,
|
|
"question",
|
|
)
|
|
|
|
|
|
class BertTensorizer:
|
|
def __init__(self, pad_to_max=True, max_length=256):
|
|
self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
|
|
self.max_length = max_length
|
|
self.pad_to_max = pad_to_max
|
|
|
|
def text_to_tensor(
|
|
self,
|
|
text: str,
|
|
title=None,
|
|
):
|
|
text = text.strip()
|
|
|
|
if title:
|
|
token_ids = self.tokenizer.encode(
|
|
text,
|
|
text_pair=title,
|
|
max_seq_len=self.max_length,
|
|
pad_to_max_seq_len=False,
|
|
truncation_strategy="longest_first",
|
|
)["input_ids"]
|
|
else:
|
|
token_ids = self.tokenizer.encode(
|
|
text,
|
|
max_seq_len=self.max_length,
|
|
pad_to_max_seq_len=False,
|
|
truncation_strategy="longest_first",
|
|
)["input_ids"]
|
|
|
|
seq_len = self.max_length
|
|
if self.pad_to_max and len(token_ids) < seq_len:
|
|
token_ids = token_ids + [self.tokenizer.pad_token_type_id] * (seq_len - len(token_ids))
|
|
if len(token_ids) >= seq_len:
|
|
token_ids = token_ids[0:seq_len]
|
|
token_ids[-1] = 102
|
|
|
|
return paddle.to_tensor(token_ids)
|