MobileVLM と OpenVINO によるモバイル言語アシスタント#

%pip install -q "torch>=2.1.0" "timm>=0.9.12" --extra-index-url "https://download.pytorch.org/whl/cpu" 
%pip install -q "transformers>=4.33.1,<4.35.0" accelerate "sentencepiece>=0.1.99" "openvino>=2023.2.0" "nncf>=2.7.0" ipywidgets numpy "gradio>=4.19"

from pathlib import Path 

import sys 
MOBILEVLM_REPO_DIR = Path("./MobileVLM") 
if not MOBILEVLM_REPO_DIR.exists():
     !git clone -q "https://github.com/Meituan-AutoML/MobileVLM.git" 
sys.path.insert(0, str(MOBILEVLM_REPO_DIR))

import warnings 
import itertools 
import gc 
from typing import Optional, List, Tuple 

from mobilevlm.model.mobilevlm import load_pretrained_model 
from mobilevlm.conversation import conv_templates, SeparatorStyle 
from mobilevlm.utils import ( 
    disable_torch_init, 
    process_images, 
    tokenizer_image_token, 
    KeywordsStoppingCriteria, 
) 
from mobilevlm.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN 
import PIL 
import torch 
import transformers 
import numpy as np 
import gradio as gr 
import openvino as ov 
import nncf 
import ipywidgets as widgets

MODELS_DIR = Path("./models") 
MODEL_PATH = "mtgv/MobileVLM-1.7B" 

TEMPERATURE = 0.2 
TOP_P = None 
NUM_BEAMS = 1 
MAX_NEW_TOKENS = 512 

IMAGE_PATH = MOBILEVLM_REPO_DIR / "assets" / "samples" / "demo.jpg" 
PROMPT_STR = "Who is the author of this book?\nAnswer the question using a single word or phrase."

model_name = MODEL_PATH.split("/")[-1] 
disable_torch_init() 
with warnings.catch_warnings(): 
    warnings.simplefilter("ignore") 
    tokenizer, model, image_processor, _ = load_pretrained_model(MODEL_PATH, device="cpu") 
model = model.to(dtype=torch.float32)

def cleanup_torchscript_cache(): 
    """ 
    Helper for removing cached model representation 
    """ 
    torch._C._jit_clear_class_registry() 
    torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore() 
    torch.jit._state._clear_class_state()

compression_mode = widgets.Dropdown( 
nncf.compress_weightsoptions=["INT4", "INT8"], 
nncf.compress_weightsvalue="INT4", 
nncf.compress_weightsdescription="Compression mode:", 
nncf.compress_weightsdisabled=False, 
) 

compression_mode

stage1_xml_path = MODELS_DIR / f"stage1_{compression_mode.value}.xml" 
stage2_xml_path = MODELS_DIR / f"stage2_{compression_mode.value}.xml"

if compression_mode.value == "INT4": 
    wc_parameters = dict(mode=nncf.CompressWeightsMode.INT4_ASYM, group_size=128, ratio=0.8) 
else: 
    wc_parameters = dict(mode=nncf.CompressWeightsMode.INT8)

class ModelWrapper(torch.nn. Module): 
    def __init__(self, model): 
        super().__init__() 
        self.model = model 

    def forward( 
        self, 
        input_ids: torch.LongTensor = None, 
        attention_mask: Optional[torch.Tensor] = None, 
        past_key_values: Optional[List[torch.FloatTensor]] = None, 
        inputs_embeds: Optional[torch.FloatTensor] = None, 
    ): 
        outputs = self.model.model( 
            input_ids=input_ids, 
            attention_mask=attention_mask, 
            past_key_values=past_key_values, 
            inputs_embeds=inputs_embeds, 
        ) 
        hidden_states = outputs[0] 
        logits = self.model.lm_head(hidden_states) 

        return (logits,) + outputs[1:]

def set_input_names(model, past_key_values): 
    input_names = [ 
        "input_ids", 
        "attention_mask", 
        *itertools.chain.from_iterable([f"past_key_values.{idx}.key", f"past_key_values.{idx}.value"] for idx, _ in enumerate(past_key_values)), 
    ] 
    assert len(input_names) == len(model.inputs) 
    for _input, input_name in zip(model.inputs, input_names):         _input.get_tensor().set_names({input_name})

