hymenjj/llama-cpp-python-prebuilt

1from __future__ import annotations2 3import json4 5from typing import Dict, Optional, Union, List6 7import llama_cpp8import llama_cpp.llama_speculative as llama_speculative9import llama_cpp.llama_tokenizer as llama_tokenizer10 11from llama_cpp.server.settings import ModelSettings12 13 14class LlamaProxy:15    def __init__(self, models: List[ModelSettings]) -> None:16        assert len(models) > 0, "No models provided!"17 18        self._model_settings_dict: dict[str, ModelSettings] = {}19        for model in models:20            if not model.model_alias:21                model.model_alias = model.model22            self._model_settings_dict[model.model_alias] = model23 24        self._current_model: Optional[llama_cpp.Llama] = None25        self._current_model_alias: Optional[str] = None26 27        self._default_model_settings: ModelSettings = models[0]28        self._default_model_alias: str = self._default_model_settings.model_alias  # type: ignore29 30        # Load default model31        self._current_model = self.load_llama_from_model_settings(32            self._default_model_settings33        )34        self._current_model_alias = self._default_model_alias35 36    def __call__(self, model: Optional[str] = None) -> llama_cpp.Llama:37        if model is None:38            model = self._default_model_alias39 40        if model not in self._model_settings_dict:41            model = self._default_model_alias42 43        if model == self._current_model_alias:44            if self._current_model is not None:45                return self._current_model46 47        if self._current_model:48            self._current_model.close()49        self._current_model = None50 51        settings = self._model_settings_dict[model]52        self._current_model = self.load_llama_from_model_settings(settings)53        self._current_model_alias = model54        return self._current_model55 56    def __getitem__(self, model: str):57        return self._model_settings_dict[model].model_dump()58 59    def __setitem__(self, model: str, settings: Union[ModelSettings, str, bytes]):60        if isinstance(settings, (bytes, str)):61            settings = ModelSettings.model_validate_json(settings)62        self._model_settings_dict[model] = settings63 64    def __iter__(self):65        for model in self._model_settings_dict:66            yield model67 68    def free(self):69        if self._current_model:70            self._current_model.close()71            del self._current_model72 73    @staticmethod74    def load_llama_from_model_settings(settings: ModelSettings) -> llama_cpp.Llama:75        chat_handler = None76        if settings.chat_format == "llava-1-5":77            assert settings.clip_model_path is not None, "clip model not found"78            if settings.hf_model_repo_id is not None:79                chat_handler = (80                    llama_cpp.llama_chat_format.Llava15ChatHandler.from_pretrained(81                        repo_id=settings.hf_model_repo_id,82                        filename=settings.clip_model_path,83                        verbose=settings.verbose,84                    )85                )86            else:87                chat_handler = llama_cpp.llama_chat_format.Llava15ChatHandler(88                    clip_model_path=settings.clip_model_path, verbose=settings.verbose89                )90        elif settings.chat_format == "obsidian":91            assert settings.clip_model_path is not None, "clip model not found"92            if settings.hf_model_repo_id is not None:93                chat_handler = (94                    llama_cpp.llama_chat_format.ObsidianChatHandler.from_pretrained(95                        repo_id=settings.hf_model_repo_id,96                        filename=settings.clip_model_path,97                        verbose=settings.verbose,98                    )99                )100            else:101                chat_handler = llama_cpp.llama_chat_format.ObsidianChatHandler(102                    clip_model_path=settings.clip_model_path, verbose=settings.verbose103                )104        elif settings.chat_format == "llava-1-6":105            assert settings.clip_model_path is not None, "clip model not found"106            if settings.hf_model_repo_id is not None:107                chat_handler = (108                    llama_cpp.llama_chat_format.Llava16ChatHandler.from_pretrained(109                        repo_id=settings.hf_model_repo_id,110                        filename=settings.clip_model_path,111                        verbose=settings.verbose,112                    )113                )114            else:115                chat_handler = llama_cpp.llama_chat_format.Llava16ChatHandler(116                    clip_model_path=settings.clip_model_path, verbose=settings.verbose117                )118        elif settings.chat_format == "moondream":119            assert settings.clip_model_path is not None, "clip model not found"120            if settings.hf_model_repo_id is not None:121                chat_handler = (122                    llama_cpp.llama_chat_format.MoondreamChatHandler.from_pretrained(123                        repo_id=settings.hf_model_repo_id,124                        filename=settings.clip_model_path,125                        verbose=settings.verbose,126                    )127                )128            else:129                chat_handler = llama_cpp.llama_chat_format.MoondreamChatHandler(130                    clip_model_path=settings.clip_model_path, verbose=settings.verbose131                )132        elif settings.chat_format == "nanollava":133            assert settings.clip_model_path is not None, "clip model not found"134            if settings.hf_model_repo_id is not None:135                chat_handler = (136                    llama_cpp.llama_chat_format.NanoLlavaChatHandler.from_pretrained(137                        repo_id=settings.hf_model_repo_id,138                        filename=settings.clip_model_path,139                        verbose=settings.verbose,140                    )141                )142            else:143                chat_handler = llama_cpp.llama_chat_format.NanoLlavaChatHandler(144                    clip_model_path=settings.clip_model_path, verbose=settings.verbose145                )146        elif settings.chat_format == "llama-3-vision-alpha":147            assert settings.clip_model_path is not None, "clip model not found"148            if settings.hf_model_repo_id is not None:149                chat_handler = (150                    llama_cpp.llama_chat_format.Llama3VisionAlpha.from_pretrained(151                        repo_id=settings.hf_model_repo_id,152                        filename=settings.clip_model_path,153                        verbose=settings.verbose,154                    )155                )156            else:157                chat_handler = llama_cpp.llama_chat_format.Llama3VisionAlpha(158                    clip_model_path=settings.clip_model_path, verbose=settings.verbose159                )160        elif settings.chat_format == "minicpm-v-2.6":161            assert settings.clip_model_path is not None, "clip model not found"162            if settings.hf_model_repo_id is not None:163                chat_handler = (164                    llama_cpp.llama_chat_format.MiniCPMv26ChatHandler.from_pretrained(165                        repo_id=settings.hf_model_repo_id,166                        filename=settings.clip_model_path,167                        verbose=settings.verbose,168                    )169                )170            else:171                chat_handler = llama_cpp.llama_chat_format.MiniCPMv26ChatHandler(172                    clip_model_path=settings.clip_model_path, verbose=settings.verbose173                )174        elif settings.chat_format == "qwen2.5-vl":175            assert settings.clip_model_path is not None, "clip model not found"176            if settings.hf_model_repo_id is not None:177                chat_handler = (178                    llama_cpp.llama_chat_format.Qwen25VLChatHandler.from_pretrained(179                        repo_id=settings.hf_model_repo_id,180                        filename=settings.clip_model_path,181                        verbose=settings.verbose,182                    )183                )184            else:185                chat_handler = llama_cpp.llama_chat_format.Qwen25VLChatHandler(186                    clip_model_path=settings.clip_model_path, verbose=settings.verbose187                )188        elif settings.chat_format == "hf-autotokenizer":189            assert (190                settings.hf_pretrained_model_name_or_path is not None191            ), "hf_pretrained_model_name_or_path must be set for hf-autotokenizer"192            chat_handler = (193                llama_cpp.llama_chat_format.hf_autotokenizer_to_chat_completion_handler(194                    settings.hf_pretrained_model_name_or_path195                )196            )197        elif settings.chat_format == "hf-tokenizer-config":198            assert (199                settings.hf_tokenizer_config_path is not None200            ), "hf_tokenizer_config_path must be set for hf-tokenizer-config"201            chat_handler = llama_cpp.llama_chat_format.hf_tokenizer_config_to_chat_completion_handler(202                json.load(open(settings.hf_tokenizer_config_path))203            )204 205        tokenizer: Optional[llama_cpp.BaseLlamaTokenizer] = None206        if settings.hf_pretrained_model_name_or_path is not None:207            tokenizer = llama_tokenizer.LlamaHFTokenizer.from_pretrained(208                settings.hf_pretrained_model_name_or_path209            )210 211        draft_model = None212        if settings.draft_model is not None:213            draft_model = llama_speculative.LlamaPromptLookupDecoding(214                num_pred_tokens=settings.draft_model_num_pred_tokens215            )216 217        kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]] = None218        if settings.kv_overrides is not None:219            assert isinstance(settings.kv_overrides, list)220            kv_overrides = {}221            for kv in settings.kv_overrides:222                key, value = kv.split("=")223                if ":" in value:224                    value_type, value = value.split(":")225                    if value_type == "bool":226                        kv_overrides[key] = value.lower() in ["true", "1"]227                    elif value_type == "int":228                        kv_overrides[key] = int(value)229                    elif value_type == "float":230                        kv_overrides[key] = float(value)231                    elif value_type == "str":232                        kv_overrides[key] = value233                    else:234                        raise ValueError(f"Unknown value type {value_type}")235 236        import functools237 238        kwargs = {}239 240        if settings.hf_model_repo_id is not None:241            create_fn = functools.partial(242                llama_cpp.Llama.from_pretrained,243                repo_id=settings.hf_model_repo_id,244                filename=settings.model,245            )246        else:247            create_fn = llama_cpp.Llama248            kwargs["model_path"] = settings.model249 250        _model = create_fn(251            **kwargs,252            # Model Params253            n_gpu_layers=settings.n_gpu_layers,254            split_mode=settings.split_mode,255            main_gpu=settings.main_gpu,256            tensor_split=settings.tensor_split,257            vocab_only=settings.vocab_only,258            use_mmap=settings.use_mmap,259            use_mlock=settings.use_mlock,260            kv_overrides=kv_overrides,261            rpc_servers=settings.rpc_servers,262            # Context Params263            seed=settings.seed,264            n_ctx=settings.n_ctx,265            n_batch=settings.n_batch,266            n_ubatch=settings.n_ubatch,267            n_threads=settings.n_threads,268            n_threads_batch=settings.n_threads_batch,269            rope_scaling_type=settings.rope_scaling_type,270            rope_freq_base=settings.rope_freq_base,271            rope_freq_scale=settings.rope_freq_scale,272            yarn_ext_factor=settings.yarn_ext_factor,273            yarn_attn_factor=settings.yarn_attn_factor,274            yarn_beta_fast=settings.yarn_beta_fast,275            yarn_beta_slow=settings.yarn_beta_slow,276            yarn_orig_ctx=settings.yarn_orig_ctx,277            mul_mat_q=settings.mul_mat_q,278            logits_all=settings.logits_all,279            embedding=settings.embedding,280            offload_kqv=settings.offload_kqv,281            flash_attn=settings.flash_attn,282            # Sampling Params283            last_n_tokens_size=settings.last_n_tokens_size,284            # LoRA Params285            lora_base=settings.lora_base,286            lora_path=settings.lora_path,287            # Backend Params288            numa=settings.numa,289            # Chat Format Params290            chat_format=settings.chat_format,291            chat_handler=chat_handler,292            # Speculative Decoding293            draft_model=draft_model,294            # KV Cache Quantization295            type_k=settings.type_k,296            type_v=settings.type_v,297            # Tokenizer298            tokenizer=tokenizer,299            # Misc300            verbose=settings.verbose,301        )302        if settings.cache:303            if settings.cache_type == "disk":304                if settings.verbose:305                    print(f"Using disk cache with size {settings.cache_size}")306                cache = llama_cpp.LlamaDiskCache(capacity_bytes=settings.cache_size)307            else:308                if settings.verbose:309                    print(f"Using ram cache with size {settings.cache_size}")310                cache = llama_cpp.LlamaRAMCache(capacity_bytes=settings.cache_size)311            _model.set_cache(cache)312        return _model313