DoruC/Grounded-Segment-Anything
0
1# coding=utf-82# Copyright 2018 Salesforce and HuggingFace Inc. team.3# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16""" PyTorch CTRL model."""17 18from typing import Optional, Tuple, Union19 20import numpy as np21import torch22from torch import nn23from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss24 25from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutput26from ...modeling_utils import PreTrainedModel27from ...pytorch_utils import Conv1D, find_pruneable_heads_and_indices, prune_linear_layer28from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings29from .configuration_ctrl import CTRLConfig30 31 32logger = logging.get_logger(__name__)33 34_CONFIG_FOR_DOC = "CTRLConfig"35 36CTRL_PRETRAINED_MODEL_ARCHIVE_LIST = [37 "Salesforce/ctrl"38 # See all CTRL models at https://huggingface.co/models?filter=ctrl39]40 41 42def angle_defn(pos, i, d_model_size):43 angle_rates = 1 / torch.pow(10000, (2 * (i // 2)) / d_model_size)44 return pos * angle_rates45 46 47def positional_encoding(position, d_model_size, dtype):48 # create the sinusoidal pattern for the positional encoding49 angle_rads = angle_defn(50 torch.arange(position, dtype=dtype).unsqueeze(1),51 torch.arange(d_model_size, dtype=dtype).unsqueeze(0),52 d_model_size,53 )54 55 sines = torch.sin(angle_rads[:, 0::2])56 cosines = torch.cos(angle_rads[:, 1::2])57 58 pos_encoding = torch.cat([sines, cosines], dim=-1)59 return pos_encoding60 61 62def scaled_dot_product_attention(q, k, v, mask, attention_mask=None, head_mask=None):63 # calculate attention64 matmul_qk = torch.matmul(q, k.permute(0, 1, 3, 2))65 66 dk = k.shape[-1]67 scaled_attention_logits = matmul_qk / np.sqrt(dk)68 69 if mask is not None:70 nd, ns = scaled_attention_logits.size(-2), scaled_attention_logits.size(-1)71 scaled_attention_logits += mask[ns - nd : ns, :ns] * -1e472 73 if attention_mask is not None:74 # Apply the attention mask75 scaled_attention_logits = scaled_attention_logits + attention_mask76 77 attention_weights = torch.softmax(scaled_attention_logits, dim=-1)78 79 # Mask heads if we want to80 if head_mask is not None:81 attention_weights = attention_weights * head_mask82 83 output = torch.matmul(attention_weights, v)84 85 return output, attention_weights86 87 88class MultiHeadAttention(nn.Module):89 def __init__(self, d_model_size, num_heads):90 super().__init__()91 self.num_heads = num_heads92 self.d_model_size = d_model_size93 94 self.depth = int(d_model_size / self.num_heads)95 96 self.Wq = nn.Linear(d_model_size, d_model_size)97 self.Wk = nn.Linear(d_model_size, d_model_size)98 self.Wv = nn.Linear(d_model_size, d_model_size)99 100 self.dense = nn.Linear(d_model_size, d_model_size)101 self.pruned_heads = set()102 103 def prune_heads(self, heads):104 attention_head_size = self.d_model_size // self.num_heads105 if len(heads) == 0:106 return107 heads, index = find_pruneable_heads_and_indices(heads, self.num_heads, attention_head_size, self.pruned_heads)108 109 # Prune linear layers110 self.Wq = prune_linear_layer(self.Wq, index)111 self.Wk = prune_linear_layer(self.Wk, index)112 self.Wv = prune_linear_layer(self.Wv, index)113 self.dense = prune_linear_layer(self.dense, index, dim=1)114 115 # Update hyper params116 self.num_heads = self.num_heads - len(heads)117 self.d_model_size = attention_head_size * self.num_heads118 self.pruned_heads = self.pruned_heads.union(heads)119 120 def split_into_heads(self, x, batch_size):121 x = x.reshape(batch_size, -1, self.num_heads, self.depth)122 return x.permute([0, 2, 1, 3])123 124 def forward(125 self,126 v,127 k,128 q,129 mask,130 layer_past=None,131 attention_mask=None,132 head_mask=None,133 use_cache=False,134 output_attentions=False,135 ):136 batch_size = q.shape[0]137 138 q = self.Wq(q)139 k = self.Wk(k)140 v = self.Wv(v)141 142 q = self.split_into_heads(q, batch_size)143 k = self.split_into_heads(k, batch_size)144 v = self.split_into_heads(v, batch_size)145 if layer_past is not None:146 past_key, past_value = layer_past[0], layer_past[1]147 k = torch.cat((past_key, k), dim=-2)148 v = torch.cat((past_value, v), dim=-2)149 150 if use_cache is True:151 present = torch.stack((k, v))152 else:153 present = (None,)154 155 output = scaled_dot_product_attention(q, k, v, mask, attention_mask, head_mask)156 scaled_attention = output[0].permute([0, 2, 1, 3])157 attn = output[1]158 original_size_attention = scaled_attention.reshape(batch_size, -1, self.d_model_size)159 output = self.dense(original_size_attention)160 161 outputs = (output, present)162 if output_attentions:163 outputs = outputs + (attn,)164 return outputs165 166 167def point_wise_feed_forward_network(d_model_size, dff):168 return nn.Sequential(nn.Linear(d_model_size, dff), nn.ReLU(), nn.Linear(dff, d_model_size))169 170 171class EncoderLayer(nn.Module):172 def __init__(self, d_model_size, num_heads, dff, rate=0.1):173 super().__init__()174 175 self.multi_head_attention = MultiHeadAttention(d_model_size, num_heads)176 self.ffn = point_wise_feed_forward_network(d_model_size, dff)177 178 self.layernorm1 = nn.LayerNorm(d_model_size, eps=1e-6)179 self.layernorm2 = nn.LayerNorm(d_model_size, eps=1e-6)180 181 self.dropout1 = nn.Dropout(rate)182 self.dropout2 = nn.Dropout(rate)183 184 def forward(185 self, x, mask, layer_past=None, attention_mask=None, head_mask=None, use_cache=False, output_attentions=False186 ):187 normed = self.layernorm1(x)188 attn_outputs = self.multi_head_attention(189 normed,190 normed,191 normed,192 mask,193 layer_past=layer_past,194 attention_mask=attention_mask,195 head_mask=head_mask,196 use_cache=use_cache,197 output_attentions=output_attentions,198 )199 attn_output = attn_outputs[0]200 attn_output = self.dropout1(attn_output)201 out1 = x + attn_output202 203 out2 = self.layernorm2(out1)204 ffn_output = self.ffn(out2)205 ffn_output = self.dropout2(ffn_output)206 out2 = out1 + ffn_output207 208 outputs = (out2,) + attn_outputs[1:]209 return outputs210 211 212class CTRLPreTrainedModel(PreTrainedModel):213 """214 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained215 models.216 """217 218 config_class = CTRLConfig219 base_model_prefix = "transformer"220 221 def _init_weights(self, module):222 """Initialize the weights."""223 if isinstance(module, (nn.Linear, Conv1D)):224 # Slightly different from the TF version which uses truncated_normal for initialization225 # cf https://github.com/pytorch/pytorch/pull/5617226 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)227 if module.bias is not None:228 module.bias.data.zero_()229 elif isinstance(module, nn.Embedding):230 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)231 if module.padding_idx is not None:232 module.weight.data[module.padding_idx].zero_()233 elif isinstance(module, nn.LayerNorm):234 module.bias.data.zero_()235 module.weight.data.fill_(1.0)236 237 238CTRL_START_DOCSTRING = r"""239 240 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the241 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads242 etc.)243 244 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.245 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage246 and behavior.247 248 Parameters:249 config ([`CTRLConfig`]): Model configuration class with all the parameters of the model.250 Initializing with a config file does not load the weights associated with the model, only the251 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.252"""253 254CTRL_INPUTS_DOCSTRING = r"""255 Args:256 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):257 `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values[0].shape[-2]`258 (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.259 260 If `past_key_values` is used, only input IDs that do not have their past calculated should be passed as261 `input_ids`.262 263 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and264 [`PreTrainedTokenizer.encode`] for details.265 266 [What are input IDs?](../glossary#input-ids)267 past_key_values (`Tuple[Tuple[torch.FloatTensor]]` of length `config.n_layers`):268 Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model (see269 `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have270 their past given to this model should not be passed as input ids as they have already been computed.271 attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):272 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:273 274 - 1 for tokens that are **not masked**,275 - 0 for tokens that are **masked**.276 277 [What are attention masks?](../glossary#attention-mask)278 token_type_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):279 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,280 1]`:281 282 - 0 corresponds to a *sentence A* token,283 - 1 corresponds to a *sentence B* token.284 285 [What are token type IDs?](../glossary#token-type-ids)286 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):287 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,288 config.max_position_embeddings - 1]`.289 290 [What are position IDs?](../glossary#position-ids)291 head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):292 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:293 294 - 1 indicates the head is **not masked**,295 - 0 indicates the head is **masked**.296 297 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):298 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This299 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the300 model's internal embedding lookup matrix.301 use_cache (`bool`, *optional*):302 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see303 `past_key_values`).304 output_attentions (`bool`, *optional*):305 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned306 tensors for more detail.307 output_hidden_states (`bool`, *optional*):308 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for309 more detail.310 return_dict (`bool`, *optional*):311 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.312"""313 314 315@add_start_docstrings(316 "The bare CTRL Model transformer outputting raw hidden-states without any specific head on top.",317 CTRL_START_DOCSTRING,318)319class CTRLModel(CTRLPreTrainedModel):320 def __init__(self, config):321 super().__init__(config)322 323 self.d_model_size = config.n_embd324 self.num_layers = config.n_layer325 326 self.pos_encoding = positional_encoding(config.n_positions, self.d_model_size, torch.float)327 328 self.w = nn.Embedding(config.vocab_size, config.n_embd)329 330 self.dropout = nn.Dropout(config.embd_pdrop)331 self.h = nn.ModuleList(332 [EncoderLayer(config.n_embd, config.n_head, config.dff, config.resid_pdrop) for _ in range(config.n_layer)]333 )334 self.layernorm = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)335 336 # Initialize weights and apply final processing337 self.post_init()338 339 def get_input_embeddings(self):340 return self.w341 342 def set_input_embeddings(self, new_embeddings):343 self.w = new_embeddings344 345 def _prune_heads(self, heads_to_prune):346 """347 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}348 """349 for layer, heads in heads_to_prune.items():350 self.h[layer].multi_head_attention.prune_heads(heads)351 352 @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)353 @replace_return_docstrings(output_type=BaseModelOutputWithPast, config_class=_CONFIG_FOR_DOC)354 def forward(355 self,356 input_ids: Optional[torch.LongTensor] = None,357 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,358 attention_mask: Optional[torch.FloatTensor] = None,359 token_type_ids: Optional[torch.LongTensor] = None,360 position_ids: Optional[torch.LongTensor] = None,361 head_mask: Optional[torch.FloatTensor] = None,362 inputs_embeds: Optional[torch.FloatTensor] = None,363 use_cache: Optional[bool] = None,364 output_attentions: Optional[bool] = None,365 output_hidden_states: Optional[bool] = None,366 return_dict: Optional[bool] = None,367 ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPast]:368 r"""369 Returns:370 371 Example:372 373 ```python374 >>> from transformers import AutoTokenizer, CTRLModel375 >>> import torch376 377 >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")378 >>> model = CTRLModel.from_pretrained("Salesforce/ctrl")379 380 >>> # CTRL was trained with control codes as the first token381 >>> inputs = tokenizer("Opinion My dog is cute", return_tensors="pt")382 >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()383 384 >>> outputs = model(**inputs)385 386 >>> last_hidden_states = outputs.last_hidden_state387 >>> list(last_hidden_states.shape)388 [1, 5, 1280]389 ```"""390 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions391 use_cache = use_cache if use_cache is not None else self.config.use_cache392 output_hidden_states = (393 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states394 )395 return_dict = return_dict if return_dict is not None else self.config.use_return_dict396 397 if input_ids is not None and inputs_embeds is not None:398 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")399 elif input_ids is not None:400 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)401 input_shape = input_ids.size()402 input_ids = input_ids.view(-1, input_shape[-1])403 batch_size = input_ids.shape[0]404 elif inputs_embeds is not None:405 input_shape = inputs_embeds.size()[:-1]406 batch_size = inputs_embeds.shape[0]407 else:408 raise ValueError("You have to specify either input_ids or inputs_embeds")409 410 device = input_ids.device if input_ids is not None else inputs_embeds.device411 412 if past_key_values is None:413 past_length = 0414 past_key_values = tuple([None] * len(self.h))415 else:416 past_length = past_key_values[0][0].size(-2)417 if position_ids is None:418 position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)419 position_ids = position_ids.unsqueeze(0)420 421 # Attention mask.422 if attention_mask is not None:423 if batch_size <= 0:424 raise ValueError("batch_size has to be defined and > 0")425 attention_mask = attention_mask.view(batch_size, -1)426 # We create a 3D attention mask from a 2D tensor mask.427 # Sizes are [batch_size, 1, 1, to_seq_length]428 # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]429 # this attention mask is more simple than the triangular masking of causal attention430 # used in OpenAI GPT, we just need to prepare the broadcast dimension here.431 attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)432 433 # Since attention_mask is 1.0 for positions we want to attend and 0.0 for434 # masked positions, this operation will create a tensor which is 0.0 for435 # positions we want to attend and the dtype's smallest value for masked positions.436 # Since we are adding it to the raw scores before the softmax, this is437 # effectively the same as removing these entirely.438 attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility439 attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min440 441 # Prepare head mask if needed442 head_mask = self.get_head_mask(head_mask, self.config.n_layer)443 444 if token_type_ids is not None:445 token_type_ids = token_type_ids.view(-1, input_shape[-1])446 token_type_embeds = self.w(token_type_ids)447 token_type_embeds *= np.sqrt(self.d_model_size)448 else:449 token_type_embeds = 0450 451 if inputs_embeds is None:452 inputs_embeds = self.w(input_ids)453 # inputs_embeds = embedded.unsqueeze(0) if len(input_ids.shape)<2 else embedded454 seq_len = input_shape[-1]455 mask = torch.triu(torch.ones(seq_len + past_length, seq_len + past_length), 1).to(device)456 457 inputs_embeds *= np.sqrt(self.d_model_size)458 459 # `self.pos_encoding` won't be sent to the correct device along the model, so we do it manually.460 self.pos_encoding = self.pos_encoding.to(device)461 pos_embeds = self.pos_encoding[position_ids, :]462 463 hidden_states = inputs_embeds + pos_embeds + token_type_embeds464 465 hidden_states = self.dropout(hidden_states)466 467 presents = () if use_cache else None468 all_hidden_states = () if output_hidden_states else None469 all_attentions = () if output_attentions else None470 for i, (h, layer_past) in enumerate(zip(self.h, past_key_values)):471 if output_hidden_states:472 all_hidden_states = all_hidden_states + (hidden_states,)473 outputs = h(474 hidden_states,475 mask,476 layer_past=layer_past,477 attention_mask=attention_mask,478 head_mask=head_mask[i],479 use_cache=use_cache,480 output_attentions=output_attentions,481 )482 hidden_states, present = outputs[:2]483 if use_cache is True:484 presents = presents + (present,)485 486 if output_attentions:487 all_attentions += (outputs[2],)488 489 hidden_states = self.layernorm(hidden_states)490 if output_hidden_states:491 all_hidden_states = all_hidden_states + (hidden_states,)492 493 if not return_dict:494 return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None)495 496 return BaseModelOutputWithPast(497 last_hidden_state=hidden_states,498 past_key_values=presents,499 hidden_states=all_hidden_states,500 attentions=all_attentions,501 )502 503 504@add_start_docstrings(505 """506 The CTRL Model transformer with a language modeling head on top (linear layer with weights tied to the input507 embeddings).508 """,509 CTRL_START_DOCSTRING,510)511class CTRLLMHeadModel(CTRLPreTrainedModel):512 _tied_weights_keys = ["lm_head.weight"]513 514 def __init__(self, config):515 super().__init__(config)516 self.transformer = CTRLModel(config)517 self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=True)518 519 # Initialize weights and apply final processing520 self.post_init()521 522 def get_output_embeddings(self):523 return self.lm_head524 525 def set_output_embeddings(self, new_embeddings):526 self.lm_head = new_embeddings527 528 def prepare_inputs_for_generation(self, input_ids, past_key_values=None, use_cache=None, **kwargs):529 # only last token for inputs_ids if past is defined in kwargs530 if past_key_values:531 input_ids = input_ids[:, -1].unsqueeze(-1)532 533 return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": use_cache}534 535 @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)536 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)537 def forward(538 self,539 input_ids: Optional[torch.LongTensor] = None,540 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,541 attention_mask: Optional[torch.FloatTensor] = None,542 token_type_ids: Optional[torch.LongTensor] = None,543 position_ids: Optional[torch.LongTensor] = None,544 head_mask: Optional[torch.FloatTensor] = None,545 inputs_embeds: Optional[torch.FloatTensor] = None,546 labels: Optional[torch.LongTensor] = None,547 use_cache: Optional[bool] = None,548 output_attentions: Optional[bool] = None,549 output_hidden_states: Optional[bool] = None,550 return_dict: Optional[bool] = None,551 ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithPast]:552 r"""553 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):554 Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set555 `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`556 are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`557 558 Returns:559 560 Example:561 562 ```python563 >>> import torch564 >>> from transformers import AutoTokenizer, CTRLLMHeadModel565 566 >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")567 >>> model = CTRLLMHeadModel.from_pretrained("Salesforce/ctrl")568 569 >>> # CTRL was trained with control codes as the first token570 >>> inputs = tokenizer("Wikipedia The llama is", return_tensors="pt")571 >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()572 573 >>> sequence_ids = model.generate(inputs["input_ids"])574 >>> sequences = tokenizer.batch_decode(sequence_ids)575 >>> sequences576 ['Wikipedia The llama is a member of the family Bovidae. It is native to the Andes of Peru,']577 578 >>> outputs = model(**inputs, labels=inputs["input_ids"])579 >>> round(outputs.loss.item(), 2)580 9.21581 582 >>> list(outputs.logits.shape)583 [1, 5, 246534]584 ```"""585 return_dict = return_dict if return_dict is not None else self.config.use_return_dict586 587 transformer_outputs = self.transformer(588 input_ids,589 past_key_values=past_key_values,590 attention_mask=attention_mask,591 token_type_ids=token_type_ids,592 position_ids=position_ids,593 head_mask=head_mask,594 inputs_embeds=inputs_embeds,595 use_cache=use_cache,596 output_attentions=output_attentions,597 output_hidden_states=output_hidden_states,598 return_dict=return_dict,599 )600 601 hidden_states = transformer_outputs[0]602 603 lm_logits = self.lm_head(hidden_states)604 605 loss = None606 if labels is not None:607 # Shift so that tokens < n predict n608 shift_logits = lm_logits[..., :-1, :].contiguous()609 shift_labels = labels[..., 1:].contiguous()610 # Flatten the tokens611 loss_fct = CrossEntropyLoss()612 loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))613 614 if not return_dict:615 output = (lm_logits,) + transformer_outputs[1:]616 return ((loss,) + output) if loss is not None else output617 618 return CausalLMOutputWithPast(619 loss=loss,620 logits=lm_logits,621 past_key_values=transformer_outputs.past_key_values,622 hidden_states=transformer_outputs.hidden_states,623 attentions=transformer_outputs.attentions,624 )625 626 @staticmethod627 def _reorder_cache(628 past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor629 ) -> Tuple[Tuple[torch.Tensor]]:630 """631 This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or632 [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct633 beam_idx at every generation step.634 """635 return tuple(636 tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)637 for layer_past in past_key_values638 )639 640 641@add_start_docstrings(642 """643 The CTRL Model transformer with a sequence classification head on top (linear layer).644 [`CTRLForSequenceClassification`] uses the last token in order to do the classification, as other causal models645 (e.g. GPT-2) do. Since it does classification on the last token, it requires to know the position of the last646 token. If a `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in647 each row. If no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot648 guess the padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last649 value in each row of the batch).650 """,651 CTRL_START_DOCSTRING,652)653class CTRLForSequenceClassification(CTRLPreTrainedModel):654 def __init__(self, config):655 super().__init__(config)656 self.num_labels = config.num_labels657 self.transformer = CTRLModel(config)658 self.classifier = nn.Linear(config.n_embd, self.num_labels, bias=False)659 660 # Initialize weights and apply final processing661 self.post_init()662 663 @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)664 @replace_return_docstrings(output_type=SequenceClassifierOutput, config_class=_CONFIG_FOR_DOC)665 def forward(666 self,667 input_ids: Optional[torch.LongTensor] = None,668 past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,669 attention_mask: Optional[torch.FloatTensor] = None,670 token_type_ids: Optional[torch.LongTensor] = None,671 position_ids: Optional[torch.LongTensor] = None,672 head_mask: Optional[torch.FloatTensor] = None,673 inputs_embeds: Optional[torch.FloatTensor] = None,674 labels: Optional[torch.LongTensor] = None,675 use_cache: Optional[bool] = None,676 output_attentions: Optional[bool] = None,677 output_hidden_states: Optional[bool] = None,678 return_dict: Optional[bool] = None,679 ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:680 r"""681 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):682 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,683 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If684 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).685 686 Returns:687 688 Example of single-label classification:689 690 ```python691 >>> import torch692 >>> from transformers import AutoTokenizer, CTRLForSequenceClassification693 694 >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")695 >>> model = CTRLForSequenceClassification.from_pretrained("Salesforce/ctrl")696 697 >>> # CTRL was trained with control codes as the first token698 >>> inputs = tokenizer("Opinion My dog is cute", return_tensors="pt")699 >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()700 701 >>> with torch.no_grad():702 ... logits = model(**inputs).logits703 704 >>> predicted_class_id = logits.argmax().item()705 >>> model.config.id2label[predicted_class_id]706 'LABEL_0'707 ```708 709 ```python710 >>> import torch711 712 >>> torch.manual_seed(42) # doctest: +IGNORE_RESULT713 >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`714 >>> num_labels = len(model.config.id2label)715 >>> model = CTRLForSequenceClassification.from_pretrained("Salesforce/ctrl", num_labels=num_labels)716 717 >>> labels = torch.tensor(1)718 >>> loss = model(**inputs, labels=labels).loss719 >>> round(loss.item(), 2)720 0.35721 ```722 723 Example of multi-label classification:724 725 ```python726 >>> import torch727 >>> from transformers import AutoTokenizer, CTRLForSequenceClassification728 729 >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")730 >>> model = CTRLForSequenceClassification.from_pretrained(731 ... "Salesforce/ctrl", problem_type="multi_label_classification"732 ... )733 734 >>> # CTRL was trained with control codes as the first token735 >>> inputs = tokenizer("Opinion My dog is cute", return_tensors="pt")736 >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()737 738 >>> with torch.no_grad():739 ... logits = model(**inputs).logits740 741 >>> predicted_class_id = logits.argmax().item()742 >>> model.config.id2label[predicted_class_id]743 'LABEL_0'744 ```745 746 ```python747 >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`748 >>> num_labels = len(model.config.id2label)749 >>> model = CTRLForSequenceClassification.from_pretrained("Salesforce/ctrl", num_labels=num_labels)750 751 >>> num_labels = len(model.config.id2label)752 >>> labels = torch.nn.functional.one_hot(torch.tensor([predicted_class_id]), num_classes=num_labels).to(753 ... torch.float754 ... )755 >>> loss = model(**inputs, labels=labels).loss756 >>> loss.backward() # doctest: +IGNORE_RESULT757 ```"""758 759 return_dict = return_dict if return_dict is not None else self.config.use_return_dict760 761 transformer_outputs = self.transformer(762 input_ids,763 past_key_values=past_key_values,764 attention_mask=attention_mask,765 token_type_ids=token_type_ids,766 position_ids=position_ids,767 head_mask=head_mask,768 inputs_embeds=inputs_embeds,769 use_cache=use_cache,770 output_attentions=output_attentions,771 output_hidden_states=output_hidden_states,772 return_dict=return_dict,773 )774 775 hidden_states = transformer_outputs[0]776 logits = self.classifier(hidden_states)777 778 if input_ids is not None:779 batch_size, sequence_length = input_ids.shape[:2]780 else:781 batch_size, sequence_length = inputs_embeds.shape[:2]782 783 if self.config.pad_token_id is None and batch_size != 1:784 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")785 786 if self.config.pad_token_id is None:787 sequence_lengths = -1788 else:789 if input_ids is not None:790 sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(791 logits.device792 )793 else:794 sequence_lengths = -1795 logger.warning(796 f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "797 "unexpected if using padding tokens in conjunction with `inputs_embeds.`"798 )799 800 pooled_logits = logits[range(batch_size), sequence_lengths]801 802 loss = None803 if labels is not None:804 if self.config.problem_type is None:805 if self.num_labels == 1:806 self.config.problem_type = "regression"807 elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):808 self.config.problem_type = "single_label_classification"809 else:810 self.config.problem_type = "multi_label_classification"811 812 if self.config.problem_type == "regression":813 loss_fct = MSELoss()814 if self.num_labels == 1:815 loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())816 else:817 loss = loss_fct(pooled_logits, labels)818 elif self.config.problem_type == "single_label_classification":819 loss_fct = CrossEntropyLoss()820 loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))821 elif self.config.problem_type == "multi_label_classification":822 loss_fct = BCEWithLogitsLoss()823 loss = loss_fct(pooled_logits, labels)824 if not return_dict:825 output = (pooled_logits,) + transformer_outputs[2:]826 return ((loss,) + output) if loss is not None else output827 828 return SequenceClassifierOutput(829 loss=loss,830 logits=pooled_logits,831 hidden_states=transformer_outputs.hidden_states,832 attentions=transformer_outputs.attentions,833 )834