|
from dataclasses import dataclass |
|
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union |
|
import copy |
|
|
|
from .formatter import EmptyFormatter, StringFormatter |
|
from .base import Template |
|
from .formatter import Formatter |
|
from ...utils.constants import * |
|
from . import register_template |
|
|
|
from transformers import PreTrainedTokenizer |
|
import torch |
|
|
|
|
|
|
|
@register_template('pretrain') |
|
@dataclass |
|
class PretrainTemplate(Template): |
|
format_image_token: "Formatter" = EmptyFormatter(slot="") |
|
format_user: "Formatter" = EmptyFormatter(slot="<image>") |
|
format_assistant: "Formatter" = StringFormatter(slot="{{content}}\n") |
|
system: "Formatter" = EmptyFormatter(slot="") |
|
separator: "Formatter" = EmptyFormatter(slot=['', '']) |
|
|
|
def make_labels(self, input_ids, prompt, tokenizer): |
|
labels = copy.deepcopy(input_ids) |
|
mask_len = len(self.tokenizer_image_token("<image>", tokenizer)) |
|
labels[:mask_len] = IGNORE_INDEX |
|
return labels |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