def set_output_names(model, past_key_values): 
    output_names = [ 
        "logits", 
        *itertools.chain.from_iterable([f"present.{idx}.key", f"present.{idx}.value"] for idx, _ in enumerate(past_key_values)), 
    ] 
    assert len(output_names) == len(model.outputs) 
    for out, out_name in zip(ov_model.outputs, output_names): 
        out.get_tensor().set_names({out_name})

example_input = { 
    "inputs_embeds": torch.zeros((1, 205, 2048)), 
    "attention_mask": torch.ones((1, 205), dtype=torch.long), 
} 

wrapped = ModelWrapper(model) 
past_key_values = wrapped(**example_input)[1] 

if not stage1_xml_path.exists(): 
    ov_model = ov.convert_model(wrapped, example_input=example_input) 
    set_output_names(ov_model, past_key_values) 
    ov_model = nncf.compress_weights(ov_model, **wc_parameters) 
    ov.save_model(ov_model, stage1_xml_path) 
    cleanup_torchscript_cache() 
    del ov_model 
    gc.collect()

example_input = { 
    "input_ids": torch.ones((1, 1), dtype=torch.long), 
    "past_key_values": past_key_values, 
    "attention_mask": torch.ones((1, past_key_values[-1][-1].shape[-2] + 1), dtype=torch.long), 
} 

if not stage2_xml_path.exists(): 
    ov_model = ov.convert_model( 
        wrapped, 
        example_input=example_input, 
    ) 
    set_input_names(ov_model, past_key_values) 
    set_output_names(ov_model, past_key_values) 
    ov_model = nncf.compress_weights(ov_model, **wc_parameters) 
    ov.save_model(ov_model, stage2_xml_path) 
    cleanup_torchscript_cache() 
    del ov_model 
    gc.collect()

prepare_inputs_labels_for_multimodal = model.prepare_inputs_labels_for_multimodal 
prepare_inputs_for_generation = model.prepare_inputs_for_generation 
config = model.config 
config.save_pretrained(MODELS_DIR)

del wrapped 
del model 
gc.collect();

