wehub-resource-sync
257 lines
9.1 KiB
Python
257 lines
9.1 KiB
Python
# Copyright (c) 2024 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.
|
|
"""Embedding dataset."""
|
|
|
|
import random
|
|
from dataclasses import dataclass
|
|
from typing import List
|
|
|
|
from paddle.io import Dataset, IterableDataset
|
|
|
|
from ..utils.log import logger
|
|
|
|
|
|
@dataclass
|
|
class Example:
|
|
"""Dataset example."""
|
|
|
|
query: str
|
|
pos_passage: List[str]
|
|
neg_passage: List[str] = None
|
|
|
|
|
|
@dataclass
|
|
class Sequence:
|
|
"""Sequence."""
|
|
|
|
token_ids: List[int]
|
|
position_ids: List[int]
|
|
|
|
|
|
@dataclass
|
|
class Pair:
|
|
"""Pair."""
|
|
|
|
query: Sequence
|
|
passages: List[Sequence]
|
|
|
|
|
|
class EmbeddingDatasetMixin:
|
|
"""EmbeddingDatasetMixin."""
|
|
|
|
def convert_example(tokenizer, example):
|
|
"""Convert raw json format example to Example."""
|
|
|
|
assert all(
|
|
(key in example for key in ["query", "pos_passage", "neg_passage"])
|
|
), "query, pos_passage, neg_passage are needed"
|
|
|
|
if not isinstance(example["query"], str):
|
|
raise ValueError("query must be a string.")
|
|
if isinstance(example["pos_passage"], str):
|
|
example["pos_passage"] = [example["pos_passage"]]
|
|
if isinstance(example["neg_passage"], str):
|
|
example["neg_passage"] = [example["neg_passage"]]
|
|
|
|
if len(example["neg_passage"]) > 0:
|
|
for item in [example["query"]] + example["pos_passage"] + example["neg_passage"]:
|
|
if not isinstance(item, str):
|
|
raise ValueError("The item in pos_passage / neg_passage must be a string.")
|
|
if len(item.strip()) == 0:
|
|
raise ValueError("Example with empty string in query / pos_passage / neg_passage field.")
|
|
|
|
query = example["query"]
|
|
pos_passage = example["pos_passage"]
|
|
neg_passage = example["neg_passage"]
|
|
return Example(query=query, pos_passage=pos_passage, neg_passage=neg_passage)
|
|
|
|
def tokenize_template(cls, tokenizer, template: str):
|
|
"""Tokenize a given template using the provided tokenizer."""
|
|
assert template.count("{text}") == 1, "Template must contain exactly one {text} placeholder"
|
|
|
|
template_prefix, template_suffix = template.split("{text}")
|
|
|
|
prefix_tokens = tokenizer(template_prefix, add_special_tokens=False).input_ids
|
|
suffix_tokens = tokenizer(template_suffix, add_special_tokens=False).input_ids
|
|
return prefix_tokens, suffix_tokens
|
|
|
|
def _process_truncation(self, tokens, text_type):
|
|
"""
|
|
Process tokens by converting them into a complete token sequence with prefix and suffix,
|
|
and generate corresponding position ids.
|
|
"""
|
|
if text_type not in ["query", "passage"]:
|
|
raise ValueError("text_type must be either 'query' or 'passage'")
|
|
|
|
prefix_key = f"{text_type}_template_prefix"
|
|
suffix_key = f"{text_type}_template_suffix"
|
|
max_len_key = f"max_{text_type}_len"
|
|
|
|
# If the template does not contain a suffix token, add the EOS token to the end
|
|
if getattr(self, suffix_key) == []:
|
|
setattr(self, suffix_key, [self.tokenizer.eos_token_id])
|
|
if getattr(self, prefix_key) == []:
|
|
setattr(self, prefix_key, [self.tokenizer.bos_token_id])
|
|
|
|
# Calculate the available length
|
|
max_len = getattr(self, max_len_key)
|
|
prefix_tokens = getattr(self, prefix_key)
|
|
suffix_tokens = getattr(self, suffix_key)
|
|
available_len = int(max_len - len(prefix_tokens) - len(suffix_tokens))
|
|
|
|
# Convert tokens to ids and truncate
|
|
token_ids_converted = self.tokenizer.convert_tokens_to_ids(tokens)
|
|
truncated_token_ids = token_ids_converted[:available_len]
|
|
|
|
# Combine prefix, truncated tokens, and suffix
|
|
token_ids = prefix_tokens + truncated_token_ids + suffix_tokens
|
|
pos_ids = list(range(len(token_ids)))
|
|
return token_ids, pos_ids
|
|
|
|
def _postprocess_sequence(self, example: Example, rng):
|
|
"""Post process sequence: tokenization & truncation."""
|
|
query = example.query
|
|
pos_passage = rng.choice(example.pos_passage)
|
|
neg_passage = example.neg_passage
|
|
if len(neg_passage) > 0:
|
|
if len(neg_passage) < self.group_size - 1:
|
|
# Calculate how many full sets are needed to ensure each element appears at least once
|
|
full_sets_needed = (self.group_size - 1) // len(neg_passage)
|
|
remainder = (self.group_size - 1) % len(neg_passage)
|
|
|
|
# Initialize the list and add complete sets
|
|
selected_neg_passage = neg_passage * full_sets_needed
|
|
|
|
# Ensure the remainder part is filled; randomly select from neg_passage
|
|
selected_neg_passage += rng.sample(neg_passage, remainder)
|
|
|
|
# Shuffle the result to ensure randomness
|
|
rng.shuffle(selected_neg_passage)
|
|
else:
|
|
selected_neg_passage = rng.sample(neg_passage, self.group_size - 1)
|
|
else:
|
|
selected_neg_passage = []
|
|
# Process query tokens
|
|
query_tokens = self.tokenizer.tokenize(query)
|
|
query_token_ids, query_pos_ids = self._process_truncation(query_tokens, "query")
|
|
|
|
query = Sequence(
|
|
token_ids=query_token_ids,
|
|
position_ids=query_pos_ids,
|
|
)
|
|
|
|
# Process passage tokens
|
|
passages = []
|
|
for passage in [pos_passage] + selected_neg_passage:
|
|
passage_tokens = self.tokenizer.tokenize(passage)
|
|
passage_token_ids, passage_pos_ids = self._process_truncation(passage_tokens, "passage")
|
|
passages.append(
|
|
Sequence(
|
|
token_ids=passage_token_ids,
|
|
position_ids=passage_pos_ids,
|
|
)
|
|
)
|
|
return Pair(query=query, passages=passages)
|
|
|
|
|
|
class EmbeddingDataset(EmbeddingDatasetMixin, Dataset):
|
|
def __init__(
|
|
self,
|
|
dataset,
|
|
tokenizer,
|
|
max_query_len: int = 64,
|
|
max_passage_len: int = 256,
|
|
group_size: int = 2,
|
|
query_template: str = "{text}",
|
|
passage_template: str = "{text}",
|
|
):
|
|
super().__init__()
|
|
self.example_dataset = dataset
|
|
self.tokenizer = tokenizer
|
|
self.max_query_len = max_query_len
|
|
self.max_passage_len = max_passage_len
|
|
self.group_size = group_size
|
|
self.query_template = query_template
|
|
self.passage_template = passage_template
|
|
self.query_template_prefix, self.query_template_suffix = self.tokenize_template(
|
|
self.tokenizer, self.query_template
|
|
)
|
|
self.passage_template_prefix, self.passage_template_suffix = self.tokenize_template(
|
|
self.tokenizer, self.passage_template
|
|
)
|
|
|
|
for index, data in enumerate(self.example_dataset):
|
|
self.example_dataset[index] = self.convert_example(data)
|
|
|
|
def __getitem__(self, index):
|
|
return self._postprocess_sequence(self.example_dataset[index])
|
|
|
|
def __len__(self):
|
|
raise len(self.example_dataset)
|
|
|
|
|
|
class EmbeddingIterableDataset(EmbeddingDatasetMixin, IterableDataset):
|
|
"""Create sequences from Example Dataset.
|
|
|
|
This is a stateful dataset.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataset,
|
|
tokenizer,
|
|
max_query_len: int = 64,
|
|
max_passage_len: int = 256,
|
|
group_size: int = 2,
|
|
query_template: str = "{text}",
|
|
passage_template: str = "{text}",
|
|
):
|
|
super().__init__()
|
|
self.example_dataset = dataset
|
|
self.tokenizer = tokenizer
|
|
self.max_query_len = max_query_len
|
|
self.max_passage_len = max_passage_len
|
|
self.group_size = group_size
|
|
self.query_template = query_template
|
|
self.passage_template = passage_template
|
|
self.query_template_prefix, self.query_template_suffix = self.tokenize_template(
|
|
self.tokenizer, self.query_template
|
|
)
|
|
self.passage_template_prefix, self.passage_template_suffix = self.tokenize_template(
|
|
self.tokenizer, self.passage_template
|
|
)
|
|
|
|
self.epoch_index = 0
|
|
|
|
def __iter__(self):
|
|
while True:
|
|
logger.info(f"Start to load dataset on epoch={self.epoch_index}")
|
|
yield from self.iter_one_epoch()
|
|
|
|
def iter_one_epoch(self):
|
|
"""Iterates through one epoch of the dataset."""
|
|
|
|
num_sequences = 0
|
|
rng = random.Random()
|
|
for _, example in enumerate(self.example_dataset):
|
|
example = self.convert_example(example)
|
|
rng.seed(num_sequences)
|
|
sequence = self._postprocess_sequence(example, rng)
|
|
if sequence is None:
|
|
continue
|
|
num_sequences += 1
|
|
yield [sequence]
|
|
|
|
self.epoch_index += 1
|