Databricks Dolly 2.0 と OpenVINO を使用した手順#

from IPython.display import Markdown, display 
import ipywidgets as widgets 

prepare_int4_model = widgets.Checkbox( 
    value=True, 
    description="Prepare INT4 model", 
    disabled=False, 
) 
prepare_int8_model = widgets.Checkbox( 
    value=False, 
    description="Prepare INT8 model", 
    disabled=False, 
) 
prepare_fp16_model = widgets.Checkbox( 
    value=False, 
    description="Prepare FP16 model", 
    disabled=False, 
) 

display(prepare_int4_model) 
display(prepare_int8_model) 
display(prepare_fp16_model)

from pathlib import Path 

model_id = "databricks/dolly-v2-3b" 
model_path = Path("dolly-v2-3b") 

fp16_model_dir = model_path / "FP16" 
int8_model_dir = model_path / "INT8_compressed_weights" 
int4_model_dir = model_path / "INT4_compressed_weights" 

def convert_to_fp16(): 
    if (fp16_model_dir / "openvino_model.xml").exists(): 
        return 
    fp16_model_dir.mkdir(parents=True, exist_ok=True) 
    export_command_base = "optimum-cli export openvino --model {} --task text-generation-with-past --weight-format fp16".format(model_id) 
    export_command = export_command_base + " " + str(fp16_model_dir) 
    display(Markdown("**Export command:**")) 
    display(Markdown(f"`{export_command}`")) 
    ! $export_command 

def convert_to_int8(): 
    if (int8_model_dir / "openvino_model.xml").exists(): 
        return 
    int8_model_dir.mkdir(parents=True, exist_ok=True) 
    export_command_base = "optimum-cli export openvino --model {} --task text-generation-with-past --weight-format int8".format(model_id) 
    export_command = export_command_base + " " + str(int8_model_dir) 
    display(Markdown("**Export command:**")) 
    display(Markdown(f"`{export_command}`")) 
    ! $export_command 

def convert_to_int4(): 
    if (int4_model_dir / "openvino_model.xml").exists(): 
        return 
    int4_model_dir.mkdir(parents=True, exist_ok=True) 
    export_command_base = "optimum-cli export openvino --model {} --task text-generation-with-past --weight-format int4".format(model_id) 
    export_command = export_command_base + " " + str(int4_model_dir) 
    display(Markdown("**Export command:**")) 
    display(Markdown(f"`{export_command}`")) 
    ! $export_command 

if prepare_fp16_model.value: 
    convert_to_fp16() 
if prepare_int8_model.value: 
    convert_to_int8() 
if prepare_int4_model.value: 
    convert_to_int4()

fp16_weights = fp16_model_dir / "openvino_model.bin" 
int8_weights = int8_model_dir / "openvino_model.bin" 
int4_weights = int4_model_dir / "openvino_model.bin" 

if fp16_weights.exists(): 
    print(f"Size of FP16 model is {fp16_weights.stat().st_size / 1024 / 1024:.2f} MB") 
for precision, compressed_weights in zip([8, 4], [int8_weights, int4_weights]): 
    if compressed_weights.exists(): 
        print(f"Size of model with INT{precision} compressed weights is {compressed_weights.stat().st_size / 1024 / 1024:.2f} MB") 
    if compressed_weights.exists() and fp16_weights.exists(): 
        print(f"Compression rate for INT{precision} model: {fp16_weights.stat().st_size / compressed_weights.stat().st_size:.3f}")

available_models = [] 
if int4_model_dir.exists(): 
    available_models.append("INT4") 
if int8_model_dir.exists(): 
    available_models.append("INT8") 
if fp16_model_dir.exists(): 
    available_models.append("FP16") 

model_to_run = widgets.Dropdown( 
    options=available_models, 
    value=available_models[0], 
    description="Model to run:", 
    disabled=False, 
) 

model_to_run

import ipywidgets as widgets 
import openvino as ov 

core = ov.Core() 

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

device

from threading import Thread 
from time import perf_counter 
from typing import List 
import gradio as gr 
from transformers import AutoTokenizer, TextIteratorStreamer 
import numpy as np

INSTRUCTION_KEY = "### Instruction:" 
RESPONSE_KEY = "### Response:" 
END_KEY = "### End" 
INTRO_BLURB = "Below is an instruction that describes a task. Write a response that appropriately completes the request." 

# これは、すでにトレーニング済みのモデルを使用して応答を生成するのに使用されるプロンプトです。これは応答キーで終了し、 
# モデルの役割はそれに続く補完 (つまり応答自体) を提供することです。
PROMPT_FOR_GENERATION_FORMAT =
"""{intro} 

{instruction_key} 
{instruction} 

{response_key} 
""".format( 
    intro=INTRO_BLURB, 
    instruction_key=INSTRUCTION_KEY, 
    instruction="{instruction}", 
    response_key=RESPONSE_KEY, 
)

