1+ # Copyright 2022-2025 XProbe Inc.
2+ #
3+ # Licensed under the Apache License, Version 2.0 (the "License");
4+ # you may not use this file except in compliance with the License.
5+ # You may obtain a copy of the License at
6+ #
7+ # http://www.apache.org/licenses/LICENSE-2.0
8+ #
9+ # Unless required by applicable law or agreed to in writing, software
10+ # distributed under the License is distributed on an "AS IS" BASIS,
11+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ # See the License for the specific language governing permissions and
13+ # limitations under the License.
14+
15+ import json
16+ import logging
17+ import os
18+ from typing import List
19+
20+ logger = logging .getLogger (__name__ )
21+
22+
23+ class BuiltinEmbeddingModelRegistry :
24+ """
25+ Registry for built-in embedding models downloaded from official model hub.
26+
27+ These models are treated as built-in models.
28+ They are stored in ~/.xinference/model/v2/builtin/embedding/ directory.
29+ """
30+
31+ def __init__ (self ):
32+ from ...constants import XINFERENCE_MODEL_DIR
33+
34+ self .builtin_dir = os .path .join (XINFERENCE_MODEL_DIR , "v2" , "builtin" , "embedding" )
35+ os .makedirs (self .builtin_dir , exist_ok = True )
36+
37+ def get_builtin_models (self ) -> List :
38+ """Load all built-in embedding models from the builtin directory."""
39+ from .custom import EmbeddingModelFamilyV2
40+
41+ models = []
42+
43+ if not os .path .exists (self .builtin_dir ):
44+ return models
45+
46+ for filename in os .listdir (self .builtin_dir ):
47+ if filename .endswith (".json" ):
48+ file_path = os .path .join (self .builtin_dir , filename )
49+ try :
50+ with open (file_path , "r" , encoding = "utf-8" ) as f :
51+ model_data = json .load (f )
52+
53+ # Parse using EmbeddingFamilyV2
54+ model = EmbeddingModelFamilyV2 .parse_obj (model_data )
55+ models .append (model )
56+ logger .info (f"Loaded built-in embedding model: { model .model_name } )
57+
58+ except Exception as e :
59+ logger .warning (f"Failed to load built-in model from { filename } { e } )
60+
61+ return models
62+
63+ def register_builtin_model (self , model ) -> None :
64+ """Register a built-in embedding model by saving it to the builtin directory."""
65+ persist_path = os .path .join (self .builtin_dir , f"{ model .model_name } )
66+
67+ try :
68+ with open (persist_path , "w" , encoding = "utf-8" ) as f :
69+ f .write (model .json (exclude_none = True ))
70+ logger .info (f"Registered built-in embedding model: { model .model_name } )
71+ except Exception as e :
72+ logger .error (f"Failed to register built-in model { model .model_name } { e } )
73+ raise
74+
75+ def unregister_builtin_model (self , model_name : str ) -> None :
76+ """Unregister a built-in embedding model by removing its JSON file."""
77+ persist_path = os .path .join (self .builtin_dir , f"{ model_name } )
78+
79+ if os .path .exists (persist_path ):
80+ os .remove (persist_path )
81+ logger .info (f"Unregistered built-in embedding model: { model_name } )
82+ else :
83+ logger .warning (f"Built-in model file not found: { persist_path } )
84+
85+
86+ # Global registry instance
87+ _builtin_registry = None
88+
89+
90+ def get_builtin_embedding_registry () -> BuiltinEmbeddingModelRegistry :
91+ """Get the global built-in embedding model registry instance."""
92+ global _builtin_registry
93+ if _builtin_registry is None :
94+ _builtin_registry = BuiltinEmbeddingModelRegistry ()
95+ return _builtin_registry
96+
97+
98+ def get_builtin_embedding_families () -> List :
99+ """Get all built-in embedding model families."""
100+ return get_builtin_embedding_registry ().get_builtin_models ()
101+
102+
103+ def register_builtin_embedding (embedding_family ) -> None :
104+ """Register a built-in embedding model family."""
105+ return get_builtin_embedding_registry ().register_builtin_model (embedding_family )
106+
107+
108+ def unregister_builtin_embedding (model_name : str ) -> None :
109+ """Unregister a built-in embedding model family."""
110+ return get_builtin_embedding_registry ().unregister_builtin_model (model_name )
0 commit comments