optimum-rbln 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- optimum/rbln/__init__.py +115 -0
- optimum/rbln/__version__.py +1 -0
- optimum/rbln/diffusers/__init__.py +64 -0
- optimum/rbln/diffusers/models/__init__.py +26 -0
- optimum/rbln/diffusers/models/autoencoder_kl.py +313 -0
- optimum/rbln/diffusers/models/controlnet.py +180 -0
- optimum/rbln/diffusers/models/unet_2d_condition.py +352 -0
- optimum/rbln/diffusers/pipelines/__init__.py +30 -0
- optimum/rbln/diffusers/pipelines/controlnet/__init__.py +24 -0
- optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py +266 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion/__init__.py +26 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_controlnet_img2img.py +731 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +106 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +116 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion_xl/__init__.py +2 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +109 -0
- optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +111 -0
- optimum/rbln/modeling.py +0 -0
- optimum/rbln/modeling_alias.py +49 -0
- optimum/rbln/modeling_base.py +645 -0
- optimum/rbln/modeling_config.py +169 -0
- optimum/rbln/modeling_seq2seq.py +469 -0
- optimum/rbln/transformers/__init__.py +59 -0
- optimum/rbln/transformers/generation/__init__.py +24 -0
- optimum/rbln/transformers/generation/streamers.py +122 -0
- optimum/rbln/transformers/models/__init__.py +28 -0
- optimum/rbln/transformers/models/bart/__init__.py +24 -0
- optimum/rbln/transformers/models/bart/bart_architecture.py +377 -0
- optimum/rbln/transformers/models/clip/__init__.py +24 -0
- optimum/rbln/transformers/models/clip/modeling_clip.py +116 -0
- optimum/rbln/transformers/models/gpt2/__init__.py +24 -0
- optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +253 -0
- optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +700 -0
- optimum/rbln/transformers/models/llama/__init__.py +24 -0
- optimum/rbln/transformers/models/llama/llama_architecture.py +607 -0
- optimum/rbln/transformers/models/llama/modeling_llama.py +409 -0
- optimum/rbln/transformers/models/t5/__init__.py +24 -0
- optimum/rbln/transformers/models/t5/t5_architecture.py +439 -0
- optimum/rbln/transformers/models/wav2vec2/__init__.py +24 -0
- optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py +121 -0
- optimum/rbln/transformers/models/whisper/__init__.py +24 -0
- optimum/rbln/transformers/models/whisper/modeling_whisper.py +374 -0
- optimum/rbln/transformers/models/whisper/whisper_architecture.py +406 -0
- optimum/rbln/utils/__init__.py +25 -0
- optimum/rbln/utils/import_utils.py +28 -0
- optimum/rbln/utils/runtime_utils.py +71 -0
- optimum/rbln/utils/save_utils.py +92 -0
- optimum_rbln-0.1.0.dist-info/METADATA +144 -0
- optimum_rbln-0.1.0.dist-info/RECORD +51 -0
- optimum_rbln-0.1.0.dist-info/WHEEL +4 -0
- optimum_rbln-0.1.0.dist-info/licenses/LICENSE +201 -0
@@ -0,0 +1,409 @@
|
|
1
|
+
# Copyright 2024 Rebellions Inc.
|
2
|
+
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at:
|
6
|
+
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
+
# See the License for the specific language governing permissions and
|
13
|
+
# limitations under the License.
|
14
|
+
|
15
|
+
# Portions of this software are licensed under the Apache License,
|
16
|
+
# Version 2.0. See the NOTICE file distributed with this work for
|
17
|
+
# additional information regarding copyright ownership.
|
18
|
+
|
19
|
+
# All other portions of this software, including proprietary code,
|
20
|
+
# are the intellectual property of Rebellions Inc. and may not be
|
21
|
+
# copied, modified, or distributed without prior written permission
|
22
|
+
# from Rebellions Inc.
|
23
|
+
|
24
|
+
import inspect # noqa: I001
|
25
|
+
import logging
|
26
|
+
from pathlib import Path
|
27
|
+
from tempfile import TemporaryDirectory
|
28
|
+
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
|
29
|
+
|
30
|
+
import torch # noqa: F401
|
31
|
+
import rebel # noqa: F401
|
32
|
+
|
33
|
+
from optimum.exporters import TasksManager
|
34
|
+
from transformers import AutoModelForCausalLM, LlamaForCausalLM, PretrainedConfig, AutoConfig
|
35
|
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
36
|
+
|
37
|
+
from ....modeling_base import RBLNBaseModel
|
38
|
+
from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
|
39
|
+
from ....utils.runtime_utils import RBLNPytorchRuntime
|
40
|
+
from ....utils.save_utils import maybe_save_preprocessors
|
41
|
+
from .llama_architecture import (
|
42
|
+
LlamaWrapper,
|
43
|
+
wrap_llama,
|
44
|
+
unwrap_llama,
|
45
|
+
)
|
46
|
+
|
47
|
+
|
48
|
+
logger = logging.getLogger(__name__)
|
49
|
+
|
50
|
+
if TYPE_CHECKING:
|
51
|
+
from transformers import (
|
52
|
+
AutoFeatureExtractor,
|
53
|
+
AutoProcessor,
|
54
|
+
AutoTokenizer,
|
55
|
+
PretrainedConfig,
|
56
|
+
)
|
57
|
+
|
58
|
+
|
59
|
+
class RBLNRuntimeModel(RBLNPytorchRuntime):
|
60
|
+
mandatory_members = ["main_input_name"]
|
61
|
+
|
62
|
+
# RBLN_Runtimemodule
|
63
|
+
def forward(
|
64
|
+
self,
|
65
|
+
input_ids: torch.LongTensor = None,
|
66
|
+
attention_mask: torch.LongTensor = None,
|
67
|
+
cache_position: torch.Tensor = None,
|
68
|
+
**kwargs: Dict[str, Any],
|
69
|
+
):
|
70
|
+
logits = super().forward(
|
71
|
+
input_ids=input_ids,
|
72
|
+
attention_mask=attention_mask,
|
73
|
+
cache_position=cache_position,
|
74
|
+
)
|
75
|
+
return logits
|
76
|
+
|
77
|
+
|
78
|
+
class RBLNLlamaForCausalLM(RBLNBaseModel):
|
79
|
+
"""
|
80
|
+
The Llama Model transformer with a language modeling head (linear layer) on top.
|
81
|
+
This model inherits from [`RBLNBaseModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
|
82
|
+
|
83
|
+
A class to convert and run pre-trained transformers based LlamaForCausalLM model on RBLN devices.
|
84
|
+
It implements the methods to convert a pre-trained transformers LlamaForCausalLM model into a RBLN transformer model by:
|
85
|
+
- transferring the checkpoint weights of the original into an optimized RBLN graph,
|
86
|
+
- compiling the resulting graph using the RBLN compiler.
|
87
|
+
"""
|
88
|
+
|
89
|
+
model_type = "rbln_model"
|
90
|
+
main_input_name = "input_ids"
|
91
|
+
auto_model_class = AutoModelForCausalLM
|
92
|
+
|
93
|
+
def __post_init__(self, **kwargs):
|
94
|
+
|
95
|
+
self.batch_size = self.rbln_config.meta["rbln_batch_size"]
|
96
|
+
self.max_seq_len = self.rbln_config.meta["rbln_max_seq_len"]
|
97
|
+
self.prefill_chunk_size = self.rbln_config.meta["rbln_prefill_chunk_size"]
|
98
|
+
|
99
|
+
self.prefill_attention_mask = torch.zeros(
|
100
|
+
self.batch_size, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64
|
101
|
+
)
|
102
|
+
self.causal_mask = 1 - torch.triu(
|
103
|
+
torch.ones(self.batch_size, 1, self.prefill_chunk_size, self.prefill_chunk_size), diagonal=1
|
104
|
+
)
|
105
|
+
|
106
|
+
self.prefill_decoder = RBLNRuntimeModel(runtime=self.runtimes[0], main_input_name="input_ids")
|
107
|
+
self.decoder = RBLNRuntimeModel(runtime=self.runtimes[1], main_input_name="input_ids")
|
108
|
+
self.past_cached_length = 0
|
109
|
+
|
110
|
+
@classmethod
|
111
|
+
@torch.no_grad()
|
112
|
+
def _export(
|
113
|
+
cls,
|
114
|
+
model_id: str,
|
115
|
+
config: "PretrainedConfig",
|
116
|
+
use_auth_token: Optional[Union[bool, str]] = None,
|
117
|
+
revision: Optional[str] = None,
|
118
|
+
force_download: bool = False,
|
119
|
+
cache_dir: Optional[str] = None,
|
120
|
+
subfolder: str = "",
|
121
|
+
local_files_only: bool = False,
|
122
|
+
trust_remote_code: bool = False,
|
123
|
+
**kwargs,
|
124
|
+
) -> "RBLNLlamaForCausalLM":
|
125
|
+
task = kwargs.pop("task", None)
|
126
|
+
if task is None:
|
127
|
+
task = TasksManager.infer_task_from_model(cls.auto_model_class)
|
128
|
+
|
129
|
+
save_dir = TemporaryDirectory()
|
130
|
+
save_dir_path = Path(save_dir.name)
|
131
|
+
|
132
|
+
def update_configs(kwargs):
|
133
|
+
hf_max_position_embeddings = getattr(AutoConfig.from_pretrained(model_id), "max_position_embeddings", None)
|
134
|
+
max_seq_len = kwargs.get("rbln_max_seq_len", None)
|
135
|
+
if max_seq_len is not None:
|
136
|
+
if max_seq_len <= hf_max_position_embeddings:
|
137
|
+
kwargs.update({"max_position_embeddings": max_seq_len})
|
138
|
+
else:
|
139
|
+
raise ValueError("`max_seq_len` should be less or equal than max_position_embeddings!")
|
140
|
+
|
141
|
+
kwargs.update(
|
142
|
+
{
|
143
|
+
"torchscript": True,
|
144
|
+
"return_dict": False,
|
145
|
+
"use_cache": True,
|
146
|
+
"torch_dtype": torch.float32,
|
147
|
+
"_attn_implementation": "eager",
|
148
|
+
}
|
149
|
+
)
|
150
|
+
|
151
|
+
return kwargs
|
152
|
+
|
153
|
+
kwargs = update_configs(kwargs)
|
154
|
+
|
155
|
+
rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
|
156
|
+
|
157
|
+
origin_mehtods = wrap_llama()
|
158
|
+
model: LlamaForCausalLM = TasksManager.get_model_from_task(
|
159
|
+
task=task,
|
160
|
+
model_name_or_path=model_id,
|
161
|
+
subfolder=subfolder,
|
162
|
+
revision=revision,
|
163
|
+
framework="pt",
|
164
|
+
cache_dir=cache_dir,
|
165
|
+
use_auth_token=use_auth_token,
|
166
|
+
local_files_only=local_files_only,
|
167
|
+
force_download=force_download,
|
168
|
+
trust_remote_code=trust_remote_code,
|
169
|
+
**kwargs,
|
170
|
+
)
|
171
|
+
|
172
|
+
if config is None:
|
173
|
+
config = model.config
|
174
|
+
|
175
|
+
config.save_pretrained(save_dir_path)
|
176
|
+
preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
|
177
|
+
|
178
|
+
# Get compilation arguments
|
179
|
+
if rbln_config_kwargs.get("rbln_config", None) is None:
|
180
|
+
rbln_config = cls.get_rbln_config(
|
181
|
+
preprocessors=preprocessors, model_config=model.config, **rbln_config_kwargs
|
182
|
+
)
|
183
|
+
|
184
|
+
def compile_llama():
|
185
|
+
wrapped_model = LlamaWrapper(model).eval()
|
186
|
+
|
187
|
+
prefill_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
|
188
|
+
dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
|
189
|
+
|
190
|
+
prefill_example_inputs = prefill_rbln_runtime_config.get_dummy_inputs(fill=0)
|
191
|
+
dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=0)
|
192
|
+
|
193
|
+
prefill_scripted_model = torch.jit.trace(wrapped_model, prefill_example_inputs)
|
194
|
+
dec_scripted_model = torch.jit.trace(wrapped_model, dec_example_inputs)
|
195
|
+
|
196
|
+
prefill_ir = rebel.torchscript_to_ir(
|
197
|
+
prefill_scripted_model,
|
198
|
+
input_names=[v[0] for v in prefill_rbln_runtime_config.input_info],
|
199
|
+
)
|
200
|
+
dec_ir = rebel.torchscript_to_ir(
|
201
|
+
dec_scripted_model,
|
202
|
+
input_names=[v[0] for v in dec_rbln_runtime_config.input_info],
|
203
|
+
)
|
204
|
+
|
205
|
+
# Caching prefill_decoder/decoder I/O
|
206
|
+
connections = [
|
207
|
+
(prefill_ir.outputs[1 + i], prefill_ir.inputs[3 + i])
|
208
|
+
for i in range(model.config.num_hidden_layers * 2)
|
209
|
+
]
|
210
|
+
|
211
|
+
compiled_model = rebel.compile(
|
212
|
+
prefill_ir,
|
213
|
+
dec_ir,
|
214
|
+
connections=connections,
|
215
|
+
fusion=prefill_rbln_runtime_config.fusion,
|
216
|
+
npu=prefill_rbln_runtime_config.npu,
|
217
|
+
tensor_parallel_size=prefill_rbln_runtime_config.tensor_parallel_size,
|
218
|
+
use_weight_sharing=True,
|
219
|
+
)
|
220
|
+
compiled_model.save(save_dir_path / f"{DEFAULT_COMPILED_MODEL_NAME}.rbln")
|
221
|
+
|
222
|
+
compile_llama()
|
223
|
+
unwrap_llama(origin_mehtods)
|
224
|
+
|
225
|
+
rbln_config.save(save_dir_path)
|
226
|
+
|
227
|
+
return cls._from_pretrained(
|
228
|
+
model_id=save_dir_path,
|
229
|
+
config=config,
|
230
|
+
model_save_dir=save_dir,
|
231
|
+
**rbln_constructor_kwargs,
|
232
|
+
**kwargs,
|
233
|
+
)
|
234
|
+
|
235
|
+
@classmethod
|
236
|
+
def _get_rbln_config(
|
237
|
+
cls,
|
238
|
+
preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
|
239
|
+
model_config: "PretrainedConfig",
|
240
|
+
rbln_max_seq_len: Optional[int] = None,
|
241
|
+
rbln_batch_size: Optional[int] = None,
|
242
|
+
) -> RBLNConfig:
|
243
|
+
meta = {}
|
244
|
+
|
245
|
+
prefill_chunk_size = 128
|
246
|
+
if rbln_max_seq_len is None:
|
247
|
+
rbln_max_seq_len = getattr(model_config, "max_position_embeddings", None)
|
248
|
+
|
249
|
+
meta["rbln_max_seq_len"] = rbln_max_seq_len
|
250
|
+
meta["rbln_batch_size"] = rbln_batch_size
|
251
|
+
meta["rbln_prefill_chunk_size"] = prefill_chunk_size
|
252
|
+
|
253
|
+
def get_input_info(query_length):
|
254
|
+
input_info = [
|
255
|
+
("input_ids", [rbln_batch_size, query_length], "int64"),
|
256
|
+
("attention_mask", [rbln_batch_size, 1, query_length, rbln_max_seq_len], "int64"),
|
257
|
+
(
|
258
|
+
"cache_position",
|
259
|
+
[],
|
260
|
+
"int32",
|
261
|
+
),
|
262
|
+
]
|
263
|
+
input_info.extend(
|
264
|
+
[
|
265
|
+
(
|
266
|
+
f"past_key_values_{i}",
|
267
|
+
[
|
268
|
+
rbln_batch_size,
|
269
|
+
model_config.num_key_value_heads,
|
270
|
+
rbln_max_seq_len,
|
271
|
+
model_config.hidden_size // model_config.num_attention_heads,
|
272
|
+
],
|
273
|
+
"float32",
|
274
|
+
)
|
275
|
+
for i in range(model_config.num_hidden_layers * 2)
|
276
|
+
]
|
277
|
+
)
|
278
|
+
return input_info
|
279
|
+
|
280
|
+
prefill_input_info = get_input_info(query_length=prefill_chunk_size)
|
281
|
+
dec_input_info = get_input_info(query_length=1)
|
282
|
+
|
283
|
+
prefill_rbln_runtime_config = RBLNRuntimeConfig(input_info=prefill_input_info)
|
284
|
+
dec_rbln_runtime_config = RBLNRuntimeConfig(input_info=dec_input_info)
|
285
|
+
|
286
|
+
dec_rbln_runtime_config.batch_size = rbln_batch_size
|
287
|
+
|
288
|
+
rbln_config = RBLNConfig.from_rbln_runtime_configs(
|
289
|
+
[prefill_rbln_runtime_config, dec_rbln_runtime_config],
|
290
|
+
_rbln_meta=meta,
|
291
|
+
)
|
292
|
+
|
293
|
+
return rbln_config
|
294
|
+
|
295
|
+
def _create_runtimes(self, rbln_device_map: Dict[str, int]) -> List[rebel.Runtime]:
|
296
|
+
device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
|
297
|
+
return [
|
298
|
+
self.compiled_models[0].create_runtime(input_info_index=0, tensor_type="pt", device=device_val),
|
299
|
+
self.compiled_models[0].create_runtime(input_info_index=1, tensor_type="pt", device=device_val),
|
300
|
+
]
|
301
|
+
|
302
|
+
def get_decoder(self):
|
303
|
+
return self.decoder
|
304
|
+
|
305
|
+
def can_generate(self):
|
306
|
+
return True
|
307
|
+
|
308
|
+
def __getattr__(self, __name: str) -> Any:
|
309
|
+
def redirect(func):
|
310
|
+
return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
|
311
|
+
|
312
|
+
val = getattr(LlamaForCausalLM, __name)
|
313
|
+
|
314
|
+
if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
|
315
|
+
return redirect(val)
|
316
|
+
|
317
|
+
return val
|
318
|
+
|
319
|
+
def _reorder_cache(self, past_key_values, beam_idx):
|
320
|
+
raise NotImplementedError
|
321
|
+
|
322
|
+
# args input_ids, past_key_values and attention_mask are updated by _update_model_kwargs_for_generation() in _greedy_search() in GenerationMixin
|
323
|
+
def prepare_inputs_for_generation(self, input_ids, past_key_values=0, attention_mask=None, **kwargs):
|
324
|
+
batch_size, hf_input_length = input_ids.shape
|
325
|
+
past_cached_length = past_key_values
|
326
|
+
query_length = hf_input_length - past_cached_length
|
327
|
+
|
328
|
+
# In greedy decoding
|
329
|
+
if past_key_values == 0:
|
330
|
+
self.prompt_length = query_length
|
331
|
+
self.prompt_ids = input_ids
|
332
|
+
self.prompt_attn_mask = attention_mask.unsqueeze(1).unsqueeze(1).contiguous()
|
333
|
+
|
334
|
+
attention_mask = torch.zeros(batch_size, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64)
|
335
|
+
cache_position = torch.tensor(0, dtype=torch.int32)
|
336
|
+
else:
|
337
|
+
attention_mask = torch.nn.functional.pad(attention_mask, (0, self.max_seq_len - hf_input_length))
|
338
|
+
attention_mask = attention_mask.reshape(batch_size, 1, 1, -1).contiguous()
|
339
|
+
cache_position = torch.tensor(past_cached_length, dtype=torch.int32)
|
340
|
+
input_ids = input_ids[:, -1:]
|
341
|
+
|
342
|
+
model_inputs = {
|
343
|
+
"input_ids": input_ids,
|
344
|
+
"past_key_values": past_key_values,
|
345
|
+
"attention_mask": attention_mask,
|
346
|
+
"cache_position": cache_position,
|
347
|
+
"query_length": query_length,
|
348
|
+
}
|
349
|
+
|
350
|
+
return model_inputs
|
351
|
+
|
352
|
+
def forward(
|
353
|
+
self,
|
354
|
+
input_ids: torch.LongTensor = None,
|
355
|
+
attention_mask: Optional[torch.Tensor] = None,
|
356
|
+
past_key_values: int = None,
|
357
|
+
cache_position: Optional[torch.Tensor] = None,
|
358
|
+
query_length: Optional[torch.Tensor] = None,
|
359
|
+
**kwargs,
|
360
|
+
) -> Tuple[torch.FloatTensor]:
|
361
|
+
if past_key_values is not None:
|
362
|
+
past_key_values += query_length
|
363
|
+
|
364
|
+
# prefill_decoder
|
365
|
+
if cache_position == 0:
|
366
|
+
while query_length > self.prefill_chunk_size:
|
367
|
+
# prepare input_ids & attention_mask
|
368
|
+
sliced_input_ids = input_ids[:, cache_position : cache_position + self.prefill_chunk_size].contiguous()
|
369
|
+
attention_mask[:, :, :, :cache_position] = 1
|
370
|
+
attention_mask[:, :, :, cache_position : cache_position + self.prefill_chunk_size] = self.causal_mask
|
371
|
+
attention_mask[:, :, :, : self.prompt_length] *= self.prompt_attn_mask[:, :, :, :]
|
372
|
+
|
373
|
+
_ = self.prefill_decoder(
|
374
|
+
sliced_input_ids,
|
375
|
+
attention_mask,
|
376
|
+
cache_position,
|
377
|
+
)
|
378
|
+
# update query_length & cache_position
|
379
|
+
query_length -= self.prefill_chunk_size
|
380
|
+
cache_position += self.prefill_chunk_size
|
381
|
+
|
382
|
+
# prepare input_ids & attention_mask
|
383
|
+
last_input_ids = input_ids[:, cache_position : cache_position + query_length]
|
384
|
+
last_input_ids = torch.nn.functional.pad(last_input_ids, (0, self.prefill_chunk_size - query_length))
|
385
|
+
|
386
|
+
attention_mask[:, :, :, :cache_position] = 1
|
387
|
+
mask_slice = self.causal_mask[:, :, :query_length, :query_length]
|
388
|
+
attention_mask[:, :, :query_length, cache_position : cache_position + query_length] = mask_slice
|
389
|
+
attention_mask[:, :, :, : self.prompt_length] *= self.prompt_attn_mask[:, :, :, :]
|
390
|
+
|
391
|
+
outputs = self.prefill_decoder(
|
392
|
+
last_input_ids.contiguous(),
|
393
|
+
attention_mask.contiguous(),
|
394
|
+
cache_position,
|
395
|
+
)
|
396
|
+
|
397
|
+
outputs = outputs[:, query_length - 1].unsqueeze(1)
|
398
|
+
# decoder
|
399
|
+
else:
|
400
|
+
outputs = self.decoder(
|
401
|
+
input_ids.contiguous(),
|
402
|
+
attention_mask.contiguous(),
|
403
|
+
cache_position=cache_position,
|
404
|
+
)
|
405
|
+
|
406
|
+
return CausalLMOutputWithPast(
|
407
|
+
logits=outputs,
|
408
|
+
past_key_values=past_key_values,
|
409
|
+
)
|
@@ -0,0 +1,24 @@
|
|
1
|
+
# Copyright 2024 Rebellions Inc.
|
2
|
+
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at:
|
6
|
+
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
+
# See the License for the specific language governing permissions and
|
13
|
+
# limitations under the License.
|
14
|
+
|
15
|
+
# Portions of this software are licensed under the Apache License,
|
16
|
+
# Version 2.0. See the NOTICE file distributed with this work for
|
17
|
+
# additional information regarding copyright ownership.
|
18
|
+
|
19
|
+
# All other portions of this software, including proprietary code,
|
20
|
+
# are the intellectual property of Rebellions Inc. and may not be
|
21
|
+
# copied, modified, or distributed without prior written permission
|
22
|
+
# from Rebellions Inc.
|
23
|
+
|
24
|
+
from .t5_architecture import T5DecoderWrapper, T5EncoderWrapper
|