diff --git a/outlines/models/exllamav2.py b/outlines/models/exllamav2.py index d2aa84b0d..af7b6a058 100644 --- a/outlines/models/exllamav2.py +++ b/outlines/models/exllamav2.py @@ -42,6 +42,8 @@ def decode(self, token_ids: "torch.LongTensor") -> List[str]: class ExLlamaV2Model: """Represents a `exl2` model.""" + model_name = "ExLlamaV2" + def __init__( self, generator: "ExLlamaV2DynamicGenerator", diff --git a/outlines/models/llamacpp.py b/outlines/models/llamacpp.py index 904b193c4..3a847d385 100644 --- a/outlines/models/llamacpp.py +++ b/outlines/models/llamacpp.py @@ -138,6 +138,8 @@ class LlamaCpp: """ + model_name = "LlamaCpp" + def __init__(self, model: "Llama"): self.model = model diff --git a/outlines/models/mlxlm.py b/outlines/models/mlxlm.py index d8b7e032c..369d959f0 100644 --- a/outlines/models/mlxlm.py +++ b/outlines/models/mlxlm.py @@ -17,6 +17,8 @@ class MLXLM: Represents an `mlx_lm` model """ + model_name = "MLXLM" + def __init__( self, model: "nn.Module", diff --git a/outlines/models/transformers.py b/outlines/models/transformers.py index 444492500..c080b06f3 100644 --- a/outlines/models/transformers.py +++ b/outlines/models/transformers.py @@ -129,6 +129,8 @@ def __setstate__(self, state): class Transformers: """Represents a `transformers` model.""" + model_name = "Transformers" + def __init__( self, model: "PreTrainedModel", diff --git a/outlines/models/vllm.py b/outlines/models/vllm.py index 778c27c6f..a1f5332bc 100644 --- a/outlines/models/vllm.py +++ b/outlines/models/vllm.py @@ -19,6 +19,8 @@ class VLLM: """ + model_name = "VLLM" + def __init__(self, model: "LLM"): self.model = model self.lora_request = None