from __future__ import annotations

import logging

import argparse

import concurrent.futures

import enum

import faulthandler

import functools

import itertools

import json

import math

import mmap

import os

import pickle

import re

import signal

import struct

import sys

import textwrap

import time

import zipfile

from abc import ABC, abstractmethod

from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor

from dataclasses import dataclass

from pathlib import Path

from typing import TYPE_CHECKING, Any, Callable, ClassVar, IO, Iterable, Literal, Protocol, TypeVar, runtime_checkable, Tuple

import configparser

import numpy as np

from sentencepiece import SentencePieceProcessor

if 'NO_LOCAL_GGUF' not in os.environ:

sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))

import gguf

if TYPE_CHECKING:

from typing_extensions import Self, TypeAlias

logger = logging.getLogger("convert")

if hasattr(faulthandler, 'register') and hasattr(signal, 'SIGUSR1'):

faulthandler.register(signal.SIGUSR1)

NDArray: TypeAlias = 'np.ndarray[Any, Any]'

ARCH = gguf.MODEL_ARCH.BITNET

DEFAULT_CONCURRENCY = 16

ADDED_TOKENS_FILE = 'added_tokens.json'

FAST_TOKENIZER_FILE = 'tokenizer.json'

class DataType:

return n_elements * self.dtype.itemsize

class UnquantizedDataType(DataType):

pass

DT_F16 = UnquantizedDataType('F16', dtype = np.dtype(np.float16), valid_conversions = ['F32', 'Q8_0'])

DT_F32 = UnquantizedDataType('F32', dtype = np.dtype(np.float32), valid_conversions = ['F16', 'Q8_0', 'I2'])

DT_I32 = UnquantizedDataType('I32', dtype = np.dtype(np.int16), valid_conversions = [])

DT_BF16 = UnquantizedDataType('BF16', dtype = np.dtype(np.uint16), valid_conversions = ['F32', 'F16', 'Q8_0'])

DT_I2 = UnquantizedDataType('I2', dtype = np.dtype(np.uint8), valid_conversions = ['F32', 'F16', 'Q8_0'])

class QuantizedDataType(DataType):

return self.quantized_dtype.itemsize * (n_elements // self.block_size)

class Q8_0QuantizedDataType(QuantizedDataType):

return np.fromiter(quantize_blocks_q8_0(blocks), count = n_blocks, dtype = self.quantized_dtype)

DT_Q8_0 = Q8_0QuantizedDataType('Q8_0',

dtype = np.dtype(np.float32), valid_conversions = [],

ggml_type = gguf.GGMLQuantizationType.Q8_0, block_size = 32,

quantized_dtype = np.dtype([('d', '<f2'), ('qs', 'i1', (32,))]))

NUMPY_TYPE_TO_DATA_TYPE: dict[np.dtype[Any], DataType] = {}

for dt in (DT_BF16, DT_F16, DT_F32, DT_I32, DT_I2):

if dt.dtype in NUMPY_TYPE_TO_DATA_TYPE:

raise ValueError(f'Invalid duplicate data type {dt}')

NUMPY_TYPE_TO_DATA_TYPE[dt.dtype] = dt

SAFETENSORS_DATA_TYPES: dict[str, DataType] = {

}

class GGMLFileType(enum.IntEnum):

return dt

GGML_FILE_TYPE_TO_DATA_TYPE: dict[GGMLFileType, DataType] = {

}

class Params:

return params

class BaseVocab(Protocol):

tokenizer_model: ClassVar[str]

name: ClassVar[str]

class NoVocab(BaseVocab):

return "<NoVocab for a model without integrated vocabulary>"

class Vocab(BaseVocab, Protocol):

def all_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]: ...

class BpeVocab(Vocab):

return f"<BpeVocab with {self.vocab_size_base} base tokens and {len(self.added_tokens_list)} added tokens>"

class SentencePieceVocab(Vocab):

return f"<SentencePieceVocab with {self.vocab_size_base} base tokens and {len(self.added_tokens_list)} added tokens>"

class LlamaHfVocab(Vocab):

return f"<LlamaHfVocab with {self.vocab_size_base} base tokens and {len(self.added_tokens_list)} added tokens>"

def permute(weights: NDArray, n_head: int, n_head_kv: int) -> NDArray:

.reshape(weights.shape))

class Tensor(ABC):

def to_ggml(self) -> GGMLCompatibleTensor: ...

def bf16_to_fp32(bf16_arr: np.ndarray[Any, np.dtype[np.uint16]]) -> NDArray:

assert bf16_arr.dtype == np.uint16, f"Input array should be of dtype uint16, but got {bf16_arr.dtype}"

fp32_arr = bf16_arr.astype(np.uint32) << 16

return fp32_arr.view(np.float32)

def preprocess_weights(

return w

def transform_to_i2(x : NDArray):

return ans, scale

class UnquantizedTensor(Tensor):

return UnquantizedTensor(permute(self.ndarray, n_head, n_head_kv))

def load_unquantized(lazy_tensor: LazyTensor, expected_dtype: Any = None, convert: bool = False) -> NDArray:

return tensor.ndarray

GGMLCompatibleTensor = UnquantizedTensor

class LazyTensor:

raise ValueError(f'Cannot validate conversion from {self.data_type} to {data_type}.')

LazyModel: TypeAlias = 'dict[str, LazyTensor]'

class ModelPlus:

vocab: BaseVocab | None # For GGML models (which have vocab built in), the vocab.

def merge_sharded(models: list[LazyModel]) -> LazyModel:

return {name: convert(name) for name in names}

def merge_multifile_models(models_plus: list[ModelPlus]) -> ModelPlus:

return ModelPlus(model, paths, format, vocab) # pytype: disable=wrong-arg-types

def permute_lazy(lazy_tensor: LazyTensor, n_head: int, n_head_kv: int) -> LazyTensor:

def load() -> Tensor:

return lazy_tensor.load().permute(n_head, n_head_kv)

return LazyTensor(load, lazy_tensor.shape, lazy_tensor.data_type, f'permute({n_head}, {n_head_kv}) ' + lazy_tensor.description)

def permute_part_lazy(lazy_tensor: LazyTensor, n_part: int, n_head: int, n_head_kv: int) -> LazyTensor:

return LazyTensor(load, s, lazy_tensor.data_type, f'permute({n_head}, {n_head_kv}) ' + lazy_tensor.description)

def part_lazy(lazy_tensor: LazyTensor, n_part: int) -> LazyTensor:

return LazyTensor(load, s, lazy_tensor.data_type, 'part ' + lazy_tensor.description)

def pack_experts_lazy(lazy_tensors: list[LazyTensor]) -> LazyTensor:

return LazyTensor(load, s, lazy_tensors[0].data_type, 'pack_experts ' + ' | '.join(lt.description for lt in lazy_tensors))