Add a __call__ method to models

outlines/models/__init__.py

@@ -9,6 +9,7 @@
 from typing import Union
 
 from .anthropic import Anthropic
+from .base import Model
 from .exllamav2 import ExLlamaV2Model, exl2
 from .gemini import Gemini
 from .llamacpp import LlamaCpp

outlines/models/anthropic.py

@@ -2,6 +2,7 @@
 from functools import singledispatchmethod
 from typing import Union
 
+from outlines.models import Model
 from outlines.prompts import Vision
 
 __all__ = ["Anthropic"]
@@ -63,7 +64,7 @@ def format_vision_input(self, model_input: Vision):
         }
 
 
-class Anthropic(AnthropicBase):
+class Anthropic(Model, AnthropicBase):
     def __init__(self, model_name: str, *args, **kwargs):
         from anthropic import Anthropic
 

outlines/models/base.py

@@ -0,0 +1,10 @@
+from abc import ABC
+
+
+class Model(ABC):
+    """Base class for all models."""
+
+    def __call__(self, model_input, output_type=None, **inference_kwargs):
+        from outlines.generate import Generator
+
+        return Generator(self, output_type)(model_input, **inference_kwargs)

outlines/models/gemini.py

@@ -7,6 +7,7 @@
 from pydantic import BaseModel
 from typing_extensions import _TypedDictMeta  # type: ignore
 
+from outlines.models import Model
 from outlines.prompts import Vision
 from outlines.types import Choice, Json, List
 
@@ -91,7 +92,7 @@ def format_enum_output_type(self, output_type):
         }
 
 
-class Gemini(GeminiBase):
+class Gemini(Model, GeminiBase):
     def __init__(self, model_name: str, *args, **kwargs):
         import google.generativeai as genai
 

outlines/models/llamacpp.py

@@ -2,6 +2,7 @@
 import warnings
 from typing import TYPE_CHECKING, Dict, Iterator, List, Set, Tuple, Union
 
+from outlines.models import Model
 from outlines.models.tokenizer import Tokenizer
 
 if TYPE_CHECKING:
@@ -93,7 +94,7 @@ def __setstate__(self, state):
         raise NotImplementedError("Cannot load a pickled llamacpp tokenizer")
 
 
-class LlamaCpp:
+class LlamaCpp(Model):
     """Wraps a model provided by the `llama-cpp-python` library."""
 
     def __init__(self, model_path: Union[str, "Llama"], **kwargs):

outlines/models/openai.py

@@ -5,6 +5,7 @@
 
 from pydantic import BaseModel
 
+from outlines.models import Model
 from outlines.prompts import Vision
 from outlines.types import Json
 
@@ -114,7 +115,7 @@ def format_json_output_type(self, output_type: Json):
         }
 
 
-class OpenAI(OpenAIBase):
+class OpenAI(Model, OpenAIBase):
     """Thin wrapper around the `openai.OpenAI` client.
 
     This wrapper is used to convert the input and output types specified by the