Upload modeling_internlm2.py with huggingface_hub

modeling_internlm2.py

@@ -59,10 +59,6 @@ try:
 except:
     pass
 
-try:
-    support_bf16_triu = torch.__version__ >= "2.1.0"
-except Exception:
-    support_bf16_triu = False
 
 logger = logging.get_logger(__name__)
 
@@ -1097,11 +1093,7 @@ class InternLM2Model(InternLM2PreTrainedModel):
         else:
             causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
             if sequence_length != 1:
-                if support_bf16_triu or dtype == torch.float32:
-                    causal_mask = torch.triu(causal_mask, diagonal=1)
-                else:
-                    triu_mask = torch.triu(torch.ones(causal_mask.size(), device=device), diagonal=1).bool()
-                    causal_mask.masked_fill_(~triu_mask, 0)
+                causal_mask = torch.triu(causal_mask, diagonal=1)
             causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
             causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
             if attention_mask is not None:
@@ -1806,3 +1798,238 @@ class InternLM2ForTokenClassification(InternLM2PreTrainedModel):
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
+
+
+# Modified from transformers.models.llama.modeling_llama.LlamaForTokenClassification
+class InternLM2ForRewardModel(InternLM2PreTrainedModel):
+
+    _auto_class = "AutoModel"
+    _tied_weights_keys = ["v_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = InternLM2Model(config)
+        self.vocab_size = config.vocab_size
+        self.v_head = nn.Linear(config.hidden_size, 1, bias=False)
+        self.reward_token_id = config.reward_token_id
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.tok_embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.tok_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.v_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.v_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=SequenceClassifierOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+        """
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss.
+        
+        Returns:
+
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        hidden_states = self.v_head(hidden_states)
+        # get end reward token's score
+        ends = attention_mask.cumsum(dim=1).argmax(dim=1).view(-1,1)
+
+        reward_scores = torch.gather(hidden_states.squeeze(-1), 1, ends)
+
+        loss = None
+
+        if not return_dict:
+            output = (reward_scores,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=reward_scores,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    @torch.no_grad()
+    def get_score(
+        self,
+        tokenizer,
+        conversation: List[dict],
+        **kwargs,
+    ):
+        """
+        Computes the reward score for a given conversation.
+
+        This function takes a conversation represented as a list of dictionaries, formats it into a string using the chat 
+        template from the tokenizer, and passes it through the model to compute the score. A special token representing 
+        the reward score is appended to the input sequence. The reward score is then extracted from the model's output.
+
+        Args:
+            tokenizer: The tokenizer to be used for formatting and tokenizing the conversation.
+            conversation (List[dict]): A list of dictionaries where each dictionary represents a message in the conversation.
+
+        Returns:
+            float: The computed reward score from the model.
+        """
+        conversation_str = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False)
+        input_ids = tokenizer.encode(conversation_str, return_tensors="pt", add_special_tokens=False)
+        # add reward score token at the end of the input_ids if it is not already there
+        if input_ids[0, -1] != self.reward_token_id:
+            input_ids = torch.cat([input_ids, torch.tensor([[self.reward_token_id]], dtype=torch.long)], dim=1)
+        attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+        
+        outputs = self.forward(input_ids=input_ids.to(self.device), attention_mask=attention_mask.to(self.device), **kwargs)
+        score = outputs[0].cpu().item()
+        return score
+    
+    @torch.no_grad()
+    def get_scores(
+        self,
+        tokenizer,
+        conversations: List[List[dict]],
+        **kwargs,
+    ):
+        """
+        Computes the reward scores for multiple conversations in a batched manner.
+
+        This function takes multiple conversations, each represented as a list of dictionaries, formats them into strings using the chat 
+        template from the tokenizer, and passes these formatted strings through the model to compute scores for each conversation. 
+        Each input sequence has a special token representing the reward score appended before passing to the model. 
+        The reward scores are then extracted from the model's output.
+
+        Args:
+            tokenizer: The tokenizer to be used for formatting and tokenizing  the conversation.
+            conversations (List[List[dict]]): A list of conversations, with each conversation represented as a list of dictionaries where each dictionary contains a message.
+
+        Returns:
+            List[float]: A list of computed reward scores for each conversation in the input batch.
+        """
+        conversation_strs = [tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False) for conversation in conversations]
+        batch_input_ids = []
+        attention_masks = []
+
+        for conversation_str in conversation_strs:
+            input_ids = tokenizer.encode(conversation_str, return_tensors="pt", add_special_tokens=False)
+            # add reward score token at the end of the input_ids if it is not already there
+            if input_ids[0, -1] != self.reward_token_id:
+                input_ids = torch.cat([input_ids, torch.tensor([[self.reward_token_id]], dtype=torch.long)], dim=1)
+            input_ids = input_ids.squeeze(0)
+            attention_mask = torch.ones(input_ids.shape, dtype=torch.bool)
+            batch_input_ids.append(input_ids)
+            attention_masks.append(attention_mask)
+
+        r_pad_batch_input_ids = torch.nn.utils.rnn.pad_sequence(batch_input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
+        r_pad_attention_masks = torch.nn.utils.rnn.pad_sequence(attention_masks, batch_first=True, padding_value=False)
+
+        outputs = self.forward(input_ids=r_pad_batch_input_ids.to(self.device), attention_mask=r_pad_attention_masks.to(self.device), **kwargs)
+        scores = outputs[0].squeeze().cpu().tolist()
+        return scores
+    
+    @torch.no_grad()
+    def compare(
+        self,
+        tokenizer,
+        conversation1: List[dict],
+        conversation2: List[dict],
+        return_logits: bool = False,
+        **kwargs,
+    ):
+        """
+        Compares the reward scores of two conversations and determines which conversation has a higher score.
+
+        This function computes reward scores for two given conversations using the `get_score` method and compares the scores to determine which conversation has a higher score.
+        The function can optionally return the actual scores (logits) along with the comparison result.
+
+        Parameters:
+            tokenizer: The tokenizer used for formatting and tokenizing the conversation.
+            conversation1 (List[dict]): The first conversation to compare, represented as a list of dictionaries where each dictionary contains a message.
+            conversation2 (List[dict]): The second conversation to compare, similarly represented.
+            return_logits (bool, optional): If True, the function returns both the comparison result and the actual scores of the two conversations. Defaults to False.
+
+        Returns:
+
+            bool: True if the score of the first conversation is greater than the second, otherwise False.
+            List[float] (optional): A list containing the scores of the first and second conversations respectively.
+
+        Note:
+        - This function is designed for inference, with `@torch.no_grad()` used to disable gradient calculations to optimize performance.
+        """
+        score1 = self.get_score(tokenizer, conversation1, **kwargs)
+        score2 = self.get_score(tokenizer, conversation2, **kwargs)
+        if return_logits:
+            return score1 > score2, [score1, score2]
+        else:
+            return score1 > score2
+        
+    @torch.no_grad()
+    def rank(
+        self,
+        tokenizer,
+        conversations: List[List[dict]],
+        return_logits: bool = False,
+        **kwargs,
+    ):
+        """
+        Ranks the conversations based on their scores.
+
+        Args:
+            tokenizer: The tokenizer to be used for formatting and tokenizing  the conversation.
+            conversations: A list of conversations, where each conversation is represented as a list of dictionaries. Each dictionary contains the necessary information for the conversation.
+            return_logits: If True, returns the conversation indices along with their logits. Defaults to False.
+
+        Returns:
+            list: A list of conversation rank indices based on their scores. Smaller index means higher score.
+            List[float] (optional): If return_logits is True, a list of conversation indices and their corresponding logits.
+
+        """
+        scores = self.get_scores(tokenizer, conversations, **kwargs)
+        if return_logits:
+            return sorted(range(len(scores)), key=lambda i: scores[i], reverse=True), scores
+        else:
+            return sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)