Fix batch beam search

modeling_glm.py

@@ -29,13 +29,13 @@ from transformers.utils import (
 )
 from transformers.modeling_outputs import (
     BaseModelOutputWithPastAndCrossAttentions,
+    ModelOutput,
     SequenceClassifierOutput,
-    ModelOutput
 )
+
 from transformers.modeling_utils import (
     PreTrainedModel,
 )
-from transformers.utils import logging
 from .configuration_glm import GLMConfig
 from torch.nn.parameter import Parameter
 
@@ -781,20 +781,60 @@ 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,
         )
 
 
+@add_start_docstrings(
+    """GLM Model transformer for multiple choice classification""",
+    GLM_START_DOCSTRING
+)
+class GLMForMultipleChoice(GLMPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.glm = GLMModel(config)
+        self.post_init()
+
+    def forward(
+            self,
+            input_ids=None,
+            position_ids=None,
+            attention_mask=None,
+            choice_ids=None,
+            choice_indices=None,
+            labels=None,
+            mems=None,
+            **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):
+            log_probs_single = []
+            for choice, choice_target_id in zip(choices, choice_index):
+                tmp = output[choice_target_id, choice]
+                log_probs_single.append(tmp.sum())
+            log_probs.append(torch.stack(log_probs_single))
+        log_probs = torch.stack(log_probs)
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(log_probs, labels)
+        return ModelOutput(
+            loss=loss,
+            logits=log_probs,
+            lm_logits=lm_logits,
+            mems=model_output.mems
+        )
+
 @add_start_docstrings(
     """GLM Model transformer with a `language modeling` head on top""",
     GLM_START_DOCSTRING,
@@ -833,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,
@@ -845,10 +895,21 @@ class GLMForConditionalGeneration(GLMPreTrainedModel):
             input_ids=None,
             position_ids=None,
             attention_mask=None,
+            labels=None,
             mems=None,
             **kwargs
     ):
-        return 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:
+            loss_fct = CrossEntropyLoss(ignore_index=-100)
+            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
+        return ModelOutput(
+            loss=loss,
+            logits=lm_logits,
+            mems=model_output.mems
+        )
 
 
 @add_start_docstrings(
@@ -857,16 +918,19 @@ class GLMForConditionalGeneration(GLMPreTrainedModel):
     GLM_START_DOCSTRING,
 )
 class GLMForSequenceClassification(GLMPreTrainedModel):
-    def __init__(self, config, hidden_dropout=False, num_class=1):
+    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.pool_layer = torch.nn.Linear(config.hidden_size, config.hidden_size)
-        self.multichoice_dropout = torch.nn.Dropout(hidden_dropout)
-        self.multichoice_head = torch.nn.Linear(config.hidden_size, num_class)
+        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()
@@ -891,29 +955,21 @@ class GLMForSequenceClassification(GLMPreTrainedModel):
             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.forward(input_ids, position_ids, attention_mask)
-        outputs, mems = model_out.last_hidden_state, model_out.hidden_states
-
-        if self.pool_token == 'start':
-            output = outputs[
-                torch.arange(outputs.size(0), dtype=attention_mask.dtype, device=attention_mask.device), attention_mask]
-        elif self.pool_token == 'pad':
-            output = outputs[torch.arange(outputs.size(0), dtype=attention_mask.dtype,
-                                          device=attention_mask.device), attention_mask - 1]
-        elif self.pool_token == 'cls':
-            output = outputs[:, 0]
-        else:
-            raise NotImplementedError
+        model_out = self.glm(input_ids, position_ids, attention_mask)
+        outputs, mems = model_out.last_hidden_states, model_out.mems
 
-        output = torch.tanh(self.pool_layer(output))
-        multichoice_output = self.multichoice_dropout(output)
-        logits = self.multichoice_head(multichoice_output)
-        loss_fct = CrossEntropyLoss()
+        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)
-        # assert (labels is not None, "labels must not None!")
-        loss = loss_fct(logits, labels)
+        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=mems)
+                                        hidden_states=outputs)