hymenjj/llama-cpp-python-prebuilt

1from __future__ import annotations2 3import os4from ctypes import (5    c_bool,6    c_char_p,7    c_int,8    c_uint8,9    c_uint32,10    c_float,11    c_void_p,12    c_size_t,13    POINTER,14    _Pointer,  # type: ignore15    Structure,16    byref,17)18import pathlib19from typing import (20    Union,21    NewType,22    Optional,23    TYPE_CHECKING,24)25 26import llama_cpp.llama_cpp as llama_cpp27 28from llama_cpp._ctypes_extensions import (29    load_shared_library,30    ctypes_function_for_shared_library,31)32 33if TYPE_CHECKING:34    from llama_cpp._ctypes_extensions import (35        CtypesArray,36    )37 38 39# Specify the base name of the shared library to load40_libmtmd_base_name = "mtmd"41_libmtmd_override_path = os.environ.get("MTMD_CPP_LIB")42_libmtmd_base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib" if _libmtmd_override_path is None else pathlib.Path()43 44# Load the library45_libmtmd = load_shared_library(_libmtmd_base_name, _libmtmd_base_path)46 47ctypes_function = ctypes_function_for_shared_library(_libmtmd)48 49################################################50# mtmd.h types51################################################52 53# Opaque types54mtmd_context_p = NewType("mtmd_context_p", int)55mtmd_context_p_ctypes = c_void_p56 57mtmd_bitmap_p = NewType("mtmd_bitmap_p", int)58mtmd_bitmap_p_ctypes = c_void_p59 60mtmd_image_tokens_p = NewType("mtmd_image_tokens_p", int)61mtmd_image_tokens_p_ctypes = c_void_p62 63mtmd_input_chunk_p = NewType("mtmd_input_chunk_p", int)64mtmd_input_chunk_p_ctypes = c_void_p65 66mtmd_input_chunks_p = NewType("mtmd_input_chunks_p", int)67mtmd_input_chunks_p_ctypes = c_void_p68 69# Enums70MTMD_INPUT_CHUNK_TYPE_TEXT = 071MTMD_INPUT_CHUNK_TYPE_IMAGE = 172MTMD_INPUT_CHUNK_TYPE_AUDIO = 273 74# Structures75class mtmd_context_params(Structure):76    _fields_ = [77        ("use_gpu", c_bool),78        ("print_timings", c_bool),79        ("n_threads", c_int),80        ("verbosity", c_int),  # ggml_log_level81        ("image_marker", c_char_p),82        ("media_marker", c_char_p),83    ]84 85class mtmd_input_text(Structure):86    _fields_ = [87        ("text", c_char_p),88        ("add_special", c_bool),89        ("parse_special", c_bool),90    ]91 92################################################93# mtmd.h functions94################################################95 96# MTMD_API const char * mtmd_default_marker(void);97@ctypes_function("mtmd_default_marker", [], c_char_p)98def mtmd_default_marker() -> bytes:99    ...100 101# MTMD_API struct mtmd_context_params mtmd_context_params_default(void);102@ctypes_function("mtmd_context_params_default", [], mtmd_context_params)103def mtmd_context_params_default() -> mtmd_context_params:104    ...105 106# MTMD_API mtmd_context * mtmd_init_from_file(const char * mmproj_fname,107#                                             const struct llama_model * text_model,108#                                             const struct mtmd_context_params ctx_params);109@ctypes_function(110    "mtmd_init_from_file",111    [c_char_p, llama_cpp.llama_model_p_ctypes, mtmd_context_params],112    mtmd_context_p_ctypes113)114def mtmd_init_from_file(115    mmproj_fname: bytes,116    text_model: llama_cpp.llama_model_p,117    ctx_params: mtmd_context_params,118    /,119) -> Optional[mtmd_context_p]:120    ...121 122# MTMD_API void mtmd_free(mtmd_context * ctx);123@ctypes_function("mtmd_free", [mtmd_context_p_ctypes], None)124def mtmd_free(ctx: mtmd_context_p, /):125    ...126 127# MTMD_API bool mtmd_support_vision(mtmd_context * ctx);128@ctypes_function("mtmd_support_vision", [mtmd_context_p_ctypes], c_bool)129def mtmd_support_vision(ctx: mtmd_context_p, /) -> bool:130    ...131 132# MTMD_API mtmd_bitmap * mtmd_bitmap_init(uint32_t nx, uint32_t ny, const unsigned char * data);133@ctypes_function(134    "mtmd_bitmap_init",135    [c_uint32, c_uint32, POINTER(c_uint8)],136    mtmd_bitmap_p_ctypes137)138def mtmd_bitmap_init(139    nx: Union[c_uint32, int],140    ny: Union[c_uint32, int],141    data: CtypesArray[c_uint8],142    /,143) -> Optional[mtmd_bitmap_p]:144    ...145 146# MTMD_API void mtmd_bitmap_free(mtmd_bitmap * bitmap);147@ctypes_function("mtmd_bitmap_free", [mtmd_bitmap_p_ctypes], None)148def mtmd_bitmap_free(bitmap: mtmd_bitmap_p, /):149    ...150 151# MTMD_API mtmd_input_chunks * mtmd_input_chunks_init(void);152@ctypes_function("mtmd_input_chunks_init", [], mtmd_input_chunks_p_ctypes)153def mtmd_input_chunks_init() -> Optional[mtmd_input_chunks_p]:154    ...155 156# MTMD_API void mtmd_input_chunks_free(mtmd_input_chunks * chunks);157@ctypes_function("mtmd_input_chunks_free", [mtmd_input_chunks_p_ctypes], None)158def mtmd_input_chunks_free(chunks: mtmd_input_chunks_p, /):159    ...160 161# MTMD_API size_t mtmd_input_chunks_size(const mtmd_input_chunks * chunks);162@ctypes_function("mtmd_input_chunks_size", [mtmd_input_chunks_p_ctypes], c_size_t)163def mtmd_input_chunks_size(chunks: mtmd_input_chunks_p, /) -> int:164    ...165 166# MTMD_API const mtmd_input_chunk * mtmd_input_chunks_get(const mtmd_input_chunks * chunks, size_t idx);167@ctypes_function(168    "mtmd_input_chunks_get",169    [mtmd_input_chunks_p_ctypes, c_size_t],170    mtmd_input_chunk_p_ctypes171)172def mtmd_input_chunks_get(173    chunks: mtmd_input_chunks_p, idx: Union[c_size_t, int], /174) -> Optional[mtmd_input_chunk_p]:175    ...176 177# MTMD_API int32_t mtmd_tokenize(mtmd_context * ctx,178#                                mtmd_input_chunks * output,179#                                const mtmd_input_text * text,180#                                const mtmd_bitmap ** bitmaps,181#                                size_t n_bitmaps);182@ctypes_function(183    "mtmd_tokenize",184    [185        mtmd_context_p_ctypes,186        mtmd_input_chunks_p_ctypes,187        POINTER(mtmd_input_text),188        POINTER(mtmd_bitmap_p_ctypes),189        c_size_t,190    ],191    c_int,192)193def mtmd_tokenize(194    ctx: mtmd_context_p,195    output: mtmd_input_chunks_p,196    text: "_Pointer[mtmd_input_text]",197    bitmaps: CtypesArray[mtmd_bitmap_p_ctypes],198    n_bitmaps: Union[c_size_t, int],199    /,200) -> int:201    ...202 203# MTMD_API size_t mtmd_input_chunk_get_n_tokens(const mtmd_input_chunk * chunk);204@ctypes_function("mtmd_input_chunk_get_n_tokens", [mtmd_input_chunk_p_ctypes], c_size_t)205def mtmd_input_chunk_get_n_tokens(chunk: mtmd_input_chunk_p, /) -> int:206    ...207 208# MTMD_API enum mtmd_input_chunk_type mtmd_input_chunk_get_type(const mtmd_input_chunk * chunk);209@ctypes_function("mtmd_input_chunk_get_type", [mtmd_input_chunk_p_ctypes], c_int)210def mtmd_input_chunk_get_type(chunk: mtmd_input_chunk_p, /) -> int:211    ...212 213# MTMD_API const llama_token * mtmd_input_chunk_get_tokens_text(const mtmd_input_chunk * chunk, size_t * n_tokens_output);214@ctypes_function(215    "mtmd_input_chunk_get_tokens_text",216    [mtmd_input_chunk_p_ctypes, POINTER(c_size_t)],217    POINTER(llama_cpp.llama_token)218)219def mtmd_input_chunk_get_tokens_text(220    chunk: mtmd_input_chunk_p, n_tokens_output: "_Pointer[c_size_t]", /221) -> Optional["_Pointer[llama_cpp.llama_token]"]:222    ...223 224################################################225# mtmd-helper.h functions226################################################227 228# MTMD_API mtmd_bitmap * mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len);229@ctypes_function(230    "mtmd_helper_bitmap_init_from_buf",231    [mtmd_context_p_ctypes, POINTER(c_uint8), c_size_t],232    mtmd_bitmap_p_ctypes233)234def mtmd_helper_bitmap_init_from_buf(235    ctx: mtmd_context_p,236    buf: CtypesArray[c_uint8],237    length: Union[c_size_t, int],238    /,239) -> Optional[mtmd_bitmap_p]:240    ...241 242# MTMD_API size_t mtmd_helper_get_n_tokens(const mtmd_input_chunks * chunks);243@ctypes_function("mtmd_helper_get_n_tokens", [mtmd_input_chunks_p_ctypes], c_size_t)244def mtmd_helper_get_n_tokens(chunks: mtmd_input_chunks_p, /) -> int:245    ...246 247# MTMD_API int32_t mtmd_helper_eval_chunk_single(mtmd_context * ctx,248#                                                struct llama_context * lctx,249#                                                const mtmd_input_chunk * chunk,250#                                                llama_pos n_past,251#                                                llama_seq_id seq_id,252#                                                int32_t n_batch,253#                                                bool logits_last,254#                                                llama_pos * new_n_past);255@ctypes_function(256    "mtmd_helper_eval_chunk_single",257    [258        mtmd_context_p_ctypes,259        llama_cpp.llama_context_p_ctypes,260        mtmd_input_chunk_p_ctypes,261        llama_cpp.llama_pos,262        llama_cpp.llama_seq_id,263        c_int,264        c_bool,265        POINTER(llama_cpp.llama_pos),266    ],267    c_int,268)269def mtmd_helper_eval_chunk_single(270    ctx: mtmd_context_p,271    lctx: llama_cpp.llama_context_p,272    chunk: mtmd_input_chunk_p,273    n_past: llama_cpp.llama_pos,274    seq_id: llama_cpp.llama_seq_id,275    n_batch: Union[c_int, int],276    logits_last: Union[c_bool, bool],277    new_n_past: "_Pointer[llama_cpp.llama_pos]",278    /,279) -> int:280    ...281