def get_special_token_id(tokenizer: AutoTokenizer, key: str) -> int: 
    """ 
    Gets the token ID for a given string that has been added to the tokenizer as a special token.
    When training, we configure the tokenizer so that the sequences like "### Instruction:" and "### End" are 
    treated specially and converted to a single, new token. This retrieves the token ID each of these keys map to.
 
    Args: 
        tokenizer (PreTrainedTokenizer): the tokenizer 
        key (str): the key to convert to a single token 
    Raises:
        RuntimeError: if more than one ID was generated 

    Returns: 
        int: the token ID for the given key 
    """ 
    token_ids = tokenizer.encode(key) 
    if len(token_ids) > 1: 
        raise ValueError(f"Expected only a single token for '{key}' but found {token_ids}") 
    return token_ids[0] 

tokenizer_response_key = next( 
    (token for token in tokenizer.additional_special_tokens if token.startswith(RESPONSE_KEY)), 
    None, 
) 

end_key_token_id = None 
if tokenizer_response_key: 
    try: 
        end_key_token_id = get_special_token_id(tokenizer, END_KEY) 
        # "### End" が生成されたら生成が停止することを確認 
    except ValueError: 
        pass

def run_generation( 
    user_text: str, 
    top_p: float, 
    temperature: float, 
    top_k: int, 
    max_new_tokens: int, 
    perf_text: str, 
): 
    """ 
    Text generation function 

    Parameters: 
        user_text (str): User-provided instruction for a generation. 
        top_p (float): Nucleus sampling.If set to < 1, only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for a generation. 
        temperature (float): The value used to module the logits distribution. 
        top_k (int): The number of highest probability vocabulary tokens to keep for top-k-filtering. 
        max_new_tokens (int): Maximum length of generated sequence. 
        perf_text (str): Content of text field for printing performance results.
   Returns: 
        model_output (str) - model-generated text 
        perf_text (str) - updated perf text filed content 
    """ 

    # モデルの予想されるテンプレートに従って入力プロンプトを準備 
    prompt_text = PROMPT_FOR_GENERATION_FORMAT.format(instruction=user_text) 

    # ユーザーのテキストをトークン化 
    model_inputs = tokenizer(prompt_text, return_tensors="pt") 

    # UI をブロックしないように、別のスレッドで生成を開始。テキストはメインスレッドの 
    # ストリーマーから取得されます。生成スレッドで例外を処理するため、ストリーマーにタイムアウトを追加します 
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) 
    generate_kwargs = dict( 
        model_inputs, 
        streamer=streamer, 
        max_new_tokens=max_new_tokens, 
        do_sample=True, 
        top_p=top_p, 
        temperature=float(temperature), 
        top_k=top_k, 
        eos_token_id=end_key_token_id, 
    ) 
    t = Thread(target=ov_model.generate, kwargs=generate_kwargs) 
    t.start() 

    # 生成されたテキストをストリーマーから取得し、モデル出力を更新 
    model_output = "" 
    per_token_time = [] 
    num_tokens = 0 
    start = perf_counter() 
    for new_text in streamer: 
        current_time = perf_counter() - start 
        model_output += new_text 
        perf_text, num_tokens = estimate_latency(current_time, perf_text, new_text, per_token_time, num_tokens) 
        yield model_output, perf_text 
        start = perf_counter() 
    return model_output, perf_text

def estimate_latency( 
    current_time: float, 
    current_perf_text: str, 
    new_gen_text: str, 
    per_token_time: List[float], 
    num_tokens: int, 
): 
    """ 
    Helper function for performance estimation 

    Parameters: 
        current_time (float): This step time in seconds. 
        current_perf_text (str): Current content of performance UI field. 
        new_gen_text (str): New generated text. 
        per_token_time (List[float]): history of performance from previous steps. 
        num_tokens (int): Total number of generated tokens.
    Returns: 
        update for performance text field 
        update for a total number of tokens 
    """ 
    num_current_toks = len(tokenizer.encode(new_gen_text)) 
    num_tokens += num_current_toks 
    per_token_time.append(num_current_toks / current_time) 
    if len(per_token_time) > 10 and len(per_token_time) % 4 == 0: 
        current_bucket = per_token_time[:-10] 
        return ( 
            f"Average generation speed: {np.mean(current_bucket):.2f} tokens/s.Total generated tokens: {num_tokens}", 
            num_tokens, 
        ) 
    return current_perf_text, num_tokens 

def reset_textbox(instruction: str, response: str, perf: str): 
    """ 
    Helper function for resetting content of all text fields

    Parameters: 
        instruction (str): Content of user instruction field. 
        response (str): Content of model response field. 
        perf (str): Content of performance info filed 

    Returns: 
        empty string for each placeholder 
    """ 
    return "", "", "" 

def select_device(device_str: str, current_text: str = "", progress: gr.Progress = gr.Progress()): 
    """ 
    Helper function for uploading model on the device.
    Parameters: 
        device_str (str): Device name. 
        current_text (str): Current content of user instruction field (used only for backup purposes, temporally replacing it on the progress bar during model loading). 
        progress (gr.Progress): gradio progress tracker 
    Returns: 
        current_text 
    """ 
    if device_str != ov_model._device: 
        ov_model.request = None 
        ov_model._device = device_str 

        for i in progress.tqdm(range(1), desc=f"Model loading on {device_str}"): 
            ov_model.compile() 
    return current_text