553 lines
20 KiB
Python
553 lines
20 KiB
Python
# Copyright (c) 2023 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 numpy as np
|
|
import paddle
|
|
|
|
from ..taskflow import Taskflow
|
|
from ..transformers import (
|
|
AutoModelForCausalLM,
|
|
AutoModelForConditionalGeneration,
|
|
AutoTokenizer,
|
|
)
|
|
|
|
__all__ = [
|
|
"SentenceGenerate",
|
|
"SentenceSummarize",
|
|
"SentenceBackTranslate",
|
|
"SentenceBackTranslateAPI",
|
|
"SentenceContinue",
|
|
]
|
|
|
|
|
|
class SentenceGenerate:
|
|
"""
|
|
SentenceGenerate is a sentence-level data augmentation strategy
|
|
that generates similar sentences according to the input sequence.
|
|
The strategy first generates several sentences, and then chooses
|
|
the top n similar sentences by the model.
|
|
|
|
Args:
|
|
model_name (str):
|
|
Model parameter name for generation task.
|
|
create_n (int):
|
|
Number of augmented sequences.
|
|
generate_n (int):
|
|
Number of generated sequences.
|
|
max_length (int):
|
|
The max length of the prediction.
|
|
top_p (float): The cumulative probability for
|
|
top-p-filtering in the "sampling" strategy. The value should
|
|
satisfy 0 <= top_p < 1. Default to 0.95.
|
|
"""
|
|
|
|
def __init__(
|
|
self, model_name="roformer-chinese-sim-char-base", create_n=1, generate_n=5, max_length=128, top_p=0.95
|
|
):
|
|
self.model_name = model_name
|
|
self.create_n = create_n
|
|
self.generate_n = generate_n
|
|
self.max_length = max_length
|
|
self.top_p = top_p
|
|
|
|
self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
|
|
self.model.eval()
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
|
|
|
|
def augment(self, sequences):
|
|
"""
|
|
Apply augmentation strategy on input sequences.
|
|
|
|
Args:
|
|
sequences (str or list(str)):
|
|
Input sequence or list of input sequences.
|
|
|
|
"""
|
|
if isinstance(sequences, str):
|
|
sequences = [sequences]
|
|
augmented_sequences = []
|
|
for sequence in sequences:
|
|
augmented_sequences.append(self._generate_similar_sentence(sequence, self.model, self.tokenizer))
|
|
return augmented_sequences
|
|
|
|
@paddle.no_grad()
|
|
def _generate_similar_sentence(self, sequence, model, tokenizer):
|
|
"""Generates generate_n similar sentences from the provided sequence, and choose the best create_n similar sentences."""
|
|
|
|
# Generate generate_n similar sentences
|
|
generated_sequences = [sequence]
|
|
tokenized_input = tokenizer(sequence, return_tensors="pd", padding=True)
|
|
decoded_outputs = tokenizer.batch_decode(
|
|
model.generate(
|
|
**tokenized_input,
|
|
num_return_sequences=self.generate_n,
|
|
top_p=self.top_p,
|
|
decode_strategy="sampling",
|
|
max_length=self.max_length,
|
|
)[0],
|
|
skip_special_tokens=True,
|
|
)
|
|
for decoded_output in decoded_outputs:
|
|
s = decoded_output.replace(" ", "").replace(sequence, "")
|
|
if s not in generated_sequences and len(s) > 0:
|
|
generated_sequences.append(s)
|
|
tokenized_output = tokenizer(generated_sequences, return_tensors="pd", padding=True)
|
|
|
|
# Choose best create_n similar sentences
|
|
tokenized_output = tokenizer(generated_sequences, return_tensors="pd", padding=True)
|
|
Z = model.roformer(**tokenized_output)[1].cpu().numpy()
|
|
Z /= (Z**2).sum(axis=1, keepdims=True) ** 0.5
|
|
|
|
return [generated_sequences[i + 1] for i in np.dot(Z[1:], -Z[0]).argsort()[: self.create_n]]
|
|
|
|
|
|
class SentenceSummarize:
|
|
"""
|
|
SentenceSummarize is a sentence-level data augmentation strategy
|
|
that summarizes the input sequence.
|
|
|
|
Args:
|
|
create_n (int):
|
|
Number of augmented sequences.
|
|
max_length (int):
|
|
The max length of the summarization.
|
|
batch_size(int):
|
|
The sample number of a mini-batch.
|
|
top_k (int): The number of highest probability tokens to
|
|
keep for top-k-filtering in the "sampling" strategy. Default to
|
|
0, which means no effect.
|
|
top_p (float): The cumulative probability for
|
|
top-p-filtering in the "sampling" strategy. The value should
|
|
satisfy 0 <= top_p < 1. Default to 1.0, which means no
|
|
effect.
|
|
temperature (float): The value used to module the next
|
|
token probabilities in the "sampling" strategy. Default to 1.0,
|
|
which means no effect.
|
|
use_fp16_decoding: (bool): Whether to use fp16 for decoding.
|
|
Only works when faster entry is available. Default to False.
|
|
kwargs (dict): Additional keyword arguments refer to ..taskflow.text_summarization.TextSummarization
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
create_n=1,
|
|
max_length=128,
|
|
batch_size=1,
|
|
top_k=5,
|
|
top_p=1.0,
|
|
temperature=1.0,
|
|
use_fp16_decoding=False,
|
|
**kwargs
|
|
):
|
|
|
|
kwargs.setdefault("num_return_sequences", create_n)
|
|
kwargs.setdefault("num_beams", create_n * 4)
|
|
kwargs.setdefault("max_length", max_length)
|
|
kwargs.setdefault("batch_size", batch_size)
|
|
kwargs.setdefault("top_k", top_k)
|
|
kwargs.setdefault("top_p", top_p)
|
|
kwargs.setdefault("temperature", temperature)
|
|
kwargs.setdefault("use_fp16_decoding", use_fp16_decoding)
|
|
|
|
self.create_n = kwargs["num_return_sequences"]
|
|
self.summarization = Taskflow("text_summarization", **kwargs)
|
|
|
|
def augment(self, sequences):
|
|
"""
|
|
Apply augmentation strategy on input sequences.
|
|
|
|
Args:
|
|
sequences (str or list(str)):
|
|
Input sequence or list of input sequences.
|
|
|
|
"""
|
|
if isinstance(sequences, str):
|
|
sequences = [sequences]
|
|
augmented_sequences = self.summarization(sequences)
|
|
return [augmented_sequences[i * self.create_n : (i + 1) * self.create_n] for i in range(len(sequences))]
|
|
|
|
|
|
class SentenceBackTranslate:
|
|
"""
|
|
SentenceBackTranslate is a sentence-level data augmentation strategy
|
|
that translates the input sequence into one language, and backtranslate
|
|
back into the source language by the language models.
|
|
|
|
Args:
|
|
src_lang (str):
|
|
The source language of the input sequences.
|
|
tgt_lang (str):
|
|
The target language of the translated sequences.
|
|
max_length (int):
|
|
The max length of the translation.
|
|
batch_size(int):
|
|
The sample number of a mini-batch.
|
|
num_beams (int): The number of beams in the "beam_search"
|
|
strategy. Default to 4.
|
|
use_faster: (bool): Whether to use faster entry of model
|
|
for FasterGeneration. Default to False (already deprecated).
|
|
decode_strategy (str, optional): The decoding strategy in generation.
|
|
Currently, there are three decoding strategies supported:
|
|
"greedy_search", "sampling" and "beam_search". Default to
|
|
"beam_search".
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
src_lang="zh",
|
|
tgt_lang="en",
|
|
max_length=128,
|
|
batch_size=1,
|
|
num_beams=4,
|
|
use_faster=False,
|
|
decode_strategy="beam_search",
|
|
from_model_name=None,
|
|
to_model_name=None,
|
|
):
|
|
self.src_lang = src_lang
|
|
self.tgt_lang = tgt_lang
|
|
self.max_length = max_length
|
|
self.batch_size = batch_size
|
|
self.num_beams = num_beams
|
|
self.decode_strategy = decode_strategy
|
|
self.from_model_name = from_model_name
|
|
self.to_model_name = to_model_name
|
|
self.MBART_MAP = {
|
|
"ar": "ar_AR",
|
|
"cs": "cs_CZ",
|
|
"de": "de_DE",
|
|
"en": "en_XX",
|
|
"es": "es_XX",
|
|
"et": "et_EE",
|
|
"fi": "fi_FI",
|
|
"fr": "fr_XX",
|
|
"gu": "gu_IN",
|
|
"hi": "hi_IN",
|
|
"it": "it_IT",
|
|
"ja": "ja_XX",
|
|
"kk": "kk_KZ",
|
|
"ko": "ko_KR",
|
|
"lt": "lt_LT",
|
|
"lv": "lv_LV",
|
|
"my": "my_MM",
|
|
"ne": "ne_NP",
|
|
"nl": "nl_XX",
|
|
"ro": "ro_RO",
|
|
"ru": "ru_RU",
|
|
"si": "si_LK",
|
|
"tr": "tr_TR",
|
|
"vi": "vi_VN",
|
|
"zh": "zh_CN",
|
|
"af": "af_ZA",
|
|
"az": "az_AZ",
|
|
"bn": "bn_IN",
|
|
"fa": "fa_IR",
|
|
"he": "he_IL",
|
|
"hr": "hr_HR",
|
|
"id": "id_ID",
|
|
"ka": "ka_GE",
|
|
"km": "km_KH",
|
|
"mk": "mk_MK",
|
|
"ml": "ml_IN",
|
|
"mn": "mn_MN",
|
|
"mr": "mr_IN",
|
|
"pl": "pl_PL",
|
|
"ps": "ps_AF",
|
|
"pt": "pt_XX",
|
|
"sv": "sv_SE",
|
|
"sw": "sw_KE",
|
|
"ta": "ta_IN",
|
|
"te": "te_IN",
|
|
"th": "th_TH",
|
|
"tl": "tl_XX",
|
|
"uk": "uk_UA",
|
|
"ur": "ur_PK",
|
|
"xh": "xh_ZA",
|
|
"gl": "gl_ES",
|
|
"sl": "sl_SI",
|
|
}
|
|
if self.from_model_name is None:
|
|
if tgt_lang == "en":
|
|
self.from_model_name = "mbart-large-50-many-to-one-mmt"
|
|
else:
|
|
self.from_model_name = "mbart-large-50-many-to-many-mmt"
|
|
|
|
if to_model_name is None:
|
|
if tgt_lang == "en":
|
|
self.to_model_name = "mbart-large-50-one-to-many-mmt"
|
|
else:
|
|
self.to_model_name = "mbart-large-50-many-to-many-mmt"
|
|
|
|
self.from_model = AutoModelForConditionalGeneration.from_pretrained(self.from_model_name)
|
|
self.to_model = AutoModelForConditionalGeneration.from_pretrained(self.to_model_name)
|
|
self.from_tokenizer = AutoTokenizer.from_pretrained(self.from_model_name, src_lang=self.MBART_MAP[src_lang])
|
|
self.to_tokenizer = AutoTokenizer.from_pretrained(self.to_model_name, src_lang=self.MBART_MAP[tgt_lang])
|
|
self.from_model.eval()
|
|
self.to_model.eval()
|
|
|
|
def augment(self, sequences):
|
|
"""
|
|
Apply augmentation strategy on input sequences.
|
|
|
|
Args:
|
|
sequences (str or list(str)):
|
|
Input sequence or list of input sequences.
|
|
|
|
"""
|
|
if isinstance(sequences, str):
|
|
sequences = [sequences]
|
|
sequences = self._translate(self.from_model, self.from_tokenizer, sequences, self.tgt_lang)
|
|
sequences = self._translate(self.to_model, self.to_tokenizer, sequences, self.src_lang)
|
|
return [[sequence] for sequence in sequences]
|
|
|
|
@paddle.no_grad()
|
|
def _translate(self, model, tokenizer, sequences, lang):
|
|
batched_inputs = [sequences[idx : idx + self.batch_size] for idx in range(0, len(sequences), self.batch_size)]
|
|
translated_texts = []
|
|
eos_id = model.mbart.config["eos_token_id"]
|
|
for batched_input in batched_inputs:
|
|
tokenized_input = tokenizer(batched_input, return_tensors="pd", padding=True)["input_ids"]
|
|
outputs = model.generate(
|
|
input_ids=tokenized_input,
|
|
forced_bos_token_id=tokenizer.lang_code_to_id[self.MBART_MAP[lang]],
|
|
decode_strategy=self.decode_strategy,
|
|
num_beams=self.num_beams,
|
|
max_length=self.max_length,
|
|
)[0]
|
|
for output in outputs:
|
|
eos = np.where(output.cpu().numpy() == eos_id)[0]
|
|
if len(eos) == 0:
|
|
eos_pos = len(output) - 1
|
|
else:
|
|
eos_pos = eos[0]
|
|
translated_texts.append(tokenizer.convert_ids_to_string(output[1:eos_pos]))
|
|
return translated_texts
|
|
|
|
|
|
class SentenceBackTranslateAPI:
|
|
"""
|
|
SentenceBackTranslateAPI is a sentence-level data augmentation strategy
|
|
that translates the input sequence into one language, and back-translate
|
|
back into the source language by baidu translate api.
|
|
|
|
Args:
|
|
src_lang (str):
|
|
The source language of the input sequences.
|
|
tgt_lang (str):
|
|
The target language of the translated sequences.
|
|
appid (str):
|
|
Appid for requesting Baidu translation service. (if use your own appid/appkey)
|
|
secretKey (str):
|
|
Secret key for requesting Baidu translation service. (if use your own appid/appkey)
|
|
qps (int):
|
|
Queries per second. (if use your own appid/appkey)
|
|
"""
|
|
|
|
def __init__(self, src_lang="zh", tgt_lang="en", appid=None, secretKey=None, qps=1):
|
|
|
|
self.src_lang = src_lang
|
|
self.tgt_lang = tgt_lang
|
|
self.appid = appid
|
|
self.secretKey = secretKey
|
|
self.qps = qps
|
|
self.url = "http://api.fanyi.baidu.com/api/trans/vip/translate"
|
|
|
|
def augment(self, sequences):
|
|
"""
|
|
Apply augmentation strategy on input sequences.
|
|
|
|
Args:
|
|
sequences (str or list(str)):
|
|
Input sequence or list of input sequences.
|
|
|
|
"""
|
|
if isinstance(sequences, str):
|
|
sequences = [sequences]
|
|
if self.appid is None or self.secretKey is None:
|
|
return self._back_translate_hub(sequences)
|
|
else:
|
|
return self._back_translate_api(sequences)
|
|
|
|
def _back_translate_hub(self, sequences):
|
|
try:
|
|
import paddlehub as hub
|
|
except ImportError:
|
|
print(" PaddleHub not installed!")
|
|
import os
|
|
|
|
os.system("pip install paddlehub==2.3.1")
|
|
import paddlehub as hub
|
|
|
|
module = hub.Module(name="baidu_translate")
|
|
translated_texts = []
|
|
for sequence in sequences:
|
|
sequence = module.translate(sequence, self.src_lang, self.tgt_lang)
|
|
sequence = module.translate(sequence, self.tgt_lang, self.src_lang)
|
|
translated_texts.append([sequence])
|
|
return translated_texts
|
|
|
|
def _back_translate_api(self, sequences):
|
|
|
|
translated_texts = []
|
|
for sequence in sequences:
|
|
sequence = self._translate_api(sequence, self.src_lang, self.tgt_lang)
|
|
sequence = self._translate_api(sequence, self.tgt_lang, self.src_lang)
|
|
translated_texts.append(sequence)
|
|
return translated_texts
|
|
|
|
def _translate_api(self, query, from_lang, to_lang):
|
|
|
|
import hashlib
|
|
import random
|
|
import time
|
|
|
|
import requests
|
|
|
|
# Generate salt and sign
|
|
salt = str(random.randint(32768, 65536))
|
|
sign = self.appid + query + salt + self.secretKey
|
|
sign = hashlib.md5(sign.encode("utf-8")).hexdigest()
|
|
|
|
# Build request
|
|
headers = {"Content-Type": "application/x-www-form-urlencoded"}
|
|
payload = {
|
|
"appid": f"{self.appid}",
|
|
"q": f"{query}",
|
|
"from": from_lang,
|
|
"to": to_lang,
|
|
"salt": f"{salt}",
|
|
"sign": f"{sign}",
|
|
}
|
|
|
|
# Send request
|
|
time.sleep(1 / self.qps)
|
|
try:
|
|
r = requests.post(self.url, params=payload, headers=headers)
|
|
result = r.json()
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
raise RuntimeError(error_msg)
|
|
if "error_code" in result:
|
|
raise RuntimeError(result)
|
|
return result["trans_result"][0]["dst"]
|
|
|
|
|
|
class SentenceContinue:
|
|
"""
|
|
SentenceContinue is a sentence-level data augmentation strategy
|
|
that generates continuation for the input sequence.
|
|
|
|
Args:
|
|
model_name (str):
|
|
Model parameter name for summarization task.
|
|
max_length (int):
|
|
The max length of the summarization.
|
|
decode_strategy (str, optional): The decoding strategy in generation.
|
|
Currently, there are three decoding strategies supported:
|
|
"greedy_search", "sampling" and "beam_search". Default to
|
|
"beam_search".
|
|
use_faster: (bool): Whether to use faster entry of model
|
|
for FasterGeneration. Default to False (already deprecated).
|
|
create_n (int):
|
|
Number of augmented sequences.
|
|
batch_size(int):
|
|
The sample number of a mini-batch.
|
|
top_k (int): The number of highest probability tokens to
|
|
keep for top-k-filtering in the "sampling" strategy. Default to
|
|
0, which means no effect.
|
|
top_p (float): The cumulative probability for
|
|
top-p-filtering in the "sampling" strategy. The value should
|
|
satisfy 0 <= top_p < 1. Default to 1.0, which means no
|
|
effect.
|
|
temperature (float): The value used to module the next
|
|
token probabilities in the "sampling" strategy. Default to 1.0,
|
|
which means no effect.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_name="gpt-cpm-small-cn-distill",
|
|
max_length=64,
|
|
decode_strategy="sampling",
|
|
use_faster=False,
|
|
create_n=1,
|
|
top_k=50,
|
|
temperature=1.0,
|
|
top_p=0.9,
|
|
batch_size=1,
|
|
):
|
|
self.model_name = model_name
|
|
self.max_length = max_length
|
|
self.decode_strategy = decode_strategy
|
|
self.create_n = create_n
|
|
self.top_k = top_k
|
|
self.temperature = temperature
|
|
self.top_p = top_p
|
|
self.batch_size = batch_size
|
|
|
|
self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
|
|
self.model.eval()
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
|
|
self.tokenizer.add_special_tokens(
|
|
{"pad_token": self.tokenizer.convert_ids_to_tokens(self.model.config.pad_token_id)}
|
|
)
|
|
|
|
def augment(self, sequences):
|
|
"""
|
|
Apply augmentation strategy on input sequences.
|
|
|
|
Args:
|
|
sequences (str or list(str)):
|
|
Input sequence or list of input sequences.
|
|
|
|
"""
|
|
if isinstance(sequences, str):
|
|
sequences = [sequences]
|
|
return self._generate_continue(sequences, self.model, self.tokenizer)
|
|
|
|
@paddle.no_grad()
|
|
def _generate_continue(self, sequences, model, tokenizer):
|
|
batched_inputs = [sequences[idx : idx + self.batch_size] for idx in range(0, len(sequences), self.batch_size)]
|
|
generated_sequences = []
|
|
for batched_input in batched_inputs:
|
|
tokenized_inputs = tokenizer(
|
|
batched_input, return_tensors="pd", padding=True, return_attention_mask=True, return_position_ids=True
|
|
)
|
|
outputs = model.generate(
|
|
**tokenized_inputs,
|
|
max_length=self.max_length,
|
|
decode_strategy=self.decode_strategy,
|
|
num_return_sequences=self.create_n,
|
|
top_k=self.top_k,
|
|
temperature=self.temperature,
|
|
top_p=self.top_p,
|
|
)[0]
|
|
for i in range(outputs.shape[0]):
|
|
output = outputs[i].cpu().numpy()
|
|
eos = np.where(output == model.config.eos_token_id)[0]
|
|
if len(eos) == 0:
|
|
eos_pos = len(output) - 1
|
|
else:
|
|
eos_pos = eos[0]
|
|
generated_sequences.append(tokenizer.convert_ids_to_string(output[:eos_pos].tolist()))
|
|
augmented_sequences = []
|
|
for i, sequence in enumerate(sequences):
|
|
augmented_sequence = []
|
|
for ii in range(self.create_n):
|
|
continue_sequence = (
|
|
generated_sequences[i * self.create_n + ii].replace(" ", "").replace("\n", "").replace("\t", "")
|
|
)
|
|
augmented_sequence.append(sequence + continue_sequence)
|
|
augmented_sequences.append(augmented_sequence)
|
|
return augmented_sequences
|