Fix batch beam search

modeling_glm.py

@@ -30,6 +30,7 @@ from transformers.utils import (
 from transformers.modeling_outputs import (
     BaseModelOutputWithPastAndCrossAttentions,
     ModelOutput,
+    SequenceClassifierOutput,
 )
 
 from transformers.modeling_utils import (
@@ -780,17 +781,15 @@ class GLMModel(GLMPreTrainedModel):
             attention_mask = torch.zeros(batch_size)
         # Transformer.
         transformer_output = self.transformer(embeddings, position_ids, attention_mask, mems)
-        logits, hidden_layers = transformer_output
-        # outputs = hidden_layers
+        last_hidden_states, mems = transformer_output
+        logits = None
         if self.output_predict:
-            # Parallel logits.
-            # logits_parallel = mpu.copy_to_model_parallel_region(
-            #     logits)
-            logits = F.linear(logits, self.word_embeddings.weight)
+            logits = F.linear(last_hidden_states, self.word_embeddings.weight)
 
         return ModelOutput(
+            last_hidden_states=last_hidden_states,
             logits=logits,
-            mems=hidden_layers,
+            mems=mems,
         )
 
 
@@ -815,7 +814,7 @@ class GLMForMultipleChoice(GLMPreTrainedModel):
             mems=None,
             **kwargs
     ):
-        model_output = self.glm.forward(input_ids, position_ids, attention_mask, mems=mems, **kwargs)
+        model_output = self.glm(input_ids, position_ids, attention_mask, mems=mems, **kwargs)
         lm_logits = model_output.logits
         log_probs = []
         for output, choices, choice_index in zip(F.log_softmax(lm_logits, dim=-1), choice_ids, choice_indices):
@@ -874,6 +873,16 @@ class GLMForConditionalGeneration(GLMPreTrainedModel):
                 position_ids = position_ids[:, :, :seq_length]
             if attention_mask is not None:
                 attention_mask = attention_mask[:, :, :seq_length, :seq_length]
+        if position_ids is not None and input_ids.size(0) > position_ids.size(0):
+            batch_size = position_ids.size(0)
+            num_beams = input_ids.size(0) // batch_size
+            position_ids = position_ids.unsqueeze(1).expand(-1, num_beams, -1, -1)
+            position_ids = position_ids.reshape(batch_size * num_beams, *position_ids.shape[-2:])
+        if attention_mask is not None and input_ids.size(0) > attention_mask.size(0):
+            batch_size = attention_mask.size(0)
+            num_beams = input_ids.size(0) // batch_size
+            attention_mask = attention_mask.unsqueeze(1).expand(-1, num_beams, -1, -1, -1)
+            attention_mask = attention_mask.reshape(batch_size * num_beams, *attention_mask.shape[-3:])
         return {
             "input_ids": input_ids,
             "position_ids": position_ids,
@@ -890,7 +899,7 @@ class GLMForConditionalGeneration(GLMPreTrainedModel):
             mems=None,
             **kwargs
     ):
-        model_output = self.glm.forward(input_ids, position_ids, attention_mask, mems=mems, **kwargs)
+        model_output = self.glm(input_ids, position_ids, attention_mask, mems=mems, **kwargs)
         lm_logits = model_output.logits
         loss = None
         if labels is not None:
@@ -900,4 +909,67 @@ class GLMForConditionalGeneration(GLMPreTrainedModel):
             loss=loss,
             logits=lm_logits,
             mems=model_output.mems
-        )
+        )
+
+
+@add_start_docstrings(
+    """GLM Model transformer with a sequence classification/regression head on top (a linear layer on top of
+    the pooled output) e.g. for GLUE tasks. """,
+    GLM_START_DOCSTRING,
+)
+class GLMForSequenceClassification(GLMPreTrainedModel):
+    def __init__(self, config: GLMConfig, hidden_dropout=None, num_class=1):
+        super().__init__(config)
+        self.pool_token = config.pool_token
+        self.glm = GLMModel(config)
+        self.glm.output_predict = False
+        self.num_class = num_class
+        # Multi-choice head.
+        self.dense = torch.nn.Linear(config.hidden_size, config.hidden_size)
+        classifier_dropout = (
+            config.classifier_dropout if config.classifier_dropout is not None else config.output_dropout_prob
+        )
+        self.dropout = torch.nn.Dropout(classifier_dropout)
+        self.out_proj = torch.nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(GLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_code_sample_docstrings(
+        processor_class=_TOKENIZER_FOR_DOC,
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=SequenceClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(self,
+                input_ids=None,
+                position_ids=None,
+                attention_mask=None,
+                labels=None):
+
+        num_choices = None
+
+        if len(input_ids.shape) == 3:
+            batch_size, num_choices = input_ids.shape[:2]
+            input_ids = input_ids.reshape(-1, input_ids.size(-1))
+            attention_mask = attention_mask.reshape(-1, *attention_mask.size()[2:])
+            position_ids = position_ids.reshape(-1, *position_ids.size()[2:])
+        model_out = self.glm(input_ids, position_ids, attention_mask)
+        outputs, mems = model_out.last_hidden_states, model_out.mems
+
+        output = outputs[:, 0, :]
+        output = self.dropout(output)
+        output = torch.tanh(self.dense(output))
+        output = self.dropout(output)
+        logits = self.out_proj(output)
+        if num_choices is not None:
+            logits = logits.view(-1, num_choices)
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+        # loss = F.cross_entropy(logits.contiguous().float(), labels.long())
+        return SequenceClassifierOutput(loss=loss,
+                                        logits=logits,
+                                        hidden_states=outputs)