Implement tools and outputs for the Transformers model

Author: RobinPicard

outlines/models/transformers.py

@@ -9,7 +9,9 @@
 from outlines.inputs import Audio, Chat, Image, Video
 from outlines.models.base import Model, ModelTypeAdapter
 from outlines.models.tokenizer import Tokenizer
+from outlines.outputs import Output
 from outlines.processors import OutlinesLogitsProcessor
+from outlines.tools import ToolDef
 
 if TYPE_CHECKING:
     import torch
@@ -173,7 +175,7 @@ def format_chat_input(self, model_input: Chat) -> str:
 
     def format_output_type(
         self,
-        output_type: Optional[OutlinesLogitsProcessor] = None,
+        output_type: Optional[OutlinesLogitsProcessor],
     ) -> Optional["LogitsProcessorList"]:
         """Generate the logits processor argument to pass to the model.
 
@@ -194,6 +196,13 @@ def format_output_type(
             return LogitsProcessorList([output_type])
         return None
 
+    def format_tools(self, tools):
+        """Not available for Transformers."""
+        if tools:
+            raise NotImplementedError(
+                "Transformers does not support tools."
+            )
+
 
 class Transformers(Model):
     """Thin wrapper around a `transformers` model and a `transformers`
@@ -295,9 +304,10 @@ def _prepare_model_inputs(
     def generate(
         self,
         model_input: Union[str, dict, Chat],
-        output_type: Optional[OutlinesLogitsProcessor] = None,
+        output_type: Optional[OutlinesLogitsProcessor],
+        tools: Optional[List[ToolDef]],
         **inference_kwargs: Any,
-    ) -> Union[str, List[str]]:
+    ) -> Output | List[Output]:
         """Generate text using `transformers`.
 
         Parameters
@@ -310,16 +320,19 @@ def generate(
         output_type
             The logits processor the model will use to constrain the format of
             the generated text.
+        tools
+            The tools to use for the generation.
         inference_kwargs
             Additional keyword arguments to pass to the `generate` method
             of the `transformers` model.
 
         Returns
         -------
-        Union[str, List[str]]
+        Output | List[Output]
             The text generated by the model.
 
         """
+        self.type_adapter.format_tools(tools)
         prompts, inputs = self._prepare_model_inputs(model_input, False)
         logits_processor = self.type_adapter.format_output_type(output_type)
 
@@ -336,15 +349,39 @@ def generate(
         if num_samples == 1 and len(generated_ids.shape) == 2:
             generated_ids = generated_ids.squeeze(0)
 
-        return self._decode_generation(generated_ids)
+        generated_text = self._decode_generation(generated_ids)
+
+        if isinstance(generated_text, list):
+            return [Output(content=text) for text in generated_text]
+        return Output(content=generated_text)
 
     def generate_batch(
         self,
         model_input: List[Union[str, dict, Chat]],
-        output_type: Optional[OutlinesLogitsProcessor] = None,
+        output_type: Optional[OutlinesLogitsProcessor],
+        tools: Optional[List[ToolDef]],
         **inference_kwargs: Any,
-    ) -> List[Union[str, List[str]]]:
-        """"""
+    ) -> List[Output] | List[List[Output]]:
+        """Generate a batch of completions using `transformers`.
+
+        Parameters
+        ----------
+        model_input
+            The list of prompts based on which the model will generate a response.
+        output_type
+            The logits processor the model will use to constrain the format of the generated text.
+        tools
+            The tools to use for the generation.
+        **inference_kwargs
+            Additional keyword arguments to pass to the `generate` method of the `transformers` model.
+
+        Returns
+        -------
+        List[Output] | List[List[Output]]
+            The list of text generated by the model.
+
+        """
+        self.type_adapter.format_tools(tools)
         prompts, inputs = self._prepare_model_inputs(model_input, True) # type: ignore
         logits_processor = self.type_adapter.format_output_type(output_type)
 