class OVMobileLlamaForCausalLM(transformers.GenerationMixin): 
    def __init__(self, stage1_path, stage2_path, device): 
        self.stage1 = core.compile_model(stage1_path, device) 
        self.stage2 = core.read_model(stage2_path) 
        self.generation_config = transformers.GenerationConfig.from_model_config(config) 
        self.config = transformers.AutoConfig.from_pretrained(MODELS_DIR) 
        self.main_input_name = "input_ids" 
        self.device = torch.device("cpu") 
        self.prepare_inputs_for_generation = prepare_inputs_for_generation 
        self.num_pkv = 2 
        self.input_names = {key.get_any_name(): idx for idx, key in enumerate(self.stage2.inputs)} 
        self.output_names = {key.get_any_name(): idx for idx, key in enumerate(self.stage2.outputs)} 
        self.key_value_input_names = [key for key in self.input_names if "key_values" in key] 
        self.key_value_output_names = [key for key in self.output_names if "present" in key] 
        stage2 = core.compile_model(self.stage2, device) 
        self.request = stage2.create_infer_request() 
        self._supports_cache_class = False 

    def can_generate(self): 
        """Returns True to validate the check that the model using `GenerationMixin.generate()` can indeed generate.""" 
        return True 

    def __call__( 
        self, 
        input_ids: torch.LongTensor, 
        images: torch.Tensor, 
        attention_mask: Optional[torch.LongTensor] = None, 
        prefix_mask: Optional[torch.LongTensor] = None, 
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 
        **kwargs, 
    ) -> transformers.modeling_outputs.CausalLMOutputWithPast: 
        return self.forward(input_ids, images, attention_mask, prefix_mask, past_key_values) 

    def forward( 
        self, 
        input_ids: torch.LongTensor, 
        images: torch.Tensor, 
        attention_mask: Optional[torch.LongTensor] = None, 
        prefix_mask: Optional[torch.LongTensor] = None, 
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, 
        **kwargs, 
    ) -> transformers.modeling_outputs.CausalLMOutputWithPast: 
        """General inference method""" 
        inputs = {} 
        if past_key_values is not None:
            # past_key_values をフラット化 
            attention_mask = torch.ones( 
                (input_ids.shape[0], past_key_values[-1][-1].shape[-2] + 1), 
                dtype=input_ids.dtype, 
            ) 
            past_key_values = tuple(past_key_value for pkv_per_layer in past_key_values for past_key_value in pkv_per_layer) 
            # past_key_values をデコーダー入力に追加 
            inputs = dict(zip(self.key_value_input_names, past_key_values)) 

        else: 
            return self.forward_with_image(input_ids, images, attention_mask) 
        inputs["input_ids"] = np.array(input_ids) 

        if "attention_mask" in self.input_names: 
            inputs["attention_mask"] = np.array(attention_mask) 

        # 推論を実行 
        self.request.start_async(inputs, share_inputs=True) 
        self.request.wait() 

        logits = torch.from_numpy(self.request.get_tensor("logits").data) 

        # 長さが等しいタプル: レイヤー数 * デコーダーレイヤーごとの past_key_value の数 (2 は自己注意レイヤーに対応) 
        past_key_values = tuple(self.request.get_tensor(key).data for key in self.key_value_output_names) 
        # 長さが `n_layers` のタプルのタプル。各タプルの長さは 2 (自己注意の k/v) に等しい 

        past_key_values = tuple(past_key_values[i : i + self.num_pkv] for i in range(0, len(past_key_values), self.num_pkv)) 

        return transformers.modeling_outputs.CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values) 

    def forward_with_image(self, input_ids, images, attention_mask): 
        """First step inference method, that resolves multimodal data""" 
        _, attention_mask, _, input_embed, _ = prepare_inputs_labels_for_multimodal(input_ids, attention_mask, images=images, past_key_values=None, labels=None) 
        outs = self.stage1({"inputs_embeds": input_embed, "attention_mask": attention_mask}) 
        logits = outs[0] 
        pkv = list(outs.values())[1:] 
        pkv = tuple(pkv[i : i + self.num_pkv] for i in range(0, len(pkv), self.num_pkv)) 
        return transformers.modeling_outputs.CausalLMOutputWithPast(logits=torch.from_numpy(logits), past_key_values=pkv)

core = ov.Core() 

device = widgets.Dropdown( 
    options=core.available_devices + ["AUTO"], 
    value="AUTO", 
    description="Device:", 
    disabled=False, 
) 

device

def generate(img, prompt): 
    images_tensor = process_images([img], image_processor, transformers.AutoConfig.from_pretrained(MODELS_DIR)) 
    prompt = DEFAULT_IMAGE_TOKEN + "\n" + prompt 
    conv = conv_templates["v1"].copy() 
    conv.append_message(conv.roles[0], prompt) 
    conv.append_message(conv.roles[1], None) 
    prompt = conv.get_prompt() 
    stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2 
    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0) 
    stopping_criteria = KeywordsStoppingCriteria([stop_str], tokenizer, input_ids) 

    output_ids = ov_model.generate( 
        input_ids, 
        images=images_tensor, 
        do_sample=True if TEMPERATURE > 0 else False, temperature=TEMPERATURE, 
        top_p=TOP_P, 
        num_beams=NUM_BEAMS, 
        max_new_tokens=MAX_NEW_TOKENS, 
        use_cache=True, 
        stopping_criteria=[stopping_criteria], 
    ) 
    input_token_len = input_ids.shape[1] 
    outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0] 
    outputs = outputs.strip() 
    if outputs.endswith(stop_str): 
        outputs = outputs[: -len(stop_str)] 

    return outputs.strip() 

demo = gr.Interface( 
    generate, 
    [gr.Image(label="Image", type="pil"), gr.Textbox(label="Prompt")], 
    gr.Textbox(), 
    examples=[ 
        [ 
            str(IMAGE_PATH), 
            PROMPT_STR, 
        ] 
    ], 
    allow_flagging="never", 
) 

try: 
    demo.launch(debug=False) 
except Exception: 
    demo.launch(debug=False, share=True) 
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細はドキュメントをご覧ください: https://gradio.app/docs/