wehub-resource-sync
187 lines
6.3 KiB
Python
187 lines
6.3 KiB
Python
"""The Python API for MLC Embeddings."""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple # noqa: UP035
|
|
|
|
import numpy as np
|
|
import tvm
|
|
import tvm_ffi
|
|
from tvm import relax
|
|
from tvm.contrib import tvmjs
|
|
from tvm.runtime import Device, Module
|
|
from tvm.runtime.vm import VirtualMachine
|
|
|
|
from mlc_llm.serve import engine_utils
|
|
from mlc_llm.support.auto_device import detect_device
|
|
from mlc_llm.tokenizers import Tokenizer
|
|
|
|
|
|
def _extract_metadata(mod: Module):
|
|
return json.loads(VirtualMachine(mod, tvm.runtime.device("cpu"))["_metadata"]())
|
|
|
|
|
|
def _load_params(
|
|
model_weight_path: str,
|
|
device: Device,
|
|
model_metadata: Dict[str, Any], # noqa: UP006
|
|
) -> List[tvm.runtime.Tensor]: # noqa: UP006
|
|
params, meta = tvmjs.load_tensor_cache(model_weight_path, device)
|
|
param_names = [param["name"] for param in model_metadata["params"]]
|
|
assert len(param_names) == meta["ParamSize"]
|
|
|
|
plist = []
|
|
for param_name in param_names:
|
|
plist.append(params[param_name])
|
|
return plist
|
|
|
|
|
|
def _get_tvm_module(
|
|
model_weight_path: str,
|
|
lib_path: str,
|
|
device: Device,
|
|
instrument: tvm_ffi.Function = None,
|
|
):
|
|
ex = tvm.runtime.load_module(lib_path)
|
|
vm = relax.VirtualMachine(ex, device)
|
|
if instrument:
|
|
vm.set_instrument(instrument)
|
|
metadata = _extract_metadata(ex)
|
|
params = _load_params(model_weight_path, device, metadata)
|
|
return vm.module, params, metadata
|
|
|
|
|
|
class DefaultDebugInstrument:
|
|
"""The default debug instrument to use if users don't specify
|
|
a customized one.
|
|
|
|
This debug instrument will dump the arguments and output of each
|
|
VM Call instruction into a .npz file. It will also alert the user
|
|
if any function outputs are NaN or INF.
|
|
"""
|
|
|
|
def __init__(self, debug_out: Path):
|
|
"""Constructor
|
|
|
|
Parameters
|
|
----------
|
|
debug_out : Path
|
|
the directory to dump the .npz files
|
|
"""
|
|
self.counter = 0
|
|
self.first_nan_occurred = False
|
|
self.first_inf_occurred = False
|
|
self.debug_out = debug_out
|
|
debug_out.mkdir(exist_ok=True, parents=True)
|
|
|
|
def reset(self, debug_out: Path):
|
|
"""Reset the state of the Instrument class
|
|
|
|
Parameters
|
|
----------
|
|
debug_out : Path
|
|
the directory to dump the .npz files
|
|
"""
|
|
self.counter = 0
|
|
self.first_nan_occurred = False
|
|
self.first_inf_occurred = False
|
|
self.debug_out = debug_out
|
|
debug_out.mkdir(exist_ok=True, parents=True)
|
|
|
|
def __call__(self, func, name, before_run, ret_val, *args):
|
|
# Determine what functions to look at
|
|
if before_run: # Whether before the function is called or after
|
|
return
|
|
if name.startswith("vm.builtin.") and "attention_with_fused_qkv" not in name:
|
|
return
|
|
|
|
# Decide what to print or save about the function's arguments (where args[-1] is the
|
|
# buffer we write the result to)
|
|
func_name = f"f{self.counter}_{name}"
|
|
|
|
# Save the arguments to npz
|
|
arg_dict = {}
|
|
for i, arg in enumerate(args):
|
|
if isinstance(arg, tvm.runtime.Tensor):
|
|
arg_dict[f"arg_{i}"] = arg.numpy()
|
|
|
|
np.savez(self.debug_out / f"{func_name}.npz", **arg_dict)
|
|
|
|
self.counter += 1
|
|
|
|
|
|
class MLCEmbeddings:
|
|
"""A class to embed queries using MLC LLM encoder models.
|
|
|
|
Parameters
|
|
----------
|
|
model: str
|
|
The model folder after compiling with MLC-LLM build process. The parameter
|
|
can either be the model name with its quantization scheme
|
|
(e.g. ``Llama-2-7b-chat-hf-q4f16_1``), or a full path to the model
|
|
folder. In the former case, we will use the provided name to search
|
|
for the model folder over possible paths.
|
|
|
|
model_lib_path : str
|
|
The full path to the model library file to use (e.g. a ``.so`` file).
|
|
|
|
device : Optional[str]
|
|
The description of the device to run on. User should provide a string in the
|
|
form of 'device_name:device_id' or 'device_name', where 'device_name' is one of
|
|
'cuda', 'metal', 'vulkan', 'rocm', 'opencl', 'auto' (automatically detect the
|
|
local device), and 'device_id' is the device id to run on. If no 'device_id'
|
|
is provided, it will be set to 0 by default.
|
|
|
|
debug_dir: Path
|
|
The output folder to store the dumped debug files. If None, will not dump any debug files.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
model_lib_path: str,
|
|
device: Optional[str] = "auto",
|
|
debug_dir: Optional[str] = None,
|
|
):
|
|
self.device = detect_device(device)
|
|
instrument = DefaultDebugInstrument(Path(debug_dir)) if debug_dir else None
|
|
self.mod, self.params, self.metadata = _get_tvm_module(
|
|
model, model_lib_path, self.device, instrument
|
|
)
|
|
self.model_path = model
|
|
self.tokenizer = Tokenizer(self.model_path)
|
|
self.prefill_func = self.mod["prefill"]
|
|
|
|
def embed(self, queries: List[str]) -> tvm.runtime.Tensor: # noqa: UP006
|
|
"""
|
|
Embeds a list of queries in a single batch.
|
|
|
|
Parameters
|
|
----------
|
|
queries : List[str]
|
|
A list of queries to embed.
|
|
|
|
Returns
|
|
-------
|
|
List[float]
|
|
A list of embeddings for the queries.
|
|
"""
|
|
tokens, attention_mask = self._tokenize_queries(queries)
|
|
tokens_tvm = tvm.runtime.tensor(tokens.astype("int32"), device=self.device)
|
|
attention_mask_tvm = tvm.runtime.tensor(attention_mask.astype("int32"), device=self.device)
|
|
output = self.prefill_func(tokens_tvm, attention_mask_tvm, self.params)
|
|
return output
|
|
|
|
def _tokenize_queries(self, queries: List[str]) -> Tuple[np.ndarray, np.ndarray]: # noqa: UP006
|
|
tokens = engine_utils.process_prompts(queries, self.tokenizer.encode)
|
|
max_query_length = max(len(token_seq) for token_seq in tokens)
|
|
|
|
token_inputs: np.ndarray = np.zeros((len(tokens), max_query_length), dtype=np.int32)
|
|
attention_mask: np.ndarray = np.zeros((len(tokens), max_query_length), dtype=np.int32)
|
|
|
|
for i, token_seq in enumerate(tokens):
|
|
token_inputs[i, : len(token_seq)] = token_seq
|
|
attention_mask[i, : len(token_seq)] = 1
|
|
|
|
return token_inputs, attention_mask
|