OpenVINO を使用した LLM 指示追従パイプライン#

from IPython.display import display, Markdown 

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 
import logging 
import openvino as ov 
import nncf 

nncf.set_log_level(logging.ERROR) 

pt_model_id = model_configuration["model_id"] 
fp16_model_dir = Path(model_id.value) / "FP16" 
int8_model_dir = Path(model_id.value) / "INT8_compressed_weights" 
int4_model_dir = Path(model_id.value) / "INT4_compressed_weights" 

core = ov.Core() 

def convert_to_fp16(): 
    if (fp16_model_dir / "openvino_model.xml").exists(): 
        return 
    export_command_base = "optimum-cli export openvino --model {} --task text-generation-with-past --weight-format fp16".format(pt_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(pt_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(): 
    compression_configs = { 
        "mistral-7b": { 
            "sym": True, 
            "group_size": 64, 
            "ratio": 0.6, 
        }, 
        "red-pajama-3b-instruct": { 
            "sym": False, 
            "group_size": 128, 
            "ratio": 0.5, 
        }, 
        "dolly-v2-3b": {"sym": False, "group_size": 32, 
            "ratio": 0.5}, 
        "llama-3-8b-instruct": {"sym": True, "group_size": 128, "ratio": 1.0}, 
        "default": { 
            "sym": False, 
            "group_size": 128, 
            "ratio": 0.8, 
        }, 
    } 

    model_compression_params = compression_configs.get(model_id.value, compression_configs["default"]) 
    if (int4_model_dir / "openvino_model.xml").exists(): 
        return 
    export_command_base = "optimum-cli export openvino --model {} --task text-generation-with-past --weight-format int4".format(pt_model_id) 
    int4_compression_args = " --group-size {} --ratio {}".format(model_compression_params["group_size"], model_compression_params["ratio"]) 
    if model_compression_params["sym"]: 
        int4_compression_args += " --sym" 
    export_command_base += int4_compression_args 
    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}")

from threading import Thread 
from time import perf_counter 
from typing import List 
import gradio as gr 
import numpy as np 
from openvino_genai import StreamerBase 
from queue import Queue 
import re

core = ov.Core() 

detokinizer_dir = Path(model_dir, "openvino_detokenizer.xml") 

class TextIteratorStreamer(StreamerBase): 
    def __init__(self, tokenizer): 
        super().__init__() 
        self.tokenizer = tokenizer 
        self.compiled_detokenizer = core.compile_model(detokinizer_dir.as_posix()) 
        self.text_queue = Queue() 
        self.stop_signal = None 

    def __iter__(self): 
        return self 

    def __next__(self): 
        value = self.text_queue.get() 
        if value == self.stop_signal: 
            raise StopIteration() 
        else: 
            return value 

    def put(self, token_id): 
        openvino_output = self.compiled_detokenizer(np.array([[token_id]], dtype=int)) 
        text = str(openvino_output["string_output"][0]) 
        # ラベル/特殊記号を削除 
        text = text.lstrip("!") 
        text = re.sub("<.*>", "", text) 
        self.text_queue.put(text) 

    def end(self): 
        self.text_queue.put(self.stop_signal)

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 
    """ 

    # デコードステージの設定 
    config = pipe.get_generation_config() 
    config.temperature = temperature 
    if top_k > 0: 
        config.top_k = top_k 
    config.top_p = top_p 
    config.do_sample = True 
    config.max_new_tokens = max_new_tokens 

    # UI をブロックしないように、別のスレッドで生成を開始。テキストはメインスレッドの 
    # ストリーマーから取得されます。 

    streamer = TextIteratorStreamer(pipe.get_tokenizer()) 

    t = Thread(target=pipe.generate, args=(user_text, config, streamer)) 
    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, 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, 
    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. 
        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_tokens += 1 
    per_token_time.append(1 / 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 "", "", ""