stubs for mlx, vllm, outlines and transformers

Author: lars20070

pydantic_ai_slim/pydantic_ai/models/outlines.py

@@ -60,7 +60,7 @@
     )
     from outlines.models.vllm_offline import (
         VLLMOffline,
-        from_vllm_offline,  # pyright: ignore[reportUnknownVariableType]
+        from_vllm_offline,
     )
     from outlines.types.dsl import JsonSchema
     from PIL import Image as PILImage
@@ -393,7 +393,7 @@ def _format_vllm_offline_inference_kwargs(  # pragma: no cover
         self, model_settings: dict[str, Any]
     ) -> dict[str, Any]:
         """Select the model settings supported by the vLLMOffline model."""
-        from vllm.sampling_params import SamplingParams  # pyright: ignore
+        from vllm.sampling_params import SamplingParams
 
         supported_args = [
             'max_tokens',

pyproject.toml

@@ -188,6 +188,7 @@ quote-style = "single"
 "docs/**/*.py" = ["D"]
 
 [tool.pyright]
+stubPath = "stubs"
 pythonVersion = "3.12"
 typeCheckingMode = "strict"
 reportMissingTypeStubs = false
@@ -217,6 +218,7 @@ exclude = [
 [tool.mypy]
 files = "tests/typed_agent.py,tests/typed_graph.py"
 strict = true
+mypy_path = "stubs"
 
 [tool.pytest.ini_options]
 testpaths = ["tests", "docs/.hooks"]

stubs/mlx/__init__.pyi

@@ -0,0 +1,4 @@
+from typing import Any
+
+# mlx is imported as a package, primarily for mlx.nn
+__all__: list[str]

stubs/mlx/nn.pyi

@@ -0,0 +1,5 @@
+from typing import Any
+
+
+class Module:
+    ...

stubs/mlx_lm.pyi

@@ -0,0 +1,7 @@
+from typing import Any
+
+from mlx.nn import Module
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+def load(model_path: str) -> tuple[Module, PreTrainedTokenizer]: ...
+def generate_step(*args: Any, **kwargs: Any) -> Any: ...

stubs/outlines/models/base.pyi

@@ -0,0 +1,11 @@
+from typing import Any, AsyncIterable, Iterable
+
+
+class Model:
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+    def stream(self, *args: Any, **kwargs: Any) -> Iterable[Any]: ...
+
+
+class AsyncModel(Model):
+    async def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+    def stream(self, *args: Any, **kwargs: Any) -> AsyncIterable[Any]: ...

stubs/outlines/models/mlxlm.pyi

@@ -0,0 +1,12 @@
+from typing import Any
+from outlines.models.base import Model
+
+from mlx.nn import Module
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+
+class MLXLM(Model):
+    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+
+
+def from_mlxlm(model: Module, tokenizer: PreTrainedTokenizer) -> MLXLM: ...

stubs/outlines/models/transformers.pyi

@@ -0,0 +1,23 @@
+from typing import Any
+from outlines.models.base import Model
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+
+class Transformers(Model):
+    ...
+
+
+class TransformersMultiModal(Model):
+    ...
+
+
+def from_transformers(
+    model: PreTrainedModel,
+    tokenizer_or_processor: PreTrainedTokenizer | ProcessorMixin,
+    *,
+    device_dtype: Any = None,
+) -> Transformers | TransformersMultiModal: ...
+

stubs/outlines/models/vllm_offline.pyi

@@ -0,0 +1,9 @@
+from typing import Any
+from outlines.models.base import Model
+
+
+class VLLMOffline(Model):
+    ...
+
+
+def from_vllm_offline(model: Any) -> VLLMOffline: ...

stubs/transformers/__init__.pyi

@@ -0,0 +1,29 @@
+from typing import Any
+
+from . import modeling_utils, processing_utils, tokenization_utils
+from .modeling_utils import PreTrainedModel
+from .processing_utils import ProcessorMixin
+from .tokenization_utils import PreTrainedTokenizer
+
+
+class AutoModelForCausalLM(PreTrainedModel):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> PreTrainedModel: ...
+
+
+class AutoTokenizer(PreTrainedTokenizer):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> PreTrainedTokenizer: ...
+
+
+class AutoProcessor(ProcessorMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> ProcessorMixin: ...
+
+
+class LlavaForConditionalGeneration(PreTrainedModel):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> PreTrainedModel: ...
+
+
+def from_pretrained(*args: Any, **kwargs: Any) -> Any: ...