xinference 1.7.1__py3-none-any.whl → 1.8.0__py3-none-any.whl

xinference/model/rerank/custom.py

@@ -12,98 +12,52 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-import os
-from threading import Lock
 from typing import List, Literal, Optional
 
-from ...constants import XINFERENCE_CACHE_DIR, XINFERENCE_MODEL_DIR
-from .core import RerankModelSpec
+from ..custom import ModelRegistry
+from .core import RerankModelFamilyV2
 
 logger = logging.getLogger(__name__)
 
 
-UD_RERANK_LOCK = Lock()
-
-
-class CustomRerankModelSpec(RerankModelSpec):
+class CustomRerankModelFamilyV2(RerankModelFamilyV2):
+    version: Literal[2] = 2
     model_id: Optional[str]  # type: ignore
     model_revision: Optional[str]  # type: ignore
     model_uri: Optional[str]
     model_type: Literal["rerank"] = "rerank"  # for frontend
 
 
-UD_RERANKS: List[CustomRerankModelSpec] = []
+UD_RERANKS: List[CustomRerankModelFamilyV2] = []
 
 
-def get_user_defined_reranks() -> List[CustomRerankModelSpec]:
-    with UD_RERANK_LOCK:
-        return UD_RERANKS.copy()
+class RerankModelRegistry(ModelRegistry):
+    model_type = "rerank"
 
+    def __init__(self):
+        from . import BUILTIN_RERANK_MODELS
 
-def register_rerank(model_spec: CustomRerankModelSpec, persist: bool):
-    from ...constants import XINFERENCE_MODEL_DIR
-    from ..utils import is_valid_model_name, is_valid_model_uri
-    from . import BUILTIN_RERANK_MODELS, MODELSCOPE_RERANK_MODELS
+        super().__init__()
+        self.models = UD_RERANKS
+        self.builtin_models = list(BUILTIN_RERANK_MODELS.keys())
 
-    if not is_valid_model_name(model_spec.model_name):
-        raise ValueError(f"Invalid model name {model_spec.model_name}.")
 
-    model_uri = model_spec.model_uri
-    if model_uri and not is_valid_model_uri(model_uri):
-        raise ValueError(f"Invalid model URI {model_uri}.")
+def get_user_defined_reranks() -> List[CustomRerankModelFamilyV2]:
+    from ..custom import RegistryManager
 
-    with UD_RERANK_LOCK:
-        for model_name in (
-            list(BUILTIN_RERANK_MODELS.keys())
-            + list(MODELSCOPE_RERANK_MODELS.keys())
-            + [spec.model_name for spec in UD_RERANKS]
-        ):
-            if model_spec.model_name == model_name:
-                raise ValueError(
-                    f"Model name conflicts with existing model {model_spec.model_name}"
-                )
+    registry = RegistryManager.get_registry("rerank")
+    return registry.get_custom_models()
 
-        UD_RERANKS.append(model_spec)
 
-    if persist:
-        persist_path = os.path.join(
-            XINFERENCE_MODEL_DIR, "rerank", f"{model_spec.model_name}.json"
-        )
-        os.makedirs(os.path.dirname(persist_path), exist_ok=True)
-        with open(persist_path, mode="w") as fd:
-            fd.write(model_spec.json())
+def register_rerank(model_spec: CustomRerankModelFamilyV2, persist: bool):
+    from ..custom import RegistryManager
+
+    registry = RegistryManager.get_registry("rerank")
+    registry.register(model_spec, persist)
 
 
 def unregister_rerank(model_name: str, raise_error: bool = True):
-    with UD_RERANK_LOCK:
-        model_spec = None
-        for i, f in enumerate(UD_RERANKS):
-            if f.model_name == model_name:
-                model_spec = f
-                break
-        if model_spec:
-            UD_RERANKS.remove(model_spec)
-
-            persist_path = os.path.join(
-                XINFERENCE_MODEL_DIR, "rerank", f"{model_spec.model_name}.json"
-            )
-            if os.path.exists(persist_path):
-                os.remove(persist_path)
-
-            cache_dir = os.path.join(XINFERENCE_CACHE_DIR, model_spec.model_name)
-            if os.path.exists(cache_dir):
-                logger.warning(
-                    f"Remove the cache of user-defined model {model_spec.model_name}. "
-                    f"Cache directory: {cache_dir}"
-                )
-                if os.path.islink(cache_dir):
-                    os.remove(cache_dir)
-                else:
-                    logger.warning(
-                        f"Cache directory is not a soft link, please remove it manually."
-                    )
-        else:
-            if raise_error:
-                raise ValueError(f"Model {model_name} not found")
-            else:
-                logger.warning(f"Custom rerank model {model_name} not found")
+    from ..custom import RegistryManager
+
+    registry = RegistryManager.get_registry("rerank")
+    registry.unregister(model_name, raise_error)

