chore(utils): centralize context length helper

Author: qinxuye

xinference/model/llm/mlx/core.py

@@ -56,6 +56,7 @@
     QWEN_TOOL_CALL_FAMILY,
     ChatModelMixin,
     generate_completion_chunk,
+    get_context_length_from_config,
 )
 
 logger = logging.getLogger(__name__)
@@ -436,23 +437,6 @@ class PromptCache:
     tokens: List[int] = field(default_factory=list)
 
 
-def get_context_length(config: dict) -> int:
-    """Get the context length of a model from model config."""
-    if config.get("max_sequence_length") is not None:
-        max_sequence_length = config["max_sequence_length"]
-    else:
-        max_sequence_length = 2048
-    if config.get("seq_length") is not None:
-        seq_length = config["seq_length"]
-    else:
-        seq_length = 2048
-    if config.get("max_position_embeddings") is not None:
-        max_position_embeddings = config["max_position_embeddings"]
-    else:
-        max_position_embeddings = 2048
-    return max(max_sequence_length, seq_length, max_position_embeddings)
-
-
 class MLXModel(LLM, ChatModelMixin):
     _rank_to_addresses: Optional[Dict[int, str]]
     allow_batch: bool = False
@@ -752,7 +736,7 @@ def wait_for_load(self):
         # get context length
         config = load_config(Path(self.model_path))
         config.update(self._model_config)
-        self._context_length = get_context_length(config)
+        self._context_length = get_context_length_from_config(config)
 
         # Update allow_batch based on distributed inference
         # Only enable continuous batching for non-distributed inference (single worker)
@@ -1418,7 +1402,7 @@ def load(self):
         # get context length
         config = load_config(Path(self.model_path))
         config.update(self._model_config)
-        self._context_length = get_context_length(config)
+        self._context_length = get_context_length_from_config(config)
 
     def _generate_stream_inner(self, **kwargs):
         import mlx.core as mx

xinference/model/llm/transformers/utils.py

@@ -37,6 +37,7 @@
     max_tokens_field,
 )
 from ...scheduler.request import InferenceRequest
+from ..utils import get_context_length_from_config
 
 if TYPE_CHECKING:
     from ...llm.transformers.core import PytorchModel
@@ -46,25 +47,7 @@
 
 def get_context_length(config) -> int:
     """Get the context length of a model from a huggingface model config."""
-    if (
-        hasattr(config, "max_sequence_length")
-        and config.max_sequence_length is not None
-    ):
-        max_sequence_length = config.max_sequence_length
-    else:
-        max_sequence_length = 2048
-    if hasattr(config, "seq_length") and config.seq_length is not None:
-        seq_length = config.seq_length
-    else:
-        seq_length = 2048
-    if (
-        hasattr(config, "max_position_embeddings")
-        and config.max_position_embeddings is not None
-    ):
-        max_position_embeddings = config.max_position_embeddings
-    else:
-        max_position_embeddings = 2048
-    return max(max_sequence_length, seq_length, max_position_embeddings)
+    return get_context_length_from_config(config)
 
 
 def prepare_logits_processor(

xinference/model/llm/utils.py

@@ -56,6 +56,48 @@
 logger = logging.getLogger(__name__)
 
 
+_CONTEXT_LENGTH_KEYS: Tuple[str, ...] = (
+    "max_sequence_length",
+    "seq_length",
+    "max_position_embeddings",
+    "sliding_window",
+)
+
+
+def _get_config_value(config: Union[dict, Any], key: str) -> Any:
+    if isinstance(config, dict):
+        return config.get(key)
+    return getattr(config, key, None)
+
+
+def _collect_context_length_candidates(
+    config: Union[dict, Any], nested_attrs: Iterable[str]
+) -> List[int]:
+    candidates: List[int] = []
+    for key in _CONTEXT_LENGTH_KEYS:
+        value = _get_config_value(config, key)
+        if value is not None:
+            candidates.append(value)
+    for nested_attr in nested_attrs:
+        nested = _get_config_value(config, nested_attr)
+        if nested is not None:
+            candidates.extend(_collect_context_length_candidates(nested, nested_attrs))
+    return candidates
+
+
+def get_context_length_from_config(
+    config: Union[dict, Any], nested_attrs: Iterable[str] = ("text_config",)
+) -> int:
+    """
+    Determine a reasonable context length from model config dictionaries or
+    HuggingFace config objects.
+    """
+    candidates = _collect_context_length_candidates(config, nested_attrs)
+    if not candidates:
+        return 2048
+    return max(candidates)
+
+
 QWEN_TOOL_CALL_FAMILY = [
     "qwen1.5-chat",
     "qwen1.5-moe-chat",