Add integration for AutoGPTQ and AutoAWQ

outlines/models/__init__.py

@@ -5,5 +5,7 @@
 codebase.
 
 """
+from .awq import awq
+from .gptq import gptq
 from .openai import OpenAI, openai
-from .transformers import Transformers, transformers
+from .transformers import Transformer, transformers

outlines/models/awq.py

@@ -0,0 +1,45 @@
+from typing import TYPE_CHECKING, Optional
+
+from .transformers import Transformer, TransformerTokenizer
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer
+
+
+class AWQModel(Transformer):
+    """Represents a `transformers` model."""
+
+    def __init__(
+        self,
+        model: "PreTrainedModel",
+        tokenizer: "PreTrainedTokenizer",
+    ):
+        self.device = model.model.device
+        self.model = model
+        self.tokenizer = tokenizer
+
+
+def awq(
+    model_name: str,
+    fuse_layers: bool = True,
+    device: Optional[str] = None,
+    model_kwargs: dict = {},
+    tokenizer_kwargs: dict = {},
+):
+    try:
+        from awq import AutoAWQForCausalLM
+    except ImportError:
+        raise ImportError(
+            "The `autoawq` and `transformers` library needs to be installed in order to use `AutoAWQ` models."
+        )
+
+    model_kwargs["fuse_layers"] = fuse_layers
+    model_kwargs["safetensors"] = True
+
+    if device is not None:
+        model_kwargs["device_map"] = device
+
+    model = AutoAWQForCausalLM.from_quantized(model_name, **model_kwargs)
+    tokenizer = TransformerTokenizer(model_name, trust_remote_code=True)
+
+    return AWQModel(model, tokenizer)

outlines/models/gptq.py

@@ -0,0 +1,25 @@
+from typing import Optional
+
+from .transformers import Transformer, TransformerTokenizer
+
+
+def gptq(
+    model_name: str,
+    device: Optional[str] = None,
+    model_kwargs: dict = {},
+    tokenizer_kwargs: dict = {},
+):
+    try:
+        from auto_gptq import AutoGPTQForCausalLM
+    except ImportError:
+        raise ImportError(
+            "The `auto_gptq` library needs to be installed in order to use `AutoGPTQ` models."
+        )
+
+    if device is not None:
+        model_kwargs["device_map"] = device
+
+    model = AutoGPTQForCausalLM.from_quantized(model_name, **model_kwargs)
+    tokenizer = TransformerTokenizer(model_name, **tokenizer_kwargs)
+
+    return Transformer(model, tokenizer)

outlines/models/transformers.py

@@ -55,7 +55,7 @@ class CodeLlamaTokenizerFast:  # type: ignore
     )
 
 
-class Transformers:
+class Transformer:
     """Represents a `transformers` model."""
 
     def __init__(
@@ -116,7 +116,7 @@ def __call__(
         return self.forward(input_ids, attention_mask, past_key_values)[0]
 
 
-class TransformersTokenizer(Tokenizer):
+class TransformerTokenizer(Tokenizer):
     """Represents a tokenizer for models in the `transformers` library."""
 
     def __init__(self, model_name: str, **kwargs):
@@ -215,6 +215,6 @@ def transformers(
         model_kwargs["device_map"] = device
 
     model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
-    tokenizer = TransformersTokenizer(model_name, **tokenizer_kwargs)
+    tokenizer = TransformerTokenizer(model_name, **tokenizer_kwargs)
 
-    return Transformers(model, tokenizer)
+    return Transformer(model, tokenizer)

outlines/text/generate/sequence.py

@@ -6,7 +6,7 @@
 from outlines.models import OpenAI
 
 if TYPE_CHECKING:
-    from outlines.models.transformers import KVCacheType, Transformers
+    from outlines.models.transformers import KVCacheType, Transformer
     from outlines.text.generate.sample import Sampler
 
 
@@ -15,7 +15,7 @@ class Sequence:
 
     def __init__(
         self,
-        model: "Transformers",
+        model: "Transformer",
         max_tokens: Optional[int] = None,
         sampler: Optional["Sampler"] = None,
     ):
@@ -41,6 +41,7 @@ def __init__(
         self.model = model
         self.device = model.device
         self.max_tokens = max_tokens
+
         self.pad_token_id = torch.tensor(
             model.tokenizer.pad_token_id, device=model.device
         )

pyproject.toml

@@ -87,6 +87,8 @@ exclude=["examples"]
 
 [[tool.mypy.overrides]]
 module = [
+    "awq.*",
+    "auto_gptq.*",
     "jinja2",
     "joblib.*",
     "jsonschema.*",

tests/models/test_transformers.py

@@ -2,13 +2,13 @@
 import torch
 from transformers.models.gpt2 import GPT2TokenizerFast
 
-from outlines.models.transformers import TransformersTokenizer, transformers
+from outlines.models.transformers import TransformerTokenizer, transformers
 
 TEST_MODEL = "hf-internal-testing/tiny-random-GPTJForCausalLM"
 
 
 def test_tokenizer():
-    tokenizer = TransformersTokenizer(TEST_MODEL)
+    tokenizer = TransformerTokenizer(TEST_MODEL)
     assert tokenizer.eos_token_id == 0
     assert tokenizer.pad_token_id == 0
     assert isinstance(tokenizer.tokenizer, GPT2TokenizerFast)
@@ -37,15 +37,15 @@ def test_tokenizer():
     isinstance(text[0], str)
     isinstance(text[1], str)
 
-    tokenizer = TransformersTokenizer(
+    tokenizer = TransformerTokenizer(
         TEST_MODEL, additional_special_tokens=["<t1>", "<t2>"]
     )
     assert "<t1>" in tokenizer.special_tokens
     assert "<t2>" in tokenizer.special_tokens
 
 
 def test_llama_tokenizer():
-    tokenizer = TransformersTokenizer("hf-internal-testing/llama-tokenizer")
+    tokenizer = TransformerTokenizer("hf-internal-testing/llama-tokenizer")
 
     # Broken
     assert tokenizer.tokenizer.convert_tokens_to_string(["▁baz"]) == "baz"
@@ -63,15 +63,15 @@ def test_model():
         transformers(TEST_MODEL, device="non_existent")
 
     model = transformers(TEST_MODEL, device="cpu")
-    assert isinstance(model.tokenizer, TransformersTokenizer)
+    assert isinstance(model.tokenizer, TransformerTokenizer)
     assert model.device.type == "cpu"
 
     model = transformers(TEST_MODEL, model_kwargs={"device_map": "cpu"})
-    assert isinstance(model.tokenizer, TransformersTokenizer)
+    assert isinstance(model.tokenizer, TransformerTokenizer)
     assert model.device.type == "cpu"
 
     model = transformers(TEST_MODEL, device="cpu", model_kwargs={"device_map": "cuda"})
-    assert isinstance(model.tokenizer, TransformersTokenizer)
+    assert isinstance(model.tokenizer, TransformerTokenizer)
     assert model.device.type == "cpu"
 
     input_ids = torch.tensor([[0, 1, 2]])
@@ -92,7 +92,7 @@ def test_model():
 
 
 def test_tokenizer_eq_hash():
-    tokenizer = TransformersTokenizer("gpt2")
-    tokenizer2 = TransformersTokenizer("gpt2")
+    tokenizer = TransformerTokenizer("gpt2")
+    tokenizer2 = TransformerTokenizer("gpt2")
     assert tokenizer == tokenizer2
     assert hash(tokenizer) == hash(tokenizer2)

tests/text/generate/test_integration_transfomers.py

@@ -9,7 +9,7 @@
 
 import outlines.models as models
 import outlines.text.generate as generate
-from outlines.models.transformers import TransformersTokenizer
+from outlines.models.transformers import TransformerTokenizer
 from outlines.text.fsm import reduced_vocabulary
 
 
@@ -326,8 +326,8 @@ def test_transformers_logits_vocab_size():
 
 
 def test_transformers_reduced_vocabulary_caching():
-    tokenizer = TransformersTokenizer("gpt2")
-    tokenizer2 = TransformersTokenizer("gpt2")
+    tokenizer = TransformerTokenizer("gpt2")
+    tokenizer2 = TransformerTokenizer("gpt2")
 
     # TODO: We might actually want only one copy of a given tokenizer.
     assert tokenizer is not tokenizer2

tests/text/test_fsm.py

@@ -2,7 +2,7 @@
 import numba
 import pytest
 
-from outlines.models.transformers import TransformersTokenizer
+from outlines.models.transformers import TransformerTokenizer
 from outlines.text.fsm import (
     _walk_fsm,
     create_fsm_index,
@@ -380,7 +380,7 @@ def test_create_fsm_index_tokenizer():
     num_fsm_states = len(regex_fsm.states)
     assert num_fsm_states == 220
 
-    tokenizer = TransformersTokenizer("gpt2")
+    tokenizer = TransformerTokenizer("gpt2")
 
     states_to_token_subsets, empty_token_ids = create_fsm_index_tokenizer(
         regex_fsm, tokenizer
@@ -403,7 +403,7 @@ def test_regex_index_performance():
     num_fsm_states = len(regex_fsm.states)
     assert num_fsm_states == 220
 
-    tokenizer = TransformersTokenizer("gpt2")
+    tokenizer = TransformerTokenizer("gpt2")
 
     # Pre-compile Numba functions
     res, _ = create_fsm_index_tokenizer(regex_fsm, tokenizer)