xinference/model/rerank/model_spec.json

@@ -1,90 +1,215 @@
 [
   {
+    "version": 2,
     "model_name": "bge-reranker-large",
     "type": "normal",
-    "language": ["en", "zh"],
+    "language": [
+      "en",
+      "zh"
+    ],
     "max_tokens": 512,
-    "model_id": "BAAI/bge-reranker-large",
-    "model_revision": "27c9168d479987529781de8474dff94d69beca11"
+    "model_src": {
+      "huggingface": {
+        "model_id": "BAAI/bge-reranker-large",
+        "model_revision": "27c9168d479987529781de8474dff94d69beca11"
+      },
+      "modelscope": {
+        "model_id": "Xorbits/bge-reranker-large",
+        "model_revision": "v0.0.1"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "bge-reranker-base",
     "type": "normal",
-    "language": ["en", "zh"],
+    "language": [
+      "en",
+      "zh"
+    ],
     "max_tokens": 512,
-    "model_id": "BAAI/bge-reranker-base",
-    "model_revision": "465b4b7ddf2be0a020c8ad6e525b9bb1dbb708ae"
+    "model_src": {
+      "huggingface": {
+        "model_id": "BAAI/bge-reranker-base",
+        "model_revision": "465b4b7ddf2be0a020c8ad6e525b9bb1dbb708ae"
+      },
+      "modelscope": {
+        "model_id": "Xorbits/bge-reranker-base",
+        "model_revision": "v0.0.1"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "bce-reranker-base_v1",
     "type": "normal",
-    "language": ["en", "zh"],
+    "language": [
+      "en",
+      "zh"
+    ],
     "max_tokens": 512,
-    "model_id": "maidalun1020/bce-reranker-base_v1",
-    "model_revision": "eaa31a577a0574e87a08959bd229ca14ce1b5496"
+    "model_src": {
+      "huggingface": {
+        "model_id": "maidalun1020/bce-reranker-base_v1",
+        "model_revision": "eaa31a577a0574e87a08959bd229ca14ce1b5496"
+      },
+      "modelscope": {
+        "model_id": "maidalun/bce-reranker-base_v1",
+        "model_revision": "v0.0.1"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "bge-reranker-v2-m3",
     "type": "normal",
-    "language": ["en", "zh", "multilingual"],
+    "language": [
+      "en",
+      "zh",
+      "multilingual"
+    ],
     "max_tokens": 8192,
-    "model_id": "BAAI/bge-reranker-v2-m3",
-    "model_revision": "12e974610ba9083ed95f3edf08d7e899581f4de4"
+    "model_src": {
+      "huggingface": {
+        "model_id": "BAAI/bge-reranker-v2-m3",
+        "model_revision": "12e974610ba9083ed95f3edf08d7e899581f4de4"
+      },
+      "modelscope": {
+        "model_id": "AI-ModelScope/bge-reranker-v2-m3"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "bge-reranker-v2-gemma",
     "type": "LLM-based",
-    "language": ["en", "zh", "multilingual"],
+    "language": [
+      "en",
+      "zh",
+      "multilingual"
+    ],
     "max_tokens": 8192,
-    "model_id": "BAAI/bge-reranker-v2-gemma",
-    "model_revision": "1787044f8b6fb740a9de4557c3a12377f84d9e17"
+    "model_src": {
+      "huggingface": {
+        "model_id": "BAAI/bge-reranker-v2-gemma",
+        "model_revision": "1787044f8b6fb740a9de4557c3a12377f84d9e17"
+      },
+      "modelscope": {
+        "model_id": "AI-ModelScope/bge-reranker-v2-gemma"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "bge-reranker-v2-minicpm-layerwise",
     "type": "LLM-based layerwise",
-    "language": ["en", "zh", "multilingual"],
+    "language": [
+      "en",
+      "zh",
+      "multilingual"
+    ],
     "max_tokens": 2048,
-    "model_id": "BAAI/bge-reranker-v2-minicpm-layerwise",
-    "model_revision": "47b5332b296c4d8cb6ee2c60502cc62a0d708881"
+    "model_src": {
+      "huggingface": {
+        "model_id": "BAAI/bge-reranker-v2-minicpm-layerwise",
+        "model_revision": "47b5332b296c4d8cb6ee2c60502cc62a0d708881"
+      },
+      "modelscope": {
+        "model_id": "mirror013/bge-reranker-v2-minicpm-layerwise"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "jina-reranker-v2",
     "type": "normal",
-    "language": ["en", "zh", "multilingual"],
+    "language": [
+      "en",
+      "zh",
+      "multilingual"
+    ],
     "max_tokens": 1024,
-    "model_id": "jinaai/jina-reranker-v2-base-multilingual",
-    "model_revision": "298e48cada4a9318650d7fbd795f63827f884087"
+    "model_src": {
+      "huggingface": {
+        "model_id": "jinaai/jina-reranker-v2-base-multilingual",
+        "model_revision": "298e48cada4a9318650d7fbd795f63827f884087"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "minicpm-reranker",
     "type": "normal",
-    "language": ["en", "zh"],
+    "language": [
+      "en",
+      "zh"
+    ],
     "max_tokens": 1024,
-    "model_id": "openbmb/MiniCPM-Reranker",
-    "model_revision": "5d2fd7345b6444c89d4c0fa59c92272888f3f2d0"
+    "model_src": {
+      "huggingface": {
+        "model_id": "openbmb/MiniCPM-Reranker",
+        "model_revision": "5d2fd7345b6444c89d4c0fa59c92272888f3f2d0"
+      },
+      "modelscope": {
+        "model_id": "OpenBMB/MiniCPM-Reranker"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "Qwen3-Reranker-0.6B",
     "type": "normal",
-    "language": ["en", "zh"],
+    "language": [
+      "en",
+      "zh"
+    ],
     "max_tokens": 32768,
-    "model_id": "Qwen/Qwen3-Reranker-0.6B",
-    "model_revision": "6e9e69830b95c52b5fd889b7690dda3329508de3"
+    "model_src": {
+      "huggingface": {
+        "model_id": "Qwen/Qwen3-Reranker-0.6B",
+        "model_revision": "6e9e69830b95c52b5fd889b7690dda3329508de3"
+      },
+      "modelscope": {
+        "model_id": "Qwen/Qwen3-Reranker-0.6B"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "Qwen3-Reranker-4B",
     "type": "normal",
-    "language": ["en", "zh"],
+    "language": [
+      "en",
+      "zh"
+    ],
     "max_tokens": 32768,
-    "model_id": "Qwen/Qwen3-Reranker-4B",
-    "model_revision": "f16fc5d5d2b9b1d0db8280929242745d79794ef5"
+    "model_src": {
+      "huggingface": {
+        "model_id": "Qwen/Qwen3-Reranker-4B",
+        "model_revision": "f16fc5d5d2b9b1d0db8280929242745d79794ef5"
+      },
+      "modelscope": {
+        "model_id": "Qwen/Qwen3-Reranker-4B"
+      }
+    }
   },
   {
+    "version": 2,
     "model_name": "Qwen3-Reranker-8B",
     "type": "normal",
-    "language": ["en", "zh"],
+    "language": [
+      "en",
+      "zh"
+    ],
     "max_tokens": 32768,
-    "model_id": "Qwen/Qwen3-Reranker-8B",
-    "model_revision": "5fa94080caafeaa45a15d11f969d7978e087a3db"
+    "model_src": {
+      "huggingface": {
+        "model_id": "Qwen/Qwen3-Reranker-8B",
+        "model_revision": "5fa94080caafeaa45a15d11f969d7978e087a3db"
+      },
+      "modelscope": {
+        "model_id": "Qwen/Qwen3-Reranker-8B"
+      }
+    }
   }
 ]

xinference/model/rerank/utils.py

@@ -14,10 +14,10 @@
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    from .core import RerankModelSpec
+    from .core import RerankModelFamilyV2
 
 
-def get_model_version(rerank_model: "RerankModelSpec") -> str:
+def get_model_version(rerank_model: "RerankModelFamilyV2") -> str:
     return rerank_model.model_name
 
 

xinference/model/utils.py

@@ -18,10 +18,22 @@ import logging
 import os
 import random
 import threading
+from abc import ABC, abstractmethod
 from copy import deepcopy
 from json import JSONDecodeError
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
 
 import huggingface_hub
 import numpy as np
@@ -36,6 +48,10 @@ from ..constants import (
 from ..device_utils import get_available_device, is_device_available
 from .core import CacheableModelSpec
 
+if TYPE_CHECKING:
+    from .embedding.core import LlamaCppEmbeddingSpecV1
+    from .llm.llm_family import LlamaCppLLMSpecV2
+
 logger = logging.getLogger(__name__)
 IS_NEW_HUGGINGFACE_HUB: bool = huggingface_hub.__version__ >= "0.23.0"
 
@@ -262,59 +278,6 @@ def cache_from_uri(model_spec: CacheableModelSpec) -> str:
         raise ValueError(f"Unsupported URL scheme: {src_scheme}")
 
 
-def cache(model_spec: CacheableModelSpec, model_description_type: type):
-    if (
-        hasattr(model_spec, "model_uri")
-        and getattr(model_spec, "model_uri", None) is not None
-    ):
-        logger.info(f"Model caching from URI: {model_spec.model_uri}")
-        return cache_from_uri(model_spec=model_spec)
-
-    cache_dir = os.path.realpath(
-        os.path.join(XINFERENCE_CACHE_DIR, model_spec.model_name)
-    )
-    if not os.path.exists(cache_dir):
-        os.makedirs(cache_dir, exist_ok=True)
-    meta_path = os.path.join(cache_dir, "__valid_download")
-    if valid_model_revision(meta_path, model_spec.model_revision, model_spec.model_hub):
-        return cache_dir
-
-    from_modelscope: bool = model_spec.model_hub == "modelscope"
-    if from_modelscope:
-        from modelscope.hub.snapshot_download import snapshot_download as ms_download
-
-        download_dir = retry_download(
-            ms_download,
-            model_spec.model_name,
-            None,
-            model_spec.model_id,
-            revision=model_spec.model_revision,
-        )
-        create_symlink(download_dir, cache_dir)
-    else:
-        from huggingface_hub import snapshot_download as hf_download
-
-        use_symlinks = {}
-        if not IS_NEW_HUGGINGFACE_HUB:
-            use_symlinks = {"local_dir_use_symlinks": True, "local_dir": cache_dir}
-        download_dir = retry_download(
-            hf_download,
-            model_spec.model_name,
-            None,
-            model_spec.model_id,
-            revision=model_spec.model_revision,
-            **use_symlinks,
-        )
-        if IS_NEW_HUGGINGFACE_HUB:
-            create_symlink(download_dir, cache_dir)
-    with open(meta_path, "w") as f:
-        import json
-
-        desc = model_description_type(None, None, model_spec)
-        json.dump(desc.to_dict(), f)
-    return cache_dir
-
-
 def select_device(device):
     try:
         import torch  # noqa: F401
@@ -497,3 +460,101 @@ def get_engine_params_by_name(
             f"Cannot support model_engine for {model_type}, "
             f"only available for LLM, embedding"
         )
+
+
+def generate_model_file_names_with_quantization_parts(
+    model_spec: Union["LlamaCppLLMSpecV2", "LlamaCppEmbeddingSpecV1"],
+    multimodal_projector: Optional[str] = None,
+) -> Tuple[List[str], str, bool]:
+    file_names = []
+    final_file_name = model_spec.model_file_name_template.format(
+        quantization=model_spec.quantization
+    )
+    need_merge = False
+
+    if (
+        model_spec.quantization_parts is None
+        or model_spec.quantization not in model_spec.quantization_parts
+    ):
+        file_names.append(final_file_name)
+    elif (
+        model_spec.quantization is not None
+        and model_spec.quantization in model_spec.quantization_parts
+    ):
+        parts = model_spec.quantization_parts[model_spec.quantization]
+        need_merge = True
+
+        logger.info(
+            f"Model {model_spec.model_id} {model_spec.model_format} {model_spec.quantization} has {len(parts)} parts."
+        )
+
+        if model_spec.model_file_name_split_template is None:
+            raise ValueError(
+                f"No model_file_name_split_template for model spec {model_spec.model_id}"
+            )
+
+        for part in parts:
+            file_name = model_spec.model_file_name_split_template.format(
+                quantization=model_spec.quantization, part=part
+            )
+            file_names.append(file_name)
+    if multimodal_projector:
+        file_names.append(multimodal_projector)
+
+    return file_names, final_file_name, need_merge
+
+
+def merge_cached_files(
+    cache_dir: str, input_file_names: List[str], output_file_name: str
+):
+    # now llama.cpp can find the gguf parts automatically
+    # we only need to provide the first part
+    # thus we create the symlink to the first part
+    symlink_local_file(
+        os.path.join(cache_dir, input_file_names[0]), cache_dir, output_file_name
+    )
+
+    logger.info(f"Merge complete.")
+
+
+def flatten_model_src(input_json: dict):
+    flattened = []
+    base_info = {key: value for key, value in input_json.items() if key != "model_src"}
+    for model_hub, hub_info in input_json["model_src"].items():
+        record = base_info.copy()
+        hub_info.pop("model_hub", None)
+        record.update(hub_info)
+        record["model_hub"] = model_hub
+        flattened.append(record)
+    return flattened
+
+
+def flatten_quantizations(input_json: dict):
+    flattened = []
+
+    base_info = {key: value for key, value in input_json.items() if key != "model_src"}
+
+    for model_hub, hub_info in input_json["model_src"].items():
+        quantizations = hub_info["quantizations"]
+
+        for quant in quantizations:
+            record = base_info.copy()
+            record["model_hub"] = model_hub
+            record["quantization"] = quant
+
+            for key, value in hub_info.items():
+                if key != "quantizations":
+                    record[key] = value
+
+            flattened.append(record)
+    return flattened
+
+
+class ModelInstanceInfoMixin(ABC):
+    @abstractmethod
+    def to_description(self):
+        """"""
+
+    @abstractmethod
+    def to_version_info(self):
+        """"""

xinference/model/video/__init__.py

@@ -15,40 +15,36 @@
 import codecs
 import json
 import os
-from itertools import chain
 
+from ..utils import flatten_model_src
 from .core import (
     BUILTIN_VIDEO_MODELS,
-    MODEL_NAME_TO_REVISION,
-    MODELSCOPE_VIDEO_MODELS,
     VIDEO_MODEL_DESCRIPTIONS,
-    VideoModelFamilyV1,
+    VideoModelFamilyV2,
     generate_video_description,
-    get_cache_status,
     get_video_model_descriptions,
 )
 
 
 def _install():
     load_model_family_from_json("model_spec.json", BUILTIN_VIDEO_MODELS)
-    load_model_family_from_json("model_spec_modelscope.json", MODELSCOPE_VIDEO_MODELS)
 
     # register model description
-    for model_name, model_spec in chain(
-        MODELSCOPE_VIDEO_MODELS.items(), BUILTIN_VIDEO_MODELS.items()
-    ):
+    for model_name, model_specs in BUILTIN_VIDEO_MODELS.items():
+        model_spec = [x for x in model_specs if x.model_hub == "huggingface"][0]
         VIDEO_MODEL_DESCRIPTIONS.update(generate_video_description(model_spec))
 
 
 def load_model_family_from_json(json_filename, target_families):
     json_path = os.path.join(os.path.dirname(__file__), json_filename)
-    target_families.update(
-        dict(
-            (spec["model_name"], VideoModelFamilyV1(**spec))
-            for spec in json.load(codecs.open(json_path, "r", encoding="utf-8"))
-        )
-    )
-    for model_name, model_spec in target_families.items():
-        MODEL_NAME_TO_REVISION[model_name].append(model_spec.model_revision)
+    flattened_model_specs = []
+    for spec in json.load(codecs.open(json_path, "r", encoding="utf-8")):
+        flattened_model_specs.extend(flatten_model_src(spec))
+
+    for spec in flattened_model_specs:
+        if spec["model_name"] not in target_families:
+            target_families[spec["model_name"]] = [VideoModelFamilyV2(**spec)]
+        else:
+            target_families[spec["model_name"]].append(VideoModelFamilyV2(**spec))
 
     del json_path