We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
missing 2 required positional arguments: 'top_n_tokens' and 'top_n_tokens_tensor' for the galactica model
### System Info

v1.0.3

### Information

- [ ] Docker
- [ ] The CLI directly

### Tasks

- [ ] An officially supported command
- [ ] My own modifications

### Reproduction

having a local galactica model in a folder called `facebook-galactica-30b-gptq` won't be detected since it will fail this check https://github.com/huggingface/text-generation-inference/blob/main/server/text_generation_server/models/__init__.py#L88

I suggest making it check `if "galactica" in  in model_id` instead.

### Expected behavior

expected to detect a galactica model

--- 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/galactica.py`
Content:
```
1 import re
2 import torch
3 import torch.distributed
4 
5 from typing import List, Optional, Type
6 
7 from transformers import (
8     AutoTokenizer,
9     AutoConfig,
10     PreTrainedTokenizerBase,
11 )
12 from text_generation_server.models import CausalLM
13 from text_generation_server.models.causal_lm import CausalLMBatch
14 from text_generation_server.pb import generate_pb2
15 from text_generation_server.models.custom_modeling.opt_modeling import OPTForCausalLM
16 from text_generation_server.utils import (
17     NextTokenChooser,
18     StoppingCriteria,
19     initialize_torch_distributed,
20     weight_files,
21     Weights,
22 )
23 
24 # CREDIT: Papers with code => https://github.com/paperswithcode/galai/blob/main/galai/utils.py
25 
26 # we split individual characters inside special tokens like [START_DNA]
27 CUSTOM_SEQ_RE = re.compile(r"(\[START_(DNA|SMILES|I_SMILES|AMINO)])(.*?)(\[END_\2])")
28 
29 # token added to implement a custom sequence tokenization. This token is added at
30 # corpus cleaning step and removed in pretokenization. The digits are added to increase the chance
31 # that they do not occur in the corpus. The digits are escaped so that the token does not appear
32 # literally in the source code in case we ever include it in the training data.
33 SPLIT_MARKER = f"SPL{1}T-TH{1}S-Pl3A5E"
34 
35 
36 def _insert_split_marker(m: re.Match):
37     """
38     Applies split marker based on a regex match of special tokens such as
39     [START_DNA].
40     Parameters
41     ----------
42     n : str
43         Input text to split
44     Returns
45     ----------
46     str - the text with the split token added
47     """
48     start_token, _, sequence, end_token = m.groups()
49     sequence = re.sub(r"(.)", rf"{SPLIT_MARKER}\1", sequence, flags=re.DOTALL)
50     return f"{start_token}{sequence}{SPLIT_MARKER}{end_token}"
51 
52 
53 def escape_custom_split_sequence(text):
54     """
55     Applies custom splitting to the text for GALILEO's tokenization
56     Parameters
57     ----------
58     text : str
59         Input text to split
60     Returns
61     ----------
62     str - the text with the split token added
63     """
64     return CUSTOM_SEQ_RE.sub(_insert_split_marker, text)
65 
66 
67 # END CREDIT
68 
69 
70 class GalacticaCausalLMBatch(CausalLMBatch):
71     @classmethod
72     def from_pb(
73         cls,
74         pb: generate_pb2.Batch,
75         tokenizer: PreTrainedTokenizerBase,
76         dtype: torch.dtype,
77         device: torch.device,
78     ) -> "GalacticaCausalLMBatch":
79         inputs = []
80         next_token_choosers = []
81         stopping_criterias = []
82         prefix_offsets = []
83         read_offsets = []
84         requests_idx_mapping = {}
85 
86         # Parse batch
87         max_truncation = 0
88         padding_right_offset = 0
89         max_decode_tokens = 0
90         for i, r in enumerate(pb.requests):
91             requests_idx_mapping[r.id] = i
92             # Add escape_custom_split_sequence to the CausalLMBatch logic
93             inputs.append(escape_custom_split_sequence(r.inputs))
94             next_token_choosers.append(NextTokenChooser.from_pb(r.parameters, device))
95             stopping_criteria = StoppingCriteria.from_pb(
96                 r.stopping_parameters, tokenizer
97             )
98             stopping_criterias.append(stopping_criteria)
99             max_truncation = max(max_truncation, r.truncate)
100             max_decode_tokens += stopping_criteria.max_new_tokens
101             padding_right_offset = max(
102                 padding_right_offset, stopping_criteria.max_new_tokens
103             )
104 
105         tokenized_inputs = tokenizer(
106             inputs,
107             return_tensors="pt",
108             padding=True,
109             return_token_type_ids=False,
110             truncation=True,
111             max_length=max_truncation,
112         ).to(device)
113         for _ in pb.requests:
114             input_len = tokenized_inputs["input_ids"].shape[1]
115             prefix_offsets.append(0)
116             read_offsets.append(input_len)
117 
118         input_lengths = tokenized_inputs["attention_mask"].sum(1)
119         max_input_length = input_lengths.max()
120 
121         input_ids = tokenized_inputs["input_ids"]
122         # Allocate maximum attention_mask
123         attention_mask = input_ids.new_zeros(
124             (pb.size, max_input_length + padding_right_offset)
125         )
126         # Copy tokenizer attention_mask into fully allocated attention_mask
127         attention_mask[:, :max_input_length] = tokenized_inputs["attention_mask"]
128 
129         position_ids = tokenized_inputs["attention_mask"].long().cumsum(-1) - 1
130         position_ids.masked_fill_(tokenized_inputs["attention_mask"] == 0, 1)
131         all_input_ids = tokenized_inputs["input_ids"].T.split(1, dim=1)
132 
133         max_tokens = len(inputs) * max_input_length + max_decode_tokens
134 
135         return cls(
136             batch_id=pb.id,
137             requests=pb.requests,
138             requests_idx_mapping=requests_idx_mapping,
139             input_ids=input_ids,
140             attention_mask=attention_mask,
141             position_ids=position_ids,
142             past_key_values=None,
143             all_input_ids=list(all_input_ids),
144             input_lengths=input_lengths.tolist(),
145             prefix_offsets=prefix_offsets,
146             read_offsets=read_offsets,
147             next_token_choosers=next_token_choosers,
148             stopping_criterias=stopping_criterias,
149             max_input_length=max_input_length.item(),
150             padding_right_offset=padding_right_offset,
151             max_tokens=max_tokens,
152         )
153 
154 
155 class GalacticaSharded(CausalLM):
156     def __init__(
157         self,
158         model_id: str,
159         revision: Optional[str] = None,
160         quantize: Optional[str] = None,
161         dtype: Optional[torch.dtype] = None,
162         trust_remote_code: bool = False,
163     ):
164         self.process_group, rank, world_size = initialize_torch_distributed()
165         if torch.cuda.is_available():
166             device = torch.device(f"cuda:{rank}")
167             dtype = torch.float16 if dtype is None else dtype
168         else:
169             device = torch.device("cpu")
170             dtype = torch.float32
171 
172         tokenizer = AutoTokenizer.from_pretrained(
173             model_id,
174             revision=revision,
175             padding_side="left",
176             truncation_side="left",
177             trust_remote_code=trust_remote_code,
178         )
179 
180         config = AutoConfig.from_pretrained(
181             model_id,
182             revision=revision,
183             tp_parallel=True,
184             trust_remote_code=trust_remote_code,
185         )
186         config.quantize = quantize
187         tokenizer.pad_token_id = config.pad_token_id
188 
189         torch.distributed.barrier(group=self.process_group)
190         filenames = weight_files(model_id, revision=revision, extension=".safetensors")
191         weights = Weights(
192             filenames, device=device, dtype=dtype, process_group=self.process_group
193         )
194         if config.quantize == "gptq":
195             weights._set_gptq_params(model_id)
196 
197         model = OPTForCausalLM(config, weights)
198 
199         torch.distributed.barrier(group=self.process_group)
200         super(CausalLM, self).__init__(
201             model=model,
202             tokenizer=tokenizer,
203             requires_padding=True,
204             dtype=dtype,
205             device=device,
206             rank=rank,
207             world_size=world_size,
208         )
209 
210     @property
211     def batch_type(self) -> Type[CausalLMBatch]:
212         return GalacticaCausalLMBatch
213 
214     def decode(self, generated_ids: List[int]) -> str:
215         # Do not skip special tokens as they are used for custom parsing rules of the generated text
216         return self.tokenizer.decode(
217             generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False
218         )
219 
220     def forward(
221         self, input_ids, attention_mask, position_ids, past_key_values: Optional = None
222     ):
223         outputs = self.model.forward(
224             input_ids=input_ids,
225             attention_mask=attention_mask,
226             past_key_values=past_key_values,
227             use_cache=True,
228         )
229         return outputs.logits, outputs.past_key_values
230
```
--- 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/galactica.py b/server/text_generation_server/models/galactica.py
--- a/server/text_generation_server/models/galactica.py
+++ b/server/text_generation_server/models/galactica.py
@@ -80,6 +80,7 @@
         next_token_choosers = []
         stopping_criterias = []
         prefix_offsets = []
+        top_n_tokens = []
         read_offsets = []
         requests_idx_mapping = {}
 
@@ -96,6 +97,7 @@
                 r.stopping_parameters, tokenizer
             )
             stopping_criterias.append(stopping_criteria)
+            top_n_tokens.append(r.top_n_tokens)
             max_truncation = max(max_truncation, r.truncate)
             max_decode_tokens += stopping_criteria.max_new_tokens
             padding_right_offset = max(
@@ -129,6 +131,9 @@
         position_ids = tokenized_inputs["attention_mask"].long().cumsum(-1) - 1
         position_ids.masked_fill_(tokenized_inputs["attention_mask"] == 0, 1)
         all_input_ids = tokenized_inputs["input_ids"].T.split(1, dim=1)
+        top_n_tokens_tensor = torch.tensor(
+            top_n_tokens, device=device, dtype=torch.int64
+        )
 
         max_tokens = len(inputs) * max_input_length + max_decode_tokens
 
@@ -146,6 +151,8 @@
             read_offsets=read_offsets,
             next_token_choosers=next_token_choosers,
             stopping_criterias=stopping_criterias,
+            top_n_tokens=top_n_tokens,
+            top_n_tokens_tensor=top_n_tokens_tensor,
             max_input_length=max_input_length.item(),
             padding_right_offset=padding_right_offset,
             max_tokens=max_tokens,

{"golden_diff": "diff --git a/server/text_generation_server/models/galactica.py b/server/text_generation_server/models/galactica.py\n--- a/server/text_generation_server/models/galactica.py\n+++ b/server/text_generation_server/models/galactica.py\n@@ -80,6 +80,7 @@\n         next_token_choosers = []\n         stopping_criterias = []\n         prefix_offsets = []\n+        top_n_tokens = []\n         read_offsets = []\n         requests_idx_mapping = {}\n \n@@ -96,6 +97,7 @@\n                 r.stopping_parameters, tokenizer\n             )\n             stopping_criterias.append(stopping_criteria)\n+            top_n_tokens.append(r.top_n_tokens)\n             max_truncation = max(max_truncation, r.truncate)\n             max_decode_tokens += stopping_criteria.max_new_tokens\n             padding_right_offset = max(\n@@ -129,6 +131,9 @@\n         position_ids = tokenized_inputs[\"attention_mask\"].long().cumsum(-1) - 1\n         position_ids.masked_fill_(tokenized_inputs[\"attention_mask\"] == 0, 1)\n         all_input_ids = tokenized_inputs[\"input_ids\"].T.split(1, dim=1)\n+        top_n_tokens_tensor = torch.tensor(\n+            top_n_tokens, device=device, dtype=torch.int64\n+        )\n \n         max_tokens = len(inputs) * max_input_length + max_decode_tokens\n \n@@ -146,6 +151,8 @@\n             read_offsets=read_offsets,\n             next_token_choosers=next_token_choosers,\n             stopping_criterias=stopping_criterias,\n+            top_n_tokens=top_n_tokens,\n+            top_n_tokens_tensor=top_n_tokens_tensor,\n             max_input_length=max_input_length.item(),\n             padding_right_offset=padding_right_offset,\n             max_tokens=max_tokens,\n", "issue": "missing 2 required positional arguments: 'top_n_tokens' and 'top_n_tokens_tensor' for the galactica model\n### System Info\r\n\r\nv1.0.3\r\n\r\n### Information\r\n\r\n- [ ] Docker\r\n- [ ] The CLI directly\r\n\r\n### Tasks\r\n\r\n- [ ] An officially supported command\r\n- [ ] My own modifications\r\n\r\n### Reproduction\r\n\r\nhaving a local galactica model in a folder called `facebook-galactica-30b-gptq` won't be detected since it will fail this check https://github.com/huggingface/text-generation-inference/blob/main/server/text_generation_server/models/__init__.py#L88\r\n\r\nI suggest making it check `if \"galactica\" in  in model_id` instead.\r\n\r\n### Expected behavior\r\n\r\nexpected to detect a galactica model\n", "before_files": [{"content": "import re\nimport torch\nimport torch.distributed\n\nfrom typing import List, Optional, Type\n\nfrom transformers import (\n    AutoTokenizer,\n    AutoConfig,\n    PreTrainedTokenizerBase,\n)\nfrom text_generation_server.models import CausalLM\nfrom text_generation_server.models.causal_lm import CausalLMBatch\nfrom text_generation_server.pb import generate_pb2\nfrom text_generation_server.models.custom_modeling.opt_modeling import OPTForCausalLM\nfrom text_generation_server.utils import (\n    NextTokenChooser,\n    StoppingCriteria,\n    initialize_torch_distributed,\n    weight_files,\n    Weights,\n)\n\n# CREDIT: Papers with code => https://github.com/paperswithcode/galai/blob/main/galai/utils.py\n\n# we split individual characters inside special tokens like [START_DNA]\nCUSTOM_SEQ_RE = re.compile(r\"(\\[START_(DNA|SMILES|I_SMILES|AMINO)])(.*?)(\\[END_\\2])\")\n\n# token added to implement a custom sequence tokenization. This token is added at\n# corpus cleaning step and removed in pretokenization. The digits are added to increase the chance\n# that they do not occur in the corpus. The digits are escaped so that the token does not appear\n# literally in the source code in case we ever include it in the training data.\nSPLIT_MARKER = f\"SPL{1}T-TH{1}S-Pl3A5E\"\n\n\ndef _insert_split_marker(m: re.Match):\n    \"\"\"\n    Applies split marker based on a regex match of special tokens such as\n    [START_DNA].\n    Parameters\n    ----------\n    n : str\n        Input text to split\n    Returns\n    ----------\n    str - the text with the split token added\n    \"\"\"\n    start_token, _, sequence, end_token = m.groups()\n    sequence = re.sub(r\"(.)\", rf\"{SPLIT_MARKER}\\1\", sequence, flags=re.DOTALL)\n    return f\"{start_token}{sequence}{SPLIT_MARKER}{end_token}\"\n\n\ndef escape_custom_split_sequence(text):\n    \"\"\"\n    Applies custom splitting to the text for GALILEO's tokenization\n    Parameters\n    ----------\n    text : str\n        Input text to split\n    Returns\n    ----------\n    str - the text with the split token added\n    \"\"\"\n    return CUSTOM_SEQ_RE.sub(_insert_split_marker, text)\n\n\n# END CREDIT\n\n\nclass GalacticaCausalLMBatch(CausalLMBatch):\n    @classmethod\n    def from_pb(\n        cls,\n        pb: generate_pb2.Batch,\n        tokenizer: PreTrainedTokenizerBase,\n        dtype: torch.dtype,\n        device: torch.device,\n    ) -> \"GalacticaCausalLMBatch\":\n        inputs = []\n        next_token_choosers = []\n        stopping_criterias = []\n        prefix_offsets = []\n        read_offsets = []\n        requests_idx_mapping = {}\n\n        # Parse batch\n        max_truncation = 0\n        padding_right_offset = 0\n        max_decode_tokens = 0\n        for i, r in enumerate(pb.requests):\n            requests_idx_mapping[r.id] = i\n            # Add escape_custom_split_sequence to the CausalLMBatch logic\n            inputs.append(escape_custom_split_sequence(r.inputs))\n            next_token_choosers.append(NextTokenChooser.from_pb(r.parameters, device))\n            stopping_criteria = StoppingCriteria.from_pb(\n                r.stopping_parameters, tokenizer\n            )\n            stopping_criterias.append(stopping_criteria)\n            max_truncation = max(max_truncation, r.truncate)\n            max_decode_tokens += stopping_criteria.max_new_tokens\n            padding_right_offset = max(\n                padding_right_offset, stopping_criteria.max_new_tokens\n            )\n\n        tokenized_inputs = tokenizer(\n            inputs,\n            return_tensors=\"pt\",\n            padding=True,\n            return_token_type_ids=False,\n            truncation=True,\n            max_length=max_truncation,\n        ).to(device)\n        for _ in pb.requests:\n            input_len = tokenized_inputs[\"input_ids\"].shape[1]\n            prefix_offsets.append(0)\n            read_offsets.append(input_len)\n\n        input_lengths = tokenized_inputs[\"attention_mask\"].sum(1)\n        max_input_length = input_lengths.max()\n\n        input_ids = tokenized_inputs[\"input_ids\"]\n        # Allocate maximum attention_mask\n        attention_mask = input_ids.new_zeros(\n            (pb.size, max_input_length + padding_right_offset)\n        )\n        # Copy tokenizer attention_mask into fully allocated attention_mask\n        attention_mask[:, :max_input_length] = tokenized_inputs[\"attention_mask\"]\n\n        position_ids = tokenized_inputs[\"attention_mask\"].long().cumsum(-1) - 1\n        position_ids.masked_fill_(tokenized_inputs[\"attention_mask\"] == 0, 1)\n        all_input_ids = tokenized_inputs[\"input_ids\"].T.split(1, dim=1)\n\n        max_tokens = len(inputs) * max_input_length + max_decode_tokens\n\n        return cls(\n            batch_id=pb.id,\n            requests=pb.requests,\n            requests_idx_mapping=requests_idx_mapping,\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            position_ids=position_ids,\n            past_key_values=None,\n            all_input_ids=list(all_input_ids),\n            input_lengths=input_lengths.tolist(),\n            prefix_offsets=prefix_offsets,\n            read_offsets=read_offsets,\n            next_token_choosers=next_token_choosers,\n            stopping_criterias=stopping_criterias,\n            max_input_length=max_input_length.item(),\n            padding_right_offset=padding_right_offset,\n            max_tokens=max_tokens,\n        )\n\n\nclass GalacticaSharded(CausalLM):\n    def __init__(\n        self,\n        model_id: str,\n        revision: Optional[str] = None,\n        quantize: Optional[str] = None,\n        dtype: Optional[torch.dtype] = None,\n        trust_remote_code: bool = False,\n    ):\n        self.process_group, rank, world_size = initialize_torch_distributed()\n        if torch.cuda.is_available():\n            device = torch.device(f\"cuda:{rank}\")\n            dtype = torch.float16 if dtype is None else dtype\n        else:\n            device = torch.device(\"cpu\")\n            dtype = torch.float32\n\n        tokenizer = AutoTokenizer.from_pretrained(\n            model_id,\n            revision=revision,\n            padding_side=\"left\",\n            truncation_side=\"left\",\n            trust_remote_code=trust_remote_code,\n        )\n\n        config = AutoConfig.from_pretrained(\n            model_id,\n            revision=revision,\n            tp_parallel=True,\n            trust_remote_code=trust_remote_code,\n        )\n        config.quantize = quantize\n        tokenizer.pad_token_id = config.pad_token_id\n\n        torch.distributed.barrier(group=self.process_group)\n        filenames = weight_files(model_id, revision=revision, extension=\".safetensors\")\n        weights = Weights(\n            filenames, device=device, dtype=dtype, process_group=self.process_group\n        )\n        if config.quantize == \"gptq\":\n            weights._set_gptq_params(model_id)\n\n        model = OPTForCausalLM(config, weights)\n\n        torch.distributed.barrier(group=self.process_group)\n        super(CausalLM, self).__init__(\n            model=model,\n            tokenizer=tokenizer,\n            requires_padding=True,\n            dtype=dtype,\n            device=device,\n            rank=rank,\n            world_size=world_size,\n        )\n\n    @property\n    def batch_type(self) -> Type[CausalLMBatch]:\n        return GalacticaCausalLMBatch\n\n    def decode(self, generated_ids: List[int]) -> str:\n        # Do not skip special tokens as they are used for custom parsing rules of the generated text\n        return self.tokenizer.decode(\n            generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False\n        )\n\n    def forward(\n        self, input_ids, attention_mask, position_ids, past_key_values: Optional = None\n    ):\n        outputs = self.model.forward(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            past_key_values=past_key_values,\n            use_cache=True,\n        )\n        return outputs.logits, outputs.past_key_values\n", "path": "server/text_generation_server/models/galactica.py"}], "after_files": [{"content": "import re\nimport torch\nimport torch.distributed\n\nfrom typing import List, Optional, Type\n\nfrom transformers import (\n    AutoTokenizer,\n    AutoConfig,\n    PreTrainedTokenizerBase,\n)\nfrom text_generation_server.models import CausalLM\nfrom text_generation_server.models.causal_lm import CausalLMBatch\nfrom text_generation_server.pb import generate_pb2\nfrom text_generation_server.models.custom_modeling.opt_modeling import OPTForCausalLM\nfrom text_generation_server.utils import (\n    NextTokenChooser,\n    StoppingCriteria,\n    initialize_torch_distributed,\n    weight_files,\n    Weights,\n)\n\n# CREDIT: Papers with code => https://github.com/paperswithcode/galai/blob/main/galai/utils.py\n\n# we split individual characters inside special tokens like [START_DNA]\nCUSTOM_SEQ_RE = re.compile(r\"(\\[START_(DNA|SMILES|I_SMILES|AMINO)])(.*?)(\\[END_\\2])\")\n\n# token added to implement a custom sequence tokenization. This token is added at\n# corpus cleaning step and removed in pretokenization. The digits are added to increase the chance\n# that they do not occur in the corpus. The digits are escaped so that the token does not appear\n# literally in the source code in case we ever include it in the training data.\nSPLIT_MARKER = f\"SPL{1}T-TH{1}S-Pl3A5E\"\n\n\ndef _insert_split_marker(m: re.Match):\n    \"\"\"\n    Applies split marker based on a regex match of special tokens such as\n    [START_DNA].\n    Parameters\n    ----------\n    n : str\n        Input text to split\n    Returns\n    ----------\n    str - the text with the split token added\n    \"\"\"\n    start_token, _, sequence, end_token = m.groups()\n    sequence = re.sub(r\"(.)\", rf\"{SPLIT_MARKER}\\1\", sequence, flags=re.DOTALL)\n    return f\"{start_token}{sequence}{SPLIT_MARKER}{end_token}\"\n\n\ndef escape_custom_split_sequence(text):\n    \"\"\"\n    Applies custom splitting to the text for GALILEO's tokenization\n    Parameters\n    ----------\n    text : str\n        Input text to split\n    Returns\n    ----------\n    str - the text with the split token added\n    \"\"\"\n    return CUSTOM_SEQ_RE.sub(_insert_split_marker, text)\n\n\n# END CREDIT\n\n\nclass GalacticaCausalLMBatch(CausalLMBatch):\n    @classmethod\n    def from_pb(\n        cls,\n        pb: generate_pb2.Batch,\n        tokenizer: PreTrainedTokenizerBase,\n        dtype: torch.dtype,\n        device: torch.device,\n    ) -> \"GalacticaCausalLMBatch\":\n        inputs = []\n        next_token_choosers = []\n        stopping_criterias = []\n        prefix_offsets = []\n        top_n_tokens = []\n        read_offsets = []\n        requests_idx_mapping = {}\n\n        # Parse batch\n        max_truncation = 0\n        padding_right_offset = 0\n        max_decode_tokens = 0\n        for i, r in enumerate(pb.requests):\n            requests_idx_mapping[r.id] = i\n            # Add escape_custom_split_sequence to the CausalLMBatch logic\n            inputs.append(escape_custom_split_sequence(r.inputs))\n            next_token_choosers.append(NextTokenChooser.from_pb(r.parameters, device))\n            stopping_criteria = StoppingCriteria.from_pb(\n                r.stopping_parameters, tokenizer\n            )\n            stopping_criterias.append(stopping_criteria)\n            top_n_tokens.append(r.top_n_tokens)\n            max_truncation = max(max_truncation, r.truncate)\n            max_decode_tokens += stopping_criteria.max_new_tokens\n            padding_right_offset = max(\n                padding_right_offset, stopping_criteria.max_new_tokens\n            )\n\n        tokenized_inputs = tokenizer(\n            inputs,\n            return_tensors=\"pt\",\n            padding=True,\n            return_token_type_ids=False,\n            truncation=True,\n            max_length=max_truncation,\n        ).to(device)\n        for _ in pb.requests:\n            input_len = tokenized_inputs[\"input_ids\"].shape[1]\n            prefix_offsets.append(0)\n            read_offsets.append(input_len)\n\n        input_lengths = tokenized_inputs[\"attention_mask\"].sum(1)\n        max_input_length = input_lengths.max()\n\n        input_ids = tokenized_inputs[\"input_ids\"]\n        # Allocate maximum attention_mask\n        attention_mask = input_ids.new_zeros(\n            (pb.size, max_input_length + padding_right_offset)\n        )\n        # Copy tokenizer attention_mask into fully allocated attention_mask\n        attention_mask[:, :max_input_length] = tokenized_inputs[\"attention_mask\"]\n\n        position_ids = tokenized_inputs[\"attention_mask\"].long().cumsum(-1) - 1\n        position_ids.masked_fill_(tokenized_inputs[\"attention_mask\"] == 0, 1)\n        all_input_ids = tokenized_inputs[\"input_ids\"].T.split(1, dim=1)\n        top_n_tokens_tensor = torch.tensor(\n            top_n_tokens, device=device, dtype=torch.int64\n        )\n\n        max_tokens = len(inputs) * max_input_length + max_decode_tokens\n\n        return cls(\n            batch_id=pb.id,\n            requests=pb.requests,\n            requests_idx_mapping=requests_idx_mapping,\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            position_ids=position_ids,\n            past_key_values=None,\n            all_input_ids=list(all_input_ids),\n            input_lengths=input_lengths.tolist(),\n            prefix_offsets=prefix_offsets,\n            read_offsets=read_offsets,\n            next_token_choosers=next_token_choosers,\n            stopping_criterias=stopping_criterias,\n            top_n_tokens=top_n_tokens,\n            top_n_tokens_tensor=top_n_tokens_tensor,\n            max_input_length=max_input_length.item(),\n            padding_right_offset=padding_right_offset,\n            max_tokens=max_tokens,\n        )\n\n\nclass GalacticaSharded(CausalLM):\n    def __init__(\n        self,\n        model_id: str,\n        revision: Optional[str] = None,\n        quantize: Optional[str] = None,\n        dtype: Optional[torch.dtype] = None,\n        trust_remote_code: bool = False,\n    ):\n        self.process_group, rank, world_size = initialize_torch_distributed()\n        if torch.cuda.is_available():\n            device = torch.device(f\"cuda:{rank}\")\n            dtype = torch.float16 if dtype is None else dtype\n        else:\n            device = torch.device(\"cpu\")\n            dtype = torch.float32\n\n        tokenizer = AutoTokenizer.from_pretrained(\n            model_id,\n            revision=revision,\n            padding_side=\"left\",\n            truncation_side=\"left\",\n            trust_remote_code=trust_remote_code,\n        )\n\n        config = AutoConfig.from_pretrained(\n            model_id,\n            revision=revision,\n            tp_parallel=True,\n            trust_remote_code=trust_remote_code,\n        )\n        config.quantize = quantize\n        tokenizer.pad_token_id = config.pad_token_id\n\n        torch.distributed.barrier(group=self.process_group)\n        filenames = weight_files(model_id, revision=revision, extension=\".safetensors\")\n        weights = Weights(\n            filenames, device=device, dtype=dtype, process_group=self.process_group\n        )\n        if config.quantize == \"gptq\":\n            weights._set_gptq_params(model_id)\n\n        model = OPTForCausalLM(config, weights)\n\n        torch.distributed.barrier(group=self.process_group)\n        super(CausalLM, self).__init__(\n            model=model,\n            tokenizer=tokenizer,\n            requires_padding=True,\n            dtype=dtype,\n            device=device,\n            rank=rank,\n            world_size=world_size,\n        )\n\n    @property\n    def batch_type(self) -> Type[CausalLMBatch]:\n        return GalacticaCausalLMBatch\n\n    def decode(self, generated_ids: List[int]) -> str:\n        # Do not skip special tokens as they are used for custom parsing rules of the generated text\n        return self.tokenizer.decode(\n            generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False\n        )\n\n    def forward(\n        self, input_ids, attention_mask, position_ids, past_key_values: Optional = None\n    ):\n        outputs = self.model.forward(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            past_key_values=past_key_values,\n            use_cache=True,\n        )\n        return outputs.logits, outputs.past_key_values\n", "path": "server/text_generation_server/models/galactica.py"}]}