|
from onmt.constants import DefaultTokens |
|
from onmt.transforms import register_transform |
|
from onmt.utils.logging import logger |
|
from .transform import Transform |
|
|
|
|
|
@register_transform(name="insert_mask_before_placeholder") |
|
class InsertMaskBeforePlaceholdersTransform(Transform): |
|
"""Add the `DefaultTokens.MASK_BEFORE` placeholder between |
|
the prompt and the response in an LM finetuning exemple. |
|
This is necessary to enable the 'zero-out prompt loss' mechanism. |
|
""" |
|
|
|
def __init__(self, opts): |
|
super().__init__(opts) |
|
|
|
@classmethod |
|
def add_options(cls, parser): |
|
"""Options for mask_before placeholders insertion""" |
|
|
|
group = parser.add_argument_group( |
|
"Transform/InsertMaskBeforePlaceholdersTransform" |
|
) |
|
group.add( |
|
"--response_pattern", |
|
"-response_pattern", |
|
type=str, |
|
help="Response patten to locate the end of the prompt", |
|
default="Response : ｟newline｠", |
|
) |
|
|
|
def _parse_opts(self): |
|
self.response_pattern = self.opts.response_pattern |
|
|
|
def apply(self, example, is_train=False, stats=None, **kwargs): |
|
_src = " ".join(example["src"]) |
|
if len(_src.split(self.response_pattern)) != 2: |
|
logger.info("The mask_before could not be inserted") |
|
return example |
|
prompt, response = _src.split(self.response_pattern) |
|
response = DefaultTokens.MASK_BEFORE.join([self.response_pattern, response]) |
|
_src = "".join([prompt, response]) |
|
example["src"] = _src.split(" ") |
|
example["tgt"] = _src.split(" ") |
|
return example |
|
|
