Efficient-Large-Model/VILA15-13b-hf-preview

1# Copyright 2024 NVIDIA CORPORATION & AFFILIATES2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14#15# SPDX-License-Identifier: Apache-2.016 17from typing import Any, Dict, List, Optional, Sequence18 19import torch20import transformers21 22from .constants import IGNORE_INDEX, SENTINEL_TOKEN23from .conversation import SeparatorStyle, default_conversation24from .mm_utils import tokenizer_image_token25 26# __all__ = [27#     "tokenize_conversation",28#     "preprocess_conversation",29#     "infer_stop_tokens",30# ]31 32DUMMY_CONVERSATION = [33    {"from": "human", "value": "question"},34    {"from": "gpt", "value": "answer"},35] * 1036 37 38def tokenize_conversation_legacy(39    messages: Sequence[Dict[str, str]],40    tokenizer: transformers.PreTrainedTokenizer,41    add_generation_prompt: bool = False,42    overrides: Optional[Dict[str, str]] = None,43    no_system_prompt: bool = False,44) -> torch.Tensor:45    conv = default_conversation.copy()46    roles = {"human": conv.roles[0], "gpt": conv.roles[1]}47 48    if no_system_prompt:49        conv.system = ""50 51    # Skip the first message if it is not from human52    if messages[0]["from"] != "human":53        messages = messages[1:]54 55    # Add a generation prompt if needed56    if add_generation_prompt:57        messages.append({"from": "gpt", "value": None})58 59    conv.messages = []60    for turn, message in enumerate(messages):61        role = roles[message["from"]]62        assert role == conv.roles[turn % 2]63        if overrides is not None and message["from"] in overrides:64            conv.append_message(role, overrides[message["from"]])65        else:66            conv.append_message(role, message["value"])67 68    return tokenizer_image_token(conv.get_prompt(), tokenizer, return_tensors="pt")69 70 71def tokenize_conversation(72    messages: Sequence[Dict[str, str]],73    tokenizer: transformers.PreTrainedTokenizer,74    add_generation_prompt: bool = False,75    overrides: Optional[Dict[str, str]] = None,76    no_system_prompt: bool = False,77) -> torch.Tensor:78    # Normalize the conversation before tokenization79    for message in messages:80        message["value"] = message["value"].strip()81 82    if default_conversation.sep_style != SeparatorStyle.AUTO:83        return tokenize_conversation_legacy(84            messages,85            tokenizer,86            add_generation_prompt=add_generation_prompt,87            overrides=overrides,88            no_system_prompt=no_system_prompt,89        )90 91    conversation = []92    for m in messages:93        message = {}94        if m["from"] == "human":95            message["role"] = "user"96        elif m["from"] == "gpt":97            message["role"] = "assistant"98        else:99            raise ValueError(f"Unexpected sender '{m['from']}' in conversation entry.")100 101        message["content"] = m["value"]102        if overrides is not None and m["from"] in overrides:103            message["content"] = overrides[m["from"]]104        conversation.append(message)105 106    if no_system_prompt:107        conversation = [{"role": "system", "content": ""}] + conversation108 109    text = tokenizer.apply_chat_template(110        conversation,111        add_generation_prompt=add_generation_prompt,112        tokenize=False,113    )114    return tokenizer_image_token(text, tokenizer, return_tensors="pt")115 116 117def _maybe_add_sentinel_token(tokenizer: transformers.PreTrainedTokenizer) -> None:118    if not hasattr(tokenizer, "sentinel_token"):119        tokenizer.add_tokens([SENTINEL_TOKEN], special_tokens=True)120        tokenizer.sentinel_token = SENTINEL_TOKEN121        tokenizer.sentinel_token_id = tokenizer.convert_tokens_to_ids(SENTINEL_TOKEN)122 123 124def preprocess_conversation(125    conversation: Sequence[Dict[str, str]],126    tokenizer: transformers.PreTrainedTokenizer,127    no_system_prompt: bool = False,128    retried: bool = False,129) -> Dict[str, Any]:130    inputs = tokenize_conversation(conversation, tokenizer, no_system_prompt=no_system_prompt)131    labels = torch.ones_like(inputs) * IGNORE_INDEX132 133    # Generate the template by replacing the assistant's response with a sentinel.134    _maybe_add_sentinel_token(tokenizer)135    template = tokenize_conversation(136        conversation, tokenizer, overrides={"gpt": SENTINEL_TOKEN}, no_system_prompt=no_system_prompt137    )138 139    # Remove sentinel tokens from the template.140    mask = torch.ones_like(template, dtype=torch.bool)141    for k in range(template.size(0) - 1):142        if template[k] == tokenizer.sentinel_token_id:143            mask[k : k + 2] = False144            # NOTE(zhijianl): This is to handle the corner case where there is an empty token before the sentinel token.145            if k > 0 and retried:146                mask[k - 1] = False147    template = template[mask]148 149    # Match the tokenized conversation with the template (with no assistant's response).150    # Every token that is not matched will be included in the label for training.151    p = 0152    for k in range(inputs.size(0)):153        if p < template.size(0) and inputs[k] == template[p]:154            p += 1155        else:156            labels[k] = inputs[k]157 158    # Mask all tokens in the label if the template is not fully matched.159    if p < template.size(0):160        if not retried:161            return preprocess_conversation(162                conversation,163                tokenizer,164                no_system_prompt=no_system_prompt,165                retried=True,166            )167        print(f"Failed to process the conversation: '{conversation}'. All tokens will be masked in the label.")168        labels[:] = IGNORE_INDEX169 170    return {"input_ids": inputs, "labels": labels}171 172 173def infer_stop_tokens(tokenizer: transformers.PreTrainedTokenizer) -> List[str]:174    _maybe_add_sentinel_token(tokenizer)175    template = tokenize_conversation(DUMMY_CONVERSATION, tokenizer, overrides={"gpt": SENTINEL_TOKEN})176 177    stop_tokens = {tokenizer.eos_token}178    for k in range(template.size(0) - 1):179        if template[k] == tokenizer.sentinel_token_id:180            stop_token = tokenizer.decode(template[k + 1])181            stop_tokens.add(stop_token)182    return list(stop_tokens)183