@@ -357,7 +394,17 @@ def generate_batch(
         if num_samples > 1:
             generated_ids = generated_ids.view(len(model_input), num_samples, -1)
 
-        return self._decode_generation(generated_ids)
+        generated_text = self._decode_generation(generated_ids)
+
+        return [  # type: ignore
+            [
+                Output(content=text)
+                for text in batch
+            ]
+            if isinstance(batch, list)
+            else Output(content=batch)
+            for batch in generated_text
+        ]
 
     def generate_stream(self, model_input, output_type, **inference_kwargs):
         """Not available for `transformers` models.
@@ -369,7 +416,7 @@ def generate_stream(self, model_input, output_type, **inference_kwargs):
             "Streaming is not implemented for Transformers models."
         )
 
-    def _generate_output_seq(self, prompts, inputs, **inference_kwargs):
+    def _generate_output_seq(self, prompts, inputs, **inference_kwargs): # type: ignore
         input_ids = inputs["input_ids"]
 
         output_ids = self.model.generate(
@@ -472,7 +519,7 @@ def format_chat_input(self, model_input: Chat) -> dict:
                     "content": message["content"][0],
                 })
             else:
-                messages_without_images.append(message)
+                messages_without_images.append(message)  # type: ignore
         formatted_prompt = self.tokenizer.apply_chat_template(
             messages_without_images,
             tokenize=False
@@ -513,7 +560,7 @@ def format_list_input(self, model_input: list) -> dict:
 
     def format_output_type(
         self,
-        output_type: Optional[OutlinesLogitsProcessor] = None,
+        output_type: Optional[OutlinesLogitsProcessor],
     ) -> Optional["LogitsProcessorList"]:
         """Generate the logits processor argument to pass to the model.
 
@@ -534,6 +581,13 @@ def format_output_type(
             return LogitsProcessorList([output_type])
         return None
 
+    def format_tools(self, tools):
+        """Not available for TransformersMultiModal."""
+        if tools:
+            raise NotImplementedError(
+                "TransformersMultiModal does not support tools."
+            )
+
 
 class TransformersMultiModal(Transformers):
     """Thin wrapper around a `transformers` model and a `transformers`

tests/models/test_transformers.py

@@ -12,6 +12,7 @@
     TransformerTokenizer,
     TransformersTypeAdapter,
 )
+from outlines.outputs import Output, StreamingOutput
 from outlines.types import Regex
 
 
@@ -80,16 +81,16 @@ def model_bart():
 
 
 def test_transformers_simple(model):
-    result = model.generate("Respond with one word. Not more.", None)
-    assert isinstance(result, str)
+    result = model("Respond with one word. Not more.", None)
+    assert isinstance(result, Output)
 
 
 def test_transformers_call(model, model_bart):
     result = model("Respond with one word. Not more.")
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
 
     result = model_bart("Respond with one word. Not more.")
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
 
 
 def test_transformers_chat(model):
@@ -99,12 +100,12 @@ def test_transformers_chat(model):
             {"role": "user", "content": "What is the capital of France?"},
         ])
     )
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
 
 
 def test_transformers_inference_kwargs(model):
     result = model("Respond with one word. Not more.", max_new_tokens=100)
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
 
 
 def test_transformers_invalid_inference_kwargs(model):
@@ -114,16 +115,16 @@ def test_transformers_invalid_inference_kwargs(model):
 
 def test_transformers_regex(model):
     result = model("Give a number between 0 and 9.", Regex(r"[0-9]"))
-    assert isinstance(result, str)
-    assert re.match(r"[0-9]", result)
+    assert isinstance(result, Output)
+    assert re.match(r"[0-9]", result.content)
 
 
 def test_transformers_json(model):
     class Character(BaseModel):
         name: str
 
     result = model("Create a character with a name.", Character)
-    assert "name" in result
+    assert "name" in result.content
 
 
 def test_transformers_choice(model):
@@ -132,12 +133,12 @@ class Foo(Enum):
         dog = "dog"
 
     result = model("Cat or dog?", Foo)
-    assert result in ["cat", "dog"]
+    assert result.content in ["cat", "dog"]
 
 
 def test_transformers_multiple_samples(model):
     result = model("Respond with one word. Not more.")
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
     result = model(
         "Respond with one word. Not more.", num_return_sequences=2, do_sample=True
     )
@@ -187,8 +188,8 @@ class Foo(Enum):
     result = model("Cat or dog?", Foo, num_return_sequences=2, do_sample=True)
     assert isinstance(result, list)
     assert len(result) == 2
-    assert result[0] in ["cat", "dog"]
-    assert result[1] in ["cat", "dog"]
+    assert result[0].content in ["cat", "dog"]
+    assert result[1].content in ["cat", "dog"]
 
 
 def test_transformers_batch_constrained(model):
@@ -202,8 +203,8 @@ class Foo(Enum):
     )
     assert isinstance(result, list)
     assert len(result) == 2
-    assert result[0] in ["cat", "dog"]
-    assert result[1] in ["cat", "dog"]
+    assert result[0].content in ["cat", "dog"]
+    assert result[1].content in ["cat", "dog"]
 
     result = model.batch(
         ["Cat or dog?", "Cat or dog?"],
@@ -216,8 +217,8 @@ class Foo(Enum):
     for item in result:
         assert isinstance(item, list)
         assert len(item) == 2
-        assert item[0] in ["cat", "dog"]
-        assert item[1] in ["cat", "dog"]
+        assert item[0].content in ["cat", "dog"]
+        assert item[1].content in ["cat", "dog"]
 
 
 def test_transformers_streaming(model):

tests/models/test_transformers_multimodal.py

@@ -19,6 +19,7 @@
     TransformerTokenizer,
     TransformersMultiModalTypeAdapter,
 )
+from outlines.outputs import Output, StreamingOutput
 from outlines.types import Regex
 
 TEST_MODEL = "trl-internal-testing/tiny-LlavaForConditionalGeneration"
@@ -61,20 +62,20 @@ def test_transformers_multimodal_instantiate_simple():
 
 
 def test_transformers_multimodal_simple(model, image):
-    result = model.generate(
+    result = model(
         ["<image>Describe this image in one sentence:", Image(image)],
         None,
         max_new_tokens=2,
     )
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
 
 
 def test_transformers_multimodal_call(model, image):
     result = model(
         ["<image>Describe this image in one sentence:", Image(image)],
         max_new_tokens=2,
     )
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
 
 
 def test_transformers_multimodal_wrong_number_image(model, image):
@@ -90,7 +91,7 @@ def test_transformers_multimodal_wrong_number_image(model, image):
 
 def test_transformers_multimodal_wrong_input_type(model):
     with pytest.raises(TypeError):
-        model.generate("invalid input", None)
+        model("invalid input", None)
 
 
 def test_transformers_multimodal_chat(model, image):
@@ -107,15 +108,15 @@ def test_transformers_multimodal_chat(model, image):
         ]),
         max_new_tokens=2,
     )
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
 
 
 def test_transformers_inference_kwargs(model, image):
     result = model(
         ["<image>Describe this image in one sentence:", Image(image)],
         max_new_tokens=2,
     )
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
 
 
 def test_transformers_invalid_inference_kwargs(model, image):
@@ -138,7 +139,7 @@ def test_transformers_several_image(model, image):
         ],
         max_new_tokens=2,
     )
-    assert isinstance(result, str)
+    assert isinstance(result, Output)
 
 
 def test_transformers_multimodal_json(model, image):
@@ -150,7 +151,8 @@ class Foo(BaseModel):
         Foo,
         max_new_tokens=10,
     )
-    assert "name" in result
+    assert isinstance(result, Output)
+    assert "name" in result.content
 
 
 def test_transformers_multimodal_regex(model, image):
@@ -159,8 +161,8 @@ def test_transformers_multimodal_regex(model, image):
         Regex(r"[0-9]")
     )
 
-    assert isinstance(result, str)
-    assert re.match(r"[0-9]", result)
+    assert isinstance(result, Output)
+    assert re.match(r"[0-9]", result.content)
 
 
 def test_transformers_multimodal_choice(model, image):
@@ -173,8 +175,8 @@ class Foo(Enum):
         Foo,
     )
 
-    assert isinstance(result, str)
-    assert result in ["white", "blue"]
+    assert isinstance(result, Output)
+    assert result.content in ["white", "blue"]
 
 
 def test_transformers_multimodal_multiple_samples(model, image):
@@ -245,12 +247,12 @@ def test_transformers_multimodal_batch(model, image):
 
 def test_transformers_multimodal_deprecated_input_type(model, image):
     with pytest.warns(DeprecationWarning):
-        result = model.generate(
+        result = model(
             {
                 "text": "<image>Describe this image in one sentence:",
                 "image": image,
             },
             None,
             max_new_tokens=2,
         )
-        assert isinstance(result, str)
+        assert isinstance(result, Output)