We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
small typo in galactica model loading
https://github.com/huggingface/text-generation-inference/blob/1fdc88ee908beb8ae0afe17810a17b9b4d8848e2/server/text_generation_server/models/__init__.py#L92

should be trust_remote_code

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `server/text_generation_server/models/__init__.py`
Content:
```
1 import os
2 import torch
3 
4 from loguru import logger
5 from transformers.configuration_utils import PretrainedConfig
6 from transformers.models.auto import modeling_auto
7 from typing import Optional
8 
9 from text_generation_server.models.model import Model
10 from text_generation_server.models.causal_lm import CausalLM
11 from text_generation_server.models.flash_causal_lm import FlashCausalLM
12 from text_generation_server.models.bloom import BLOOMSharded
13 from text_generation_server.models.mpt import MPTSharded
14 from text_generation_server.models.seq2seq_lm import Seq2SeqLM
15 from text_generation_server.models.rw import RW
16 from text_generation_server.models.opt import OPTSharded
17 from text_generation_server.models.galactica import GalacticaSharded
18 from text_generation_server.models.santacoder import SantaCoder
19 from text_generation_server.models.t5 import T5Sharded
20 from text_generation_server.models.gpt_neox import GPTNeoxSharded
21 
22 # The flag below controls whether to allow TF32 on matmul. This flag defaults to False
23 # in PyTorch 1.12 and later.
24 torch.backends.cuda.matmul.allow_tf32 = True
25 
26 # The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
27 torch.backends.cudnn.allow_tf32 = True
28 
29 # Disable gradients
30 torch.set_grad_enabled(False)
31 
32 __all__ = [
33     "Model",
34     "BLOOMSharded",
35     "CausalLM",
36     "FlashCausalLM",
37     "GalacticaSharded",
38     "Seq2SeqLM",
39     "SantaCoder",
40     "OPTSharded",
41     "T5Sharded",
42     "get_model",
43 ]
44 
45 FLASH_ATT_ERROR_MESSAGE = "{} requires Flash Attention enabled models."
46 
47 FLASH_ATTENTION = True
48 try:
49     from text_generation_server.models.flash_rw import FlashRWSharded
50     from text_generation_server.models.flash_neox import FlashNeoXSharded
51     from text_generation_server.models.flash_llama import (
52         FlashLlama,
53     )
54     from text_generation_server.models.flash_santacoder import (
55         FlashSantacoderSharded,
56     )
57 
58 except ImportError as e:
59     logger.warning(f"Could not import Flash Attention enabled models: {e}")
60     FLASH_ATTENTION = False
61 
62 if FLASH_ATTENTION:
63     __all__.append(FlashNeoXSharded)
64     __all__.append(FlashRWSharded)
65     __all__.append(FlashSantacoderSharded)
66     __all__.append(FlashLlama)
67 
68 
69 def get_model(
70     model_id: str,
71     revision: Optional[str],
72     sharded: bool,
73     quantize: Optional[str],
74     dtype: Optional[str],
75     trust_remote_code: bool,
76 ) -> Model:
77     if dtype is None:
78         dtype = torch.float16
79     elif dtype == "float16":
80         dtype = torch.float16
81     elif dtype == "bfloat16":
82         dtype = torch.bfloat16
83     else:
84         raise RuntimeError(f"Unknown dtype {dtype}")
85 
86     if "facebook/galactica" in model_id:
87         return GalacticaSharded(
88             model_id,
89             revision,
90             quantize=quantize,
91             dtype=dtype,
92             dtypetrust_remote_code=trust_remote_code,
93         )
94 
95     if model_id.startswith("bigcode/"):
96         if FLASH_ATTENTION:
97             return FlashSantacoderSharded(
98                 model_id,
99                 revision,
100                 quantize=quantize,
101                 dtype=dtype,
102                 trust_remote_code=trust_remote_code,
103             )
104         elif sharded:
105             raise NotImplementedError(
106                 FLASH_ATT_ERROR_MESSAGE.format("Sharded Santacoder")
107             )
108         else:
109             return SantaCoder(
110                 model_id,
111                 revision,
112                 quantize=quantize,
113                 dtype=dtype,
114                 trust_remote_code=trust_remote_code,
115             )
116 
117     config_dict, _ = PretrainedConfig.get_config_dict(
118         model_id, revision=revision, trust_remote_code=trust_remote_code
119     )
120     model_type = config_dict["model_type"]
121 
122     if model_type == "gpt_bigcode":
123         if FLASH_ATTENTION:
124             return FlashSantacoderSharded(
125                 model_id,
126                 revision,
127                 quantize=quantize,
128                 dtype=dtype,
129                 trust_remote_code=trust_remote_code,
130             )
131         elif sharded:
132             raise NotImplementedError(
133                 FLASH_ATT_ERROR_MESSAGE.format("Sharded Santacoder")
134             )
135         else:
136             return SantaCoder(
137                 model_id,
138                 revision,
139                 quantize=quantize,
140                 dtype=dtype,
141                 trust_remote_code=trust_remote_code,
142             )
143 
144     if model_type == "bloom":
145         return BLOOMSharded(
146             model_id,
147             revision,
148             quantize=quantize,
149             dtype=dtype,
150             trust_remote_code=trust_remote_code,
151         )
152     elif model_type == "mpt":
153         return MPTSharded(
154             model_id, revision, quantize=quantize, trust_remote_code=trust_remote_code
155         )
156 
157     elif model_type == "gpt_neox":
158         if FLASH_ATTENTION:
159             return FlashNeoXSharded(
160                 model_id,
161                 revision,
162                 quantize=quantize,
163                 dtype=dtype,
164                 trust_remote_code=trust_remote_code,
165             )
166         elif sharded:
167             return GPTNeoxSharded(
168                 model_id,
169                 revision,
170                 quantize=quantize,
171                 dtype=dtype,
172                 trust_remote_code=trust_remote_code,
173             )
174         else:
175             return CausalLM(
176                 model_id,
177                 revision,
178                 quantize=quantize,
179                 dtype=dtype,
180                 trust_remote_code=trust_remote_code,
181             )
182 
183     elif model_type == "llama":
184         if FLASH_ATTENTION:
185             return FlashLlama(
186                 model_id,
187                 revision,
188                 quantize=quantize,
189                 dtype=dtype,
190                 trust_remote_code=trust_remote_code,
191             )
192         elif sharded:
193             raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Llama"))
194         else:
195             return CausalLM(
196                 model_id,
197                 revision,
198                 quantize=quantize,
199                 dtype=dtype,
200                 trust_remote_code=trust_remote_code,
201             )
202 
203     if model_type in ["RefinedWeb", "RefinedWebModel", "falcon"]:
204         if sharded:
205             if FLASH_ATTENTION:
206                 if config_dict.get("alibi", False):
207                     raise NotImplementedError("sharded is not supported for this model")
208                 return FlashRWSharded(
209                     model_id,
210                     revision,
211                     quantize=quantize,
212                     dtype=dtype,
213                     trust_remote_code=trust_remote_code,
214                 )
215             raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format(f"Sharded Falcon"))
216         else:
217             if FLASH_ATTENTION and not config_dict.get("alibi", False):
218                 return FlashRWSharded(
219                     model_id,
220                     revision,
221                     quantize=quantize,
222                     dtype=dtype,
223                     trust_remote_code=trust_remote_code,
224                 )
225             else:
226                 return RW(
227                     model_id,
228                     revision,
229                     quantize=quantize,
230                     dtype=dtype,
231                     trust_remote_code=trust_remote_code,
232                 )
233 
234     elif model_type == "opt":
235         return OPTSharded(
236             model_id,
237             revision,
238             quantize=quantize,
239             dtype=dtype,
240             trust_remote_code=trust_remote_code,
241         )
242 
243     elif model_type == "t5":
244         return T5Sharded(
245             model_id,
246             revision,
247             quantize=quantize,
248             dtype=dtype,
249             trust_remote_code=trust_remote_code,
250         )
251 
252     if sharded:
253         raise ValueError("sharded is not supported for AutoModel")
254     if quantize == "gptq":
255         raise ValueError(
256             "gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
257         )
258     elif (quantize == "bitsandbytes-fp4") or (quantize == "bitsandbytes-nf4"):
259         raise ValueError(
260             "4bit quantization is not supported for AutoModel"
261         )
262     if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
263         return CausalLM(
264             model_id,
265             revision,
266             quantize=quantize,
267             dtype=dtype,
268             trust_remote_code=trust_remote_code,
269         )
270     if model_type in modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES:
271         return Seq2SeqLM(
272             model_id,
273             revision,
274             quantize=quantize,
275             dtype=dtype,
276             trust_remote_code=trust_remote_code,
277         )
278 
279     auto_map = config_dict.get("auto_map", None)
280     if trust_remote_code and auto_map is not None:
281         if "AutoModelForCausalLM" in auto_map.keys():
282             return CausalLM(
283                 model_id,
284                 revision,
285                 quantize=quantize,
286                 dtype=dtype,
287                 trust_remote_code=trust_remote_code,
288             )
289         if "AutoModelForSeq2SeqLM" in auto_map.keys():
290             return Seq2SeqLM(
291                 model_id,
292                 revision,
293                 quantize=quantize,
294                 dtype=dtype,
295                 trust_remote_code=trust_remote_code,
296             )
297 
298     raise ValueError(f"Unsupported model type {model_type}")
299
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/server/text_generation_server/models/__init__.py b/server/text_generation_server/models/__init__.py
--- a/server/text_generation_server/models/__init__.py
+++ b/server/text_generation_server/models/__init__.py
@@ -89,7 +89,7 @@
             revision,
             quantize=quantize,
             dtype=dtype,
-            dtypetrust_remote_code=trust_remote_code,
+            trust_remote_code=trust_remote_code,
         )
 
     if model_id.startswith("bigcode/"):

{"golden_diff": "diff --git a/server/text_generation_server/models/__init__.py b/server/text_generation_server/models/__init__.py\n--- a/server/text_generation_server/models/__init__.py\n+++ b/server/text_generation_server/models/__init__.py\n@@ -89,7 +89,7 @@\n             revision,\n             quantize=quantize,\n             dtype=dtype,\n-            dtypetrust_remote_code=trust_remote_code,\n+            trust_remote_code=trust_remote_code,\n         )\n \n     if model_id.startswith(\"bigcode/\"):\n", "issue": "small typo in galactica model loading\nhttps://github.com/huggingface/text-generation-inference/blob/1fdc88ee908beb8ae0afe17810a17b9b4d8848e2/server/text_generation_server/models/__init__.py#L92\r\n\r\nshould be trust_remote_code\n", "before_files": [{"content": "import os\nimport torch\n\nfrom loguru import logger\nfrom transformers.configuration_utils import PretrainedConfig\nfrom transformers.models.auto import modeling_auto\nfrom typing import Optional\n\nfrom text_generation_server.models.model import Model\nfrom text_generation_server.models.causal_lm import CausalLM\nfrom text_generation_server.models.flash_causal_lm import FlashCausalLM\nfrom text_generation_server.models.bloom import BLOOMSharded\nfrom text_generation_server.models.mpt import MPTSharded\nfrom text_generation_server.models.seq2seq_lm import Seq2SeqLM\nfrom text_generation_server.models.rw import RW\nfrom text_generation_server.models.opt import OPTSharded\nfrom text_generation_server.models.galactica import GalacticaSharded\nfrom text_generation_server.models.santacoder import SantaCoder\nfrom text_generation_server.models.t5 import T5Sharded\nfrom text_generation_server.models.gpt_neox import GPTNeoxSharded\n\n# The flag below controls whether to allow TF32 on matmul. This flag defaults to False\n# in PyTorch 1.12 and later.\ntorch.backends.cuda.matmul.allow_tf32 = True\n\n# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.\ntorch.backends.cudnn.allow_tf32 = True\n\n# Disable gradients\ntorch.set_grad_enabled(False)\n\n__all__ = [\n    \"Model\",\n    \"BLOOMSharded\",\n    \"CausalLM\",\n    \"FlashCausalLM\",\n    \"GalacticaSharded\",\n    \"Seq2SeqLM\",\n    \"SantaCoder\",\n    \"OPTSharded\",\n    \"T5Sharded\",\n    \"get_model\",\n]\n\nFLASH_ATT_ERROR_MESSAGE = \"{} requires Flash Attention enabled models.\"\n\nFLASH_ATTENTION = True\ntry:\n    from text_generation_server.models.flash_rw import FlashRWSharded\n    from text_generation_server.models.flash_neox import FlashNeoXSharded\n    from text_generation_server.models.flash_llama import (\n        FlashLlama,\n    )\n    from text_generation_server.models.flash_santacoder import (\n        FlashSantacoderSharded,\n    )\n\nexcept ImportError as e:\n    logger.warning(f\"Could not import Flash Attention enabled models: {e}\")\n    FLASH_ATTENTION = False\n\nif FLASH_ATTENTION:\n    __all__.append(FlashNeoXSharded)\n    __all__.append(FlashRWSharded)\n    __all__.append(FlashSantacoderSharded)\n    __all__.append(FlashLlama)\n\n\ndef get_model(\n    model_id: str,\n    revision: Optional[str],\n    sharded: bool,\n    quantize: Optional[str],\n    dtype: Optional[str],\n    trust_remote_code: bool,\n) -> Model:\n    if dtype is None:\n        dtype = torch.float16\n    elif dtype == \"float16\":\n        dtype = torch.float16\n    elif dtype == \"bfloat16\":\n        dtype = torch.bfloat16\n    else:\n        raise RuntimeError(f\"Unknown dtype {dtype}\")\n\n    if \"facebook/galactica\" in model_id:\n        return GalacticaSharded(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            dtypetrust_remote_code=trust_remote_code,\n        )\n\n    if model_id.startswith(\"bigcode/\"):\n        if FLASH_ATTENTION:\n            return FlashSantacoderSharded(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        elif sharded:\n            raise NotImplementedError(\n                FLASH_ATT_ERROR_MESSAGE.format(\"Sharded Santacoder\")\n            )\n        else:\n            return SantaCoder(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    config_dict, _ = PretrainedConfig.get_config_dict(\n        model_id, revision=revision, trust_remote_code=trust_remote_code\n    )\n    model_type = config_dict[\"model_type\"]\n\n    if model_type == \"gpt_bigcode\":\n        if FLASH_ATTENTION:\n            return FlashSantacoderSharded(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        elif sharded:\n            raise NotImplementedError(\n                FLASH_ATT_ERROR_MESSAGE.format(\"Sharded Santacoder\")\n            )\n        else:\n            return SantaCoder(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    if model_type == \"bloom\":\n        return BLOOMSharded(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n    elif model_type == \"mpt\":\n        return MPTSharded(\n            model_id, revision, quantize=quantize, trust_remote_code=trust_remote_code\n        )\n\n    elif model_type == \"gpt_neox\":\n        if FLASH_ATTENTION:\n            return FlashNeoXSharded(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        elif sharded:\n            return GPTNeoxSharded(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        else:\n            return CausalLM(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    elif model_type == \"llama\":\n        if FLASH_ATTENTION:\n            return FlashLlama(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        elif sharded:\n            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format(\"Sharded Llama\"))\n        else:\n            return CausalLM(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    if model_type in [\"RefinedWeb\", \"RefinedWebModel\", \"falcon\"]:\n        if sharded:\n            if FLASH_ATTENTION:\n                if config_dict.get(\"alibi\", False):\n                    raise NotImplementedError(\"sharded is not supported for this model\")\n                return FlashRWSharded(\n                    model_id,\n                    revision,\n                    quantize=quantize,\n                    dtype=dtype,\n                    trust_remote_code=trust_remote_code,\n                )\n            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format(f\"Sharded Falcon\"))\n        else:\n            if FLASH_ATTENTION and not config_dict.get(\"alibi\", False):\n                return FlashRWSharded(\n                    model_id,\n                    revision,\n                    quantize=quantize,\n                    dtype=dtype,\n                    trust_remote_code=trust_remote_code,\n                )\n            else:\n                return RW(\n                    model_id,\n                    revision,\n                    quantize=quantize,\n                    dtype=dtype,\n                    trust_remote_code=trust_remote_code,\n                )\n\n    elif model_type == \"opt\":\n        return OPTSharded(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n\n    elif model_type == \"t5\":\n        return T5Sharded(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n\n    if sharded:\n        raise ValueError(\"sharded is not supported for AutoModel\")\n    if quantize == \"gptq\":\n        raise ValueError(\n            \"gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`\"\n        )\n    elif (quantize == \"bitsandbytes-fp4\") or (quantize == \"bitsandbytes-nf4\"):\n        raise ValueError(\n            \"4bit quantization is not supported for AutoModel\"\n        )\n    if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:\n        return CausalLM(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n    if model_type in modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES:\n        return Seq2SeqLM(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n\n    auto_map = config_dict.get(\"auto_map\", None)\n    if trust_remote_code and auto_map is not None:\n        if \"AutoModelForCausalLM\" in auto_map.keys():\n            return CausalLM(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        if \"AutoModelForSeq2SeqLM\" in auto_map.keys():\n            return Seq2SeqLM(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    raise ValueError(f\"Unsupported model type {model_type}\")\n", "path": "server/text_generation_server/models/__init__.py"}], "after_files": [{"content": "import os\nimport torch\n\nfrom loguru import logger\nfrom transformers.configuration_utils import PretrainedConfig\nfrom transformers.models.auto import modeling_auto\nfrom typing import Optional\n\nfrom text_generation_server.models.model import Model\nfrom text_generation_server.models.causal_lm import CausalLM\nfrom text_generation_server.models.flash_causal_lm import FlashCausalLM\nfrom text_generation_server.models.bloom import BLOOMSharded\nfrom text_generation_server.models.mpt import MPTSharded\nfrom text_generation_server.models.seq2seq_lm import Seq2SeqLM\nfrom text_generation_server.models.rw import RW\nfrom text_generation_server.models.opt import OPTSharded\nfrom text_generation_server.models.galactica import GalacticaSharded\nfrom text_generation_server.models.santacoder import SantaCoder\nfrom text_generation_server.models.t5 import T5Sharded\nfrom text_generation_server.models.gpt_neox import GPTNeoxSharded\n\n# The flag below controls whether to allow TF32 on matmul. This flag defaults to False\n# in PyTorch 1.12 and later.\ntorch.backends.cuda.matmul.allow_tf32 = True\n\n# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.\ntorch.backends.cudnn.allow_tf32 = True\n\n# Disable gradients\ntorch.set_grad_enabled(False)\n\n__all__ = [\n    \"Model\",\n    \"BLOOMSharded\",\n    \"CausalLM\",\n    \"FlashCausalLM\",\n    \"GalacticaSharded\",\n    \"Seq2SeqLM\",\n    \"SantaCoder\",\n    \"OPTSharded\",\n    \"T5Sharded\",\n    \"get_model\",\n]\n\nFLASH_ATT_ERROR_MESSAGE = \"{} requires Flash Attention enabled models.\"\n\nFLASH_ATTENTION = True\ntry:\n    from text_generation_server.models.flash_rw import FlashRWSharded\n    from text_generation_server.models.flash_neox import FlashNeoXSharded\n    from text_generation_server.models.flash_llama import (\n        FlashLlama,\n    )\n    from text_generation_server.models.flash_santacoder import (\n        FlashSantacoderSharded,\n    )\n\nexcept ImportError as e:\n    logger.warning(f\"Could not import Flash Attention enabled models: {e}\")\n    FLASH_ATTENTION = False\n\nif FLASH_ATTENTION:\n    __all__.append(FlashNeoXSharded)\n    __all__.append(FlashRWSharded)\n    __all__.append(FlashSantacoderSharded)\n    __all__.append(FlashLlama)\n\n\ndef get_model(\n    model_id: str,\n    revision: Optional[str],\n    sharded: bool,\n    quantize: Optional[str],\n    dtype: Optional[str],\n    trust_remote_code: bool,\n) -> Model:\n    if dtype is None:\n        dtype = torch.float16\n    elif dtype == \"float16\":\n        dtype = torch.float16\n    elif dtype == \"bfloat16\":\n        dtype = torch.bfloat16\n    else:\n        raise RuntimeError(f\"Unknown dtype {dtype}\")\n\n    if \"facebook/galactica\" in model_id:\n        return GalacticaSharded(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n\n    if model_id.startswith(\"bigcode/\"):\n        if FLASH_ATTENTION:\n            return FlashSantacoderSharded(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        elif sharded:\n            raise NotImplementedError(\n                FLASH_ATT_ERROR_MESSAGE.format(\"Sharded Santacoder\")\n            )\n        else:\n            return SantaCoder(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    config_dict, _ = PretrainedConfig.get_config_dict(\n        model_id, revision=revision, trust_remote_code=trust_remote_code\n    )\n    model_type = config_dict[\"model_type\"]\n\n    if model_type == \"gpt_bigcode\":\n        if FLASH_ATTENTION:\n            return FlashSantacoderSharded(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        elif sharded:\n            raise NotImplementedError(\n                FLASH_ATT_ERROR_MESSAGE.format(\"Sharded Santacoder\")\n            )\n        else:\n            return SantaCoder(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    if model_type == \"bloom\":\n        return BLOOMSharded(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n    elif model_type == \"mpt\":\n        return MPTSharded(\n            model_id, revision, quantize=quantize, trust_remote_code=trust_remote_code\n        )\n\n    elif model_type == \"gpt_neox\":\n        if FLASH_ATTENTION:\n            return FlashNeoXSharded(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        elif sharded:\n            return GPTNeoxSharded(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        else:\n            return CausalLM(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    elif model_type == \"llama\":\n        if FLASH_ATTENTION:\n            return FlashLlama(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        elif sharded:\n            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format(\"Sharded Llama\"))\n        else:\n            return CausalLM(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    if model_type in [\"RefinedWeb\", \"RefinedWebModel\", \"falcon\"]:\n        if sharded:\n            if FLASH_ATTENTION:\n                if config_dict.get(\"alibi\", False):\n                    raise NotImplementedError(\"sharded is not supported for this model\")\n                return FlashRWSharded(\n                    model_id,\n                    revision,\n                    quantize=quantize,\n                    dtype=dtype,\n                    trust_remote_code=trust_remote_code,\n                )\n            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format(f\"Sharded Falcon\"))\n        else:\n            if FLASH_ATTENTION and not config_dict.get(\"alibi\", False):\n                return FlashRWSharded(\n                    model_id,\n                    revision,\n                    quantize=quantize,\n                    dtype=dtype,\n                    trust_remote_code=trust_remote_code,\n                )\n            else:\n                return RW(\n                    model_id,\n                    revision,\n                    quantize=quantize,\n                    dtype=dtype,\n                    trust_remote_code=trust_remote_code,\n                )\n\n    elif model_type == \"opt\":\n        return OPTSharded(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n\n    elif model_type == \"t5\":\n        return T5Sharded(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n\n    if sharded:\n        raise ValueError(\"sharded is not supported for AutoModel\")\n    if quantize == \"gptq\":\n        raise ValueError(\n            \"gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`\"\n        )\n    elif (quantize == \"bitsandbytes-fp4\") or (quantize == \"bitsandbytes-nf4\"):\n        raise ValueError(\n            \"4bit quantization is not supported for AutoModel\"\n        )\n    if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:\n        return CausalLM(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n    if model_type in modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES:\n        return Seq2SeqLM(\n            model_id,\n            revision,\n            quantize=quantize,\n            dtype=dtype,\n            trust_remote_code=trust_remote_code,\n        )\n\n    auto_map = config_dict.get(\"auto_map\", None)\n    if trust_remote_code and auto_map is not None:\n        if \"AutoModelForCausalLM\" in auto_map.keys():\n            return CausalLM(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n        if \"AutoModelForSeq2SeqLM\" in auto_map.keys():\n            return Seq2SeqLM(\n                model_id,\n                revision,\n                quantize=quantize,\n                dtype=dtype,\n                trust_remote_code=trust_remote_code,\n            )\n\n    raise ValueError(f\"Unsupported model type {model_type}\")\n", "path": "server/text_generation_server/models/__init__.py"}]}