Upload 8 files

config.json

@@ -1,9 +1,14 @@
 {
   "_name_or_path": "miniG",
-  "model_type": "chatglm",
+  "add_bias_linear": false,
+  "add_qkv_bias": true,
+  "apply_query_key_layer_scaling": true,
+  "apply_residual_connection_post_layernorm": false,
   "architectures": [
-    "ChatGLMModel"
+    "ChatGLMForConditionalGeneration"
   ],
+  "attention_dropout": 0.0,
+  "attention_softmax_in_fp32": true,
   "auto_map": {
     "AutoConfig": "configuration_chatglm.ChatGLMConfig",
     "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
@@ -11,35 +16,53 @@
     "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
     "AutoModelForSequenceClassification": "modeling_chatglm.ChatGLMForSequenceClassification"
   },
-  "add_bias_linear": false,
-  "add_qkv_bias": true,
-  "apply_query_key_layer_scaling": true,
-  "apply_residual_connection_post_layernorm": false,
-  "attention_dropout": 0.0,
-  "attention_softmax_in_fp32": true,
-  "attn_implementation": "sdpa",
   "bias_dropout_fusion": true,
+  "boi_token_id": 151339,
+  "classifier_dropout": null,
+  "eoi_token_id": 151340,
+  "eos_token_id": [
+    151329,
+    151336,
+    151338
+  ],
   "ffn_hidden_size": 13696,
   "fp32_residual_connection": false,
   "hidden_dropout": 0.0,
   "hidden_size": 4096,
   "kv_channels": 128,
   "layernorm_epsilon": 1.5625e-07,
+  "model_type": "chatglm",
   "multi_query_attention": true,
   "multi_query_group_num": 4,
   "num_attention_heads": 32,
   "num_hidden_layers": 40,
   "num_layers": 40,
-  "rope_ratio": 10000,
   "original_rope": true,
+  "pad_token_id": 151329,
   "padded_vocab_size": 151552,
   "post_layer_norm": true,
+  "pre_seq_len": null,
+  "prefix_projection": false,
   "rmsnorm": true,
+  "rope_ratio": 10000,
   "seq_length": 1048576,
-  "use_cache": true,
-  "torch_dtype": "bfloat16",
-  "transformers_version": "4.44.0",
   "tie_word_embeddings": false,
-  "eos_token_id": [151329, 151336, 151338],
-  "pad_token_id": 151329
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.43.1",
+  "use_cache": true,
+  "vision_config": {
+    "dropout_prob": 0.0,
+    "hidden_act": "gelu",
+    "hidden_size": 1792,
+    "image_size": 1120,
+    "in_channels": 3,
+    "intermediate_size": 15360,
+    "layer_norm_eps": 1e-06,
+    "num_heads": 16,
+    "num_hidden_layers": 63,
+    "num_positions": 6401,
+    "patch_size": 14,
+    "scaling_factor": 8
+  },
+  "vocab_size": 151552
 }

configuration_chatglm.py

@@ -29,6 +29,10 @@ class ChatGLMConfig(PretrainedConfig):
             apply_query_key_layer_scaling=True,
             attention_softmax_in_fp32=True,
             fp32_residual_connection=False,
+            pre_seq_len=None,
+            prefix_projection=False,
+            boi_token_id=None,
+            eoi_token_id=None,
             **kwargs
     ):
         self.num_layers = num_layers
@@ -55,4 +59,8 @@ class ChatGLMConfig(PretrainedConfig):
         self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
         self.attention_softmax_in_fp32 = attention_softmax_in_fp32
         self.fp32_residual_connection = fp32_residual_connection
+        self.pre_seq_len = pre_seq_len
+        self.prefix_projection = prefix_projection
+        self.boi_token_id = boi_token_id
+        self.eoi_token_id = eoi_token_id
         super().__init__(**kwargs)

generation_config.json

@@ -1,13 +1,13 @@
 {
+  "do_sample": true,
   "eos_token_id": [
     151329,
     151336,
     151338
   ],
+  "max_length": 8192,
   "pad_token_id": 151329,
-  "do_sample": true,
   "temperature": 0.8,
-  "max_length": 1024000,
   "top_p": 0.8,
-  "transformers_version": "4.44.0"
-}
+  "transformers_version": "4.43.1"
+}

modeling_chatglm.py

@@ -1,19 +1,13 @@
-""" PyTorch ChatGLM model. """
-import json
+""" PyTorch GLM-4V model. """
 import math
-import copy
-import warnings
-import re
 import sys
-
 import torch
 import torch.utils.checkpoint
 import torch.nn.functional as F
 from torch import nn
 from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss
 from torch.nn.utils import skip_init
-from typing import Optional, Tuple, Union, List, Callable, Dict, Any
-from copy import deepcopy
+from typing import Optional, Tuple, Union, List, Dict, Any
 
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
@@ -25,6 +19,7 @@ from transformers.utils import logging, is_torch_npu_available
 from transformers.generation.logits_process import LogitsProcessor
 from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
 
+from .visual import EVA2CLIPModel
 from .configuration_chatglm import ChatGLMConfig
 
 try:
@@ -46,6 +41,9 @@ if sys.platform != 'darwin' and not is_torch_npu_available():
 
 logger = logging.get_logger(__name__)
 
+LANGUAGE_TOKEN_TYPE = 0
+VISION_TOKEN_TYPE = 1
+
 _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"
 _CONFIG_FOR_DOC = "ChatGLMConfig"
 
@@ -62,6 +60,38 @@ class InvalidScoreLogitsProcessor(LogitsProcessor):
         return scores
 
 
+class PrefixEncoder(torch.nn.Module):
+    """
+    The torch.nn model to encode the prefix
+    Input shape: (batch-size, prefix-length)
+    Output shape: (batch-size, prefix-length, 2*layers*hidden)
+    """
+
+    def __init__(self, config: ChatGLMConfig):
+        super().__init__()
+        self.prefix_projection = config.prefix_projection
+        if self.prefix_projection:
+            # Use a two-layer MLP to encode the prefix
+            kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2
+            self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size)
+            self.trans = torch.nn.Sequential(
+                torch.nn.Linear(kv_size, config.hidden_size),
+                torch.nn.Tanh(),
+                torch.nn.Linear(config.hidden_size, kv_size)
+            )
+        else:
+            self.embedding = torch.nn.Embedding(config.pre_seq_len,
+                                                config.num_layers * config.kv_channels * config.multi_query_group_num * 2)
+
+    def forward(self, prefix: torch.Tensor):
+        if self.prefix_projection:
+            prefix_tokens = self.embedding(prefix)
+            past_key_values = self.trans(prefix_tokens)
+        else:
+            past_key_values = self.embedding(prefix)
+        return past_key_values
+
+
 def split_tensor_along_last_dim(
         tensor: torch.Tensor,
         num_partitions: int,
@@ -99,6 +129,17 @@ class RotaryEmbedding(nn.Module):
         self.original_impl = original_impl
         self.rope_ratio = rope_ratio
 
+    def impl(self, seq_length: int, dim: int, device: torch.device, dtype: torch.dtype):
+        base = 10000 * self.rope_ratio
+        inv_freq = 1.0 / (
+                base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
+        seq = torch.arange(seq_length, device=inv_freq.device, dtype=torch.float32)
+        freqs = torch.outer(seq, inv_freq)
+        # first part even vector components, second part odd vector components,
+        #  2 * dim in dimension size
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb
+
     def forward_impl(
             self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
     ):
@@ -126,9 +167,12 @@ class RotaryEmbedding(nn.Module):
         return cache
 
     def forward(self, max_seq_len, offset=0):
-        return self.forward_impl(
-            max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device
-        )
+        if self.original_impl:
+            return self.forward_impl(
+                max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device
+            )
+        else:
+            return self.impl(max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
 
 
 @torch.jit.script
@@ -166,16 +210,16 @@ class RMSNorm(torch.nn.Module):
         return (self.weight * hidden_states).to(input_dtype)
 
 
+
 class CoreAttention(torch.nn.Module):
     def __init__(self, config: ChatGLMConfig, layer_number):
         super(CoreAttention, self).__init__()
-        self.config = config
+
         self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling
         self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
         if self.apply_query_key_layer_scaling:
             self.attention_softmax_in_fp32 = True
         self.layer_number = max(1, layer_number)
-        self.is_causal = True
 
         projection_size = config.kv_channels * config.num_attention_heads
 
@@ -194,76 +238,94 @@ class CoreAttention(torch.nn.Module):
         self.attention_dropout = torch.nn.Dropout(config.attention_dropout)
 
     def forward(self, query_layer, key_layer, value_layer, attention_mask):
-        # [b, np, sq, sk]
-        output_size = (query_layer.size(0), query_layer.size(1), query_layer.size(2), key_layer.size(2))
-
-        # [b, np, sq, hn] -> [b * np, sq, hn]
-        query_layer = query_layer.view(output_size[0] * output_size[1], output_size[2], -1)
-        # [b, np, sk, hn] -> [b * np, sk, hn]
-        key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)
-
-        # preallocting input tensor: [b * np, sq, sk]
-        matmul_input_buffer = torch.empty(
-            output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,
-            device=query_layer.device
-        )
+        pytorch_major_version = int(torch.__version__.split('.')[0])
+        if pytorch_major_version >= 2:
+            if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
+                context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
+                                                                                 is_causal=True)
+            else:
+                if attention_mask is not None:
+                    attention_mask = ~attention_mask
+                context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
+                                                                                 attention_mask)
+            context_layer = context_layer.transpose(1, 2).contiguous()
+            new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
+            context_layer = context_layer.reshape(*new_context_layer_shape)
+        else:
+            # Raw attention scores
 
-        # Raw attention scores. [b * np, sq, sk]
-        matmul_result = torch.baddbmm(
-            matmul_input_buffer,
-            query_layer,  # [b * np, sq, hn]
-            key_layer.transpose(1, 2),  # [b * np, hn, sk]
-            beta=0.0,
-            alpha=(1.0 / self.norm_factor),
-        )
+            # [b, np, sq, sk]
+            output_size = (query_layer.size(0), query_layer.size(1), query_layer.size(2), key_layer.size(2))
 
-        # change view to [b, np, sq, sk]
-        attention_scores = matmul_result.view(*output_size)
-
-        # ===========================
-        # Attention probs and dropout
-        # ===========================
-
-        # attention scores and attention mask [b, np, sq, sk]
-        if self.attention_softmax_in_fp32:
-            attention_scores = attention_scores.float()
-        if self.coeff is not None:
-            attention_scores = attention_scores * self.coeff
-        if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
-            attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
-                                        device=attention_scores.device, dtype=torch.bool)
-            attention_mask.tril_()
-            attention_mask = ~attention_mask
-        if attention_mask is not None:
-            attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
-        attention_probs = F.softmax(attention_scores, dim=-1)
-        attention_probs = attention_probs.type_as(value_layer)
-
-        # This is actually dropping out entire tokens to attend to, which might
-        # seem a bit unusual, but is taken from the original Transformer paper.
-        attention_probs = self.attention_dropout(attention_probs)
-
-        # query layer shape: [b * np, sq, hn]
-        # value layer shape: [b, np, sk, hn]
-        # attention shape: [b, np, sq, sk]
-        # context layer shape: [b, np, sq, hn]
-        output_size = (value_layer.size(0), value_layer.size(1), query_layer.size(1), value_layer.size(3))
-        # change view [b * np, sk, hn]
-        value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)
-        # change view [b * np, sq, sk]
-        attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
-        # matmul: [b * np, sq, hn]
-        context_layer = torch.bmm(attention_probs, value_layer)
-        # change view [b, np, sq, hn]
-        context_layer = context_layer.view(*output_size)
-        # [b, np, sq, hn] --> [b, sq, np, hn]
-        context_layer = context_layer.transpose(1, 2).contiguous()
-        # [b, sq, np, hn] --> [b, sq, hp]
-        new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
-        context_layer = context_layer.reshape(*new_context_layer_shape)
+            # [b, np, sq, hn] -> [b * np, sq, hn]
+            query_layer = query_layer.view(output_size[0] * output_size[1], output_size[2], -1)
+            # [b, np, sk, hn] -> [b * np, sk, hn]
+            key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)
 
-        return context_layer
+            # preallocting input tensor: [b * np, sq, sk]
+            matmul_input_buffer = torch.empty(
+                output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,
+                device=query_layer.device
+            )
+
+            # Raw attention scores. [b * np, sq, sk]
+            matmul_result = torch.baddbmm(
+                matmul_input_buffer,
+                query_layer,  # [b * np, sq, hn]
+                key_layer.transpose(1, 2),  # [b * np, hn, sk]
+                beta=0.0,
+                alpha=(1.0 / self.norm_factor),
+            )
+
+            # change view to [b, np, sq, sk]
+            attention_scores = matmul_result.view(*output_size)
+
+            # ===========================
+            # Attention probs and dropout
+            # ===========================
+
+            # attention scores and attention mask [b, np, sq, sk]
+            if self.attention_softmax_in_fp32:
+                attention_scores = attention_scores.float()
+            if self.coeff is not None:
+                attention_scores = attention_scores * self.coeff
+            if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
+                attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
+                                            device=attention_scores.device, dtype=torch.bool)
+                attention_mask.tril_()
+                attention_mask = ~attention_mask
+            if attention_mask is not None:
+                attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
+            attention_probs = F.softmax(attention_scores, dim=-1)
+            attention_probs = attention_probs.type_as(value_layer)
+
+            # This is actually dropping out entire tokens to attend to, which might
+            # seem a bit unusual, but is taken from the original Transformer paper.
+            attention_probs = self.attention_dropout(attention_probs)
+            # =========================
+            # Context layer. [sq, b, hp]
+            # =========================
+
+            # value_layer -> context layer.
+            # [sk, b, np, hn] --> [b, np, sq, hn]
+
+            # context layer shape: [b, np, sq, hn]
+            output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
+            # change view [b * np, sk, hn]
+            value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)
+            # change view [b * np, sq, sk]
+            attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
+            # matmul: [b * np, sq, hn]
+            context_layer = torch.bmm(attention_probs, value_layer)
+            # change view [b, np, sq, hn]
+            context_layer = context_layer.view(*output_size)
+            # [b, np, sq, hn] --> [b, sq, np, hn]
+            context_layer = context_layer.transpose(1, 2).contiguous()
+            # [b, sq, np, hn] --> [b, sq, hp]
+            new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
+            context_layer = context_layer.reshape(*new_context_layer_shape)
 
+        return context_layer
 
 class SdpaAttention(CoreAttention):
     def forward(self, query_layer, key_layer, value_layer, attention_mask):
@@ -388,7 +450,6 @@ CORE_ATTENTION_CLASSES = {
     "flash_attention_2": FlashAttention2
 }
 
-
 class SelfAttention(torch.nn.Module):
     """Parallel self-attention layer abstract class.
 
@@ -408,6 +469,7 @@ class SelfAttention(torch.nn.Module):
 
         self.multi_query_attention = config.multi_query_attention
         self.qkv_hidden_size = 3 * self.projection_size
+        self.original_rope = config.original_rope
         if self.multi_query_attention:
             self.num_multi_query_groups_per_partition = config.multi_query_group_num
             self.qkv_hidden_size = (
@@ -418,7 +480,7 @@ class SelfAttention(torch.nn.Module):
                                          device=device, **_config_to_kwargs(config)
                                          )
 
-        self.core_attention = CORE_ATTENTION_CLASSES[config._attn_implementation](config, self.layer_number)
+        self.core_attention = CoreAttention(config, self.layer_number)
 
         # Output.
         self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,
@@ -496,11 +558,7 @@ class SelfAttention(torch.nn.Module):
             key_layer = torch.cat((cache_k, key_layer), dim=2)
             value_layer = torch.cat((cache_v, value_layer), dim=2)
         if use_cache:
-            if kv_cache is None:
-                kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)),
-                                     dim=1)
-            else:
-                kv_cache = (key_layer, value_layer)
+            kv_cache = (key_layer, value_layer)
         else:
             kv_cache = None
 
@@ -733,15 +791,7 @@ class GLMTransformer(torch.nn.Module):
                 )
             hidden_states, kv_cache = layer_ret
             if use_cache:
-                # token by token decoding, use tuple format
-                if kv_caches[0] is not None:
-                    presents = presents + (kv_cache,)
-                # prefilling in decoding, use tensor format to save cuda memory
-                else:
-                    if len(presents) == 0:
-                        presents = kv_cache
-                    else:
-                        presents = torch.cat((presents, kv_cache.to(presents.device)), dim=0)
+                presents = presents + (kv_cache,)
 
         if output_hidden_states:
             all_hidden_states = all_hidden_states + (hidden_states,)
@@ -771,20 +821,16 @@ class ChatGLMPreTrainedModel(PreTrainedModel):
         """Initialize the weights."""
         return
 
-    def get_masks(self, input_ids, past_key_values, padding_mask=None):
-        if self.config._attn_implementation == "flash_attention_2":
-            if padding_mask is not None and not padding_mask.all():
-                return padding_mask
-            return None
-        batch_size, seq_length = input_ids.shape
-        full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
+    def get_masks(self, input_embeds, past_key_values, padding_mask=None):
+        batch_size, seq_length, embed_size = input_embeds.shape
+        full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_embeds.device)
         full_attention_mask.tril_()
         past_length = 0
         if past_key_values:
             past_length = past_key_values[0][0].shape[2]
         if past_length:
             full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,
-                                                        device=input_ids.device), full_attention_mask), dim=-1)
+                                                        device=input_embeds.device), full_attention_mask), dim=-1)
         if padding_mask is not None:
             full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)
         if not past_length and padding_mask is not None:
@@ -798,6 +844,9 @@ class ChatGLMPreTrainedModel(PreTrainedModel):
         position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
         return position_ids
 
+    def get_multimodal_position_ids(self, input_ids, device):
+        batch_size, seq_length = input_ids.shape
+        position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
 
 class Embedding(torch.nn.Module):
     """Language model embeddings."""
@@ -825,6 +874,15 @@ class Embedding(torch.nn.Module):
         return embeddings
 
 
+def is_empty(images_list: Optional[List[List[torch.Tensor]]]):
+    if images_list is None or len(images_list) == 0:
+        return True
+    for image_list in images_list:
+        if image_list is not None:
+            return False
+    return True
+
+
 class ChatGLMModel(ChatGLMPreTrainedModel):
     def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):
         super().__init__(config)
@@ -852,6 +910,16 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
         self.encoder = init_method(GLMTransformer, config, **init_kwargs)
         self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
                                         dtype=config.torch_dtype, **init_kwargs)
+        self.pre_seq_len = config.pre_seq_len
+        self.prefix_projection = config.prefix_projection
+        if self.pre_seq_len is not None:
+            for param in self.parameters():
+                param.requires_grad = False
+            self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+            self.prefix_encoder = PrefixEncoder(config)
+            self.dropout = torch.nn.Dropout(0.1)
+
+        self.vision = EVA2CLIPModel(config)
 
     def get_input_embeddings(self):
         return self.embedding.word_embeddings
@@ -859,19 +927,70 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embedding.word_embeddings = value
 
+    def get_prompt(self, batch_size, device, dtype=torch.half):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
+        past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)
+        past_key_values = past_key_values.view(
+            batch_size,
+            self.pre_seq_len,
+            self.pre_seq_len,
+            self.num_layers * 2,
+            self.multi_query_group_num,
+            self.kv_channels
+        )
+        # seq_len, b, nh, hidden_size
+        past_key_values = self.dropout(past_key_values)
+        past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)
+        return past_key_values
+
     def forward(
             self,
-            input_ids,
+            input_ids: torch.LongTensor = None,
+            images: torch.Tensor = None,
             position_ids: Optional[torch.Tensor] = None,
             attention_mask: Optional[torch.BoolTensor] = None,
             full_attention_mask: Optional[torch.BoolTensor] = None,
             past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
             inputs_embeds: Optional[torch.Tensor] = None,
             use_cache: Optional[bool] = None,
-            output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
-    ):
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        """take care of image_encode, position_ids and (attention_mask = None is fine)"""
+
+        # generate mode with past_key_values. the image features are already mapped
+        if past_key_values is None:
+            # not allow for inputs_embeds, because we want to process image feature
+            assert input_ids is not None and inputs_embeds is None, f"{input_ids} {inputs_embeds}"
+            if not is_empty(images):  # multi-modality
+                image_size: int = self.config.vision_config['image_size']
+                patch_size: int = self.config.vision_config['patch_size']
+                num_patches = (image_size // patch_size // 2) ** 2
+                assert len(input_ids) == len(images), f"{len(input_ids)} {len(images)}"
+                inputs_embeds = self.embedding(input_ids)
+
+                images = images.to(dtype=inputs_embeds.dtype)
+                images_features = self.vision(images)
+
+                if position_ids is None:
+                    position_ids = self.get_position_ids(input_ids, device=inputs_embeds.device)
+                new_input_embeds, new_position_ids = [], []
+
+                for i in range(len(input_ids)):
+                    input_id = input_ids[i].tolist()
+                    boi_token_pos, eoi_token_pos = input_id.index(self.config.boi_token_id), input_id.index(
+                        self.config.eoi_token_id)
+                    assert eoi_token_pos - boi_token_pos == 2
+                    new_input_embeds.append(torch.cat(
+                        (inputs_embeds[i, :boi_token_pos], images_features[i].to(inputs_embeds.device),
+                         inputs_embeds[i, eoi_token_pos + 1:])))
+                    new_position_ids.append(torch.cat(
+                        (position_ids[i, :boi_token_pos + 1], position_ids[i, boi_token_pos + 1].repeat(num_patches),
+                         position_ids[i, eoi_token_pos:])
+                    ))
+                inputs_embeds = torch.stack(new_input_embeds, dim=0)
+                position_ids = torch.stack(new_position_ids, dim=0)
+
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
@@ -883,12 +1002,41 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
         if inputs_embeds is None:
             inputs_embeds = self.embedding(input_ids)
 
+        if self.pre_seq_len is not None:
+            if past_key_values is None:
+                past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
+                                                  dtype=inputs_embeds.dtype)
+            if attention_mask is not None:
+                attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),
+                                            attention_mask], dim=-1)
+
         if full_attention_mask is None:
             if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
-                full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)
+                if self.training:
+                    # https://github.com/THUDM/GLM-4/issues/264
+                    new_input_ids, new_attention_mask = [], []
+                    for i in range(len(input_ids)):
+                        input_id = input_ids[i].tolist()
+                        boi_token_pos, eoi_token_pos = input_id.index(self.config.boi_token_id), input_id.index(self.config.eoi_token_id)
+                        assert eoi_token_pos - boi_token_pos == 2
+
+                        new_attention_mask.append(torch.cat(
+                            (attention_mask[i, :boi_token_pos + 1], torch.ones(num_patches).to(attention_mask.device),
+                             attention_mask[i, eoi_token_pos:])))
+
+                        new_input_ids.append(torch.cat(
+                            (input_ids[i, :boi_token_pos + 1], input_ids[i, -1].repeat(num_patches),
+                             input_ids[i, eoi_token_pos:])))
+
+                    attention_mask = torch.stack(new_attention_mask, dim=0)
+                    input_ids = torch.stack(new_input_ids, dim=0)
+                    inputs_embeds = self.embedding(input_ids)
+
+                full_attention_mask = self.get_masks(inputs_embeds, past_key_values, padding_mask=attention_mask)
 
         # Rotary positional embeddings
         rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
+
         if position_ids is not None:
             rotary_pos_emb = rotary_pos_emb[position_ids]
         else:
@@ -899,12 +1047,6 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
             inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
             kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
         )
-        if presents is not None and type(presents) is torch.Tensor:
-            presents = presents.split(1, dim=0)
-            presents = list(presents)
-            presents = [list(x.squeeze(0).split(1, dim=0)) for x in presents]
-            presents = [tuple([x.squeeze(0) for x in y]) for y in presents]
-            presents = tuple(presents)
 
         if not return_dict:
             return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
@@ -917,6 +1059,16 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
         )
 
 
+def _history_to_prompt(history, query):
+    prompt = ''
+    flag = False
+    for i, (old_query, response) in enumerate(history):
+        prompt += ('<|user|>' if flag else '') + old_query + "<|assistant|>" + response + "<|endoftext|>"
+        flag = True
+    prompt += '{}{}<|assistant|>'.format('<|user|>' if flag else '', query)
+    return prompt
+
+
 class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
     def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
         super().__init__(config)
@@ -930,9 +1082,12 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             outputs: ModelOutput,
             model_kwargs: Dict[str, Any],
             is_encoder_decoder: bool = False,
+            standardize_cache_format: bool = False,
     ) -> Dict[str, Any]:
         # update past_key_values
-        cache_name, cache = self._extract_past_from_model_output(outputs)
+        cache_name, cache = self._extract_past_from_model_output(
+            outputs, standardize_cache_format=standardize_cache_format
+        )
         model_kwargs[cache_name] = cache
 
         # update attention mask
@@ -957,6 +1112,7 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
     def prepare_inputs_for_generation(
             self,
             input_ids: torch.LongTensor,
+            images: Optional[torch.Tensor] = None,
             past_key_values: Optional[torch.Tensor] = None,
             attention_mask: Optional[torch.Tensor] = None,
             position_ids: Optional[torch.Tensor] = None,
@@ -967,12 +1123,34 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
         # only last token for input_ids if past is not None
         if position_ids is None:
             position_ids = self.get_position_ids(input_ids, device=input_ids.device)
+        if attention_mask is not None:
+            image_size: int = self.config.vision_config['image_size']
+            patch_size: int = self.config.vision_config['patch_size']
+            num_patches = (image_size // patch_size // 2) ** 2
+            new_attention_masks = []
+
+            # if not image, use this default id
+            eoi_token_pos = 6
+            boi_token_pos = 4
+
+            for i in range(len(input_ids)):
+                input_id = input_ids[i].tolist()
+                if not is_empty(images):
+                    boi_token_pos, eoi_token_pos = input_id.index(self.config.boi_token_id), input_id.index(
+                        self.config.eoi_token_id)
+                assert eoi_token_pos - boi_token_pos == 2
+                new_attention_masks.append(torch.cat(
+                    (attention_mask[i, :boi_token_pos + 1], attention_mask.new_ones(num_patches),
+                     attention_mask[i, eoi_token_pos:])
+                ))
+            attention_mask = torch.stack(new_attention_masks, dim=0)
         if not is_first_forward:
             if past_key_values is not None:
                 position_ids = position_ids[..., -1:]
                 input_ids = input_ids[:, -1:]
         return {
             "input_ids": input_ids,
+            "images": images,
             "past_key_values": past_key_values,
             "position_ids": position_ids,
             "attention_mask": attention_mask,
@@ -983,6 +1161,7 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
     def forward(
             self,
             input_ids: Optional[torch.Tensor] = None,
+            images: List[List[torch.Tensor]] = None,
             position_ids: Optional[torch.Tensor] = None,
             attention_mask: Optional[torch.Tensor] = None,
             past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
@@ -999,6 +1178,7 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
 
         transformer_outputs = self.transformer(
             input_ids=input_ids,
+            images=images,
             position_ids=position_ids,
             attention_mask=attention_mask,
             past_key_values=past_key_values,
@@ -1015,12 +1195,23 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
 
         loss = None
         if labels is not None:
+            new_labels = []
+            for i in range(len(input_ids)):
+                input_id = input_ids[i].tolist()
+                boi_token_pos, eoi_token_pos = input_id.index(self.config.boi_token_id), input_id.index(
+                    self.config.eoi_token_id)
+                assert eoi_token_pos - boi_token_pos == 2
+
+                new_labels.append(torch.cat(
+                    (
+                        labels[i, :boi_token_pos + 1],
+                        torch.tensor([-100]).to(labels.device).to(labels.dtype).repeat(1600),
+                        labels[i, eoi_token_pos:])))
+
+            labels = torch.stack(new_labels, dim=0)
             lm_logits = lm_logits.to(torch.float32)
-
-            # Shift so that tokens < n predict n
             shift_logits = lm_logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()
-            # Flatten the tokens
             loss_fct = CrossEntropyLoss(ignore_index=-100)
             loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
 
@@ -1058,202 +1249,6 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             for layer_past in past
         )
 
-    def process_response(self, output, history):
-        content = ""
-        history = deepcopy(history)
-        for response in output.split("<|assistant|>"):
-            if "\n" in response:
-                metadata, content = response.split("\n", maxsplit=1)
-            else:
-                metadata, content = "", response
-            if not metadata.strip():
-                content = content.strip()
-                history.append({"role": "assistant", "metadata": metadata, "content": content})
-                content = content.replace("[[训练时间]]", "2023年")
-            else:
-                history.append({"role": "assistant", "metadata": metadata, "content": content})
-                if history[0]["role"] == "system" and "tools" in history[0]:
-                    parameters = json.loads(content)
-                    content = {"name": metadata.strip(), "parameters": parameters}
-                else:
-                    content = {"name": metadata.strip(), "content": content}
-        return content, history
-
-    @torch.inference_mode()
-    def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
-             max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,
-             **kwargs):
-        if history is None:
-            history = []
-        if logits_processor is None:
-            logits_processor = LogitsProcessorList()
-        logits_processor.append(InvalidScoreLogitsProcessor())
-        gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
-                      "temperature": temperature, "logits_processor": logits_processor, **kwargs}
-        history.append({"role": role, "content": query})
-        inputs = tokenizer.apply_chat_template(history, add_generation_prompt=True, tokenize=True,
-                                               return_tensors="pt", return_dict=True)
-        inputs = inputs.to(self.device)
-        eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),
-                        tokenizer.convert_tokens_to_ids("<|observation|>")]
-        outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id)
-        outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
-        response = tokenizer.decode(outputs)
-        response, history = self.process_response(response, history)
-        return response, history
-
-    @torch.inference_mode()
-    def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
-                    past_key_values=None, max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8,
-                    logits_processor=None, return_past_key_values=False, **kwargs):
-        if history is None:
-            history = []
-        if logits_processor is None:
-            logits_processor = LogitsProcessorList()
-        logits_processor.append(InvalidScoreLogitsProcessor())
-        eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),
-                        tokenizer.convert_tokens_to_ids("<|observation|>")]
-        gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
-                      "temperature": temperature, "logits_processor": logits_processor, **kwargs}
-        if past_key_values is None:
-            inputs = tokenizer.apply_chat_template(history + [{"role": role, "content": query}],
-                                                   add_generation_prompt=True, tokenize=True, return_tensors="pt",
-                                                   return_dict=True)
-        else:
-            inputs = tokenizer.apply_chat_template([{"role": role, "content": query}], add_special_tokens=False,
-                                                   add_generation_prompt=True, tokenize=True, return_tensors="pt",
-                                                   return_dict=True)
-        inputs = inputs.to(self.device)
-        if past_key_values is not None:
-            past_length = past_key_values[0][0].shape[2]
-            inputs.position_ids += past_length
-            attention_mask = inputs.attention_mask
-            attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
-            inputs['attention_mask'] = attention_mask
-        history.append({"role": role, "content": query})
-        for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,
-                                            eos_token_id=eos_token_id, return_past_key_values=return_past_key_values,
-                                            **gen_kwargs):
-            if return_past_key_values:
-                outputs, past_key_values = outputs
-            outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
-            response = tokenizer.decode(outputs)
-            if response and response[-1] != "�":
-                response, new_history = self.process_response(response, history)
-                if return_past_key_values:
-                    yield response, new_history, past_key_values
-                else:
-                    yield response, new_history
-
-    @torch.inference_mode()
-    def stream_generate(
-            self,
-            input_ids,
-            generation_config: Optional[GenerationConfig] = None,
-            logits_processor: Optional[LogitsProcessorList] = None,
-            stopping_criteria: Optional[StoppingCriteriaList] = None,
-            prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
-            return_past_key_values=False,
-            **kwargs,
-    ):
-        batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
-
-        if generation_config is None:
-            generation_config = self.generation_config
-        generation_config = copy.deepcopy(generation_config)
-        model_kwargs = generation_config.update(**kwargs)
-        model_kwargs["use_cache"] = generation_config.use_cache
-        bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
-
-        if isinstance(eos_token_id, int):
-            eos_token_id = [eos_token_id]
-        eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
-
-        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
-        if has_default_max_length and generation_config.max_new_tokens is None:
-            warnings.warn(
-                f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
-                "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
-                " recommend using `max_new_tokens` to control the maximum length of the generation.",
-                UserWarning,
-            )
-        elif generation_config.max_new_tokens is not None:
-            generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
-            if not has_default_max_length:
-                logger.warn(
-                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
-                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
-                    "Please refer to the documentation for more information. "
-                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
-                    UserWarning,
-                )
-
-        if input_ids_seq_length >= generation_config.max_length:
-            input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
-            logger.warning(
-                f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
-                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
-                " increasing `max_new_tokens`."
-            )
-
-        # 2. Set generation parameters if not already defined
-        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
-        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
-
-        logits_processor = self._get_logits_processor(
-            generation_config=generation_config,
-            input_ids_seq_length=input_ids_seq_length,
-            encoder_input_ids=input_ids,
-            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
-            logits_processor=logits_processor,
-        )
-
-        stopping_criteria = self._get_stopping_criteria(
-            generation_config=generation_config, stopping_criteria=stopping_criteria
-        )
-        logits_warper = self._get_logits_warper(generation_config)
-
-        unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
-        scores = None
-        while True:
-            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
-            # forward pass to get next token
-            outputs = self(
-                **model_inputs,
-                return_dict=True,
-                output_attentions=False,
-                output_hidden_states=False,
-            )
-
-            next_token_logits = outputs.logits[:, -1, :]
-
-            # pre-process distribution
-            next_token_scores = logits_processor(input_ids, next_token_logits)
-            next_token_scores = logits_warper(input_ids, next_token_scores)
-
-            # sample
-            probs = nn.functional.softmax(next_token_scores, dim=-1)
-            if generation_config.do_sample:
-                next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
-            else:
-                next_tokens = torch.argmax(probs, dim=-1)
-            # update generated ids, model inputs, and length for next step
-            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
-            model_kwargs = self._update_model_kwargs_for_generation(
-                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
-            )
-            unfinished_sequences = unfinished_sequences.mul(
-                next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
-            )
-            if return_past_key_values:
-                yield input_ids, outputs.past_key_values
-            else:
-                yield input_ids
-            # stop when each sentence is finished, or if we exceed the maximum length
-            if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
-                break
-
-
 class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
     def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
         super().__init__(config)
@@ -1261,7 +1256,7 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
         self.num_labels = config.num_labels
         self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
 
-        self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=config.torch_dtype)
+        self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=torch.half)
         if config.classifier_dropout is not None:
             self.dropout = nn.Dropout(config.classifier_dropout)
         else:
@@ -1278,7 +1273,6 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
             inputs_embeds: Optional[torch.LongTensor] = 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[torch.Tensor, ...], SequenceClassifierOutputWithPast]:
@@ -1292,13 +1286,12 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
             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 = transformer_outputs[0]
-        pooled_hidden_states = hidden_states[:, -1]
+        pooled_hidden_states = hidden_states[-1]
         if self.dropout is not None:
             pooled_hidden_states = self.dropout(pooled_hidden_states)
         logits = self.classifier_head(pooled_hidden_states)
@@ -1336,4 +1329,4 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
             past_key_values=transformer_outputs.past_key_values,
             hidden_states=transformer_outputs.hidden_states,
             attentions=transformer_outputs.attentions,
-        )
+        )

tokenization_chatglm.py

@@ -3,8 +3,10 @@ import base64
 import os
 import json
 import tiktoken
+import torch
 from torch import TensorType
 from typing import List, Optional, Union, Dict, Any
+from torchvision import transforms
 from transformers import PreTrainedTokenizer
 from transformers.utils import logging, PaddingStrategy
 from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
@@ -20,6 +22,7 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
             padding_side="left",
             clean_up_tokenization_spaces=False,
             encode_special_tokens=False,
+            image_size=None,
             **kwargs
     ):
         self.name = "GLM4Tokenizer"
@@ -27,6 +30,7 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
         pat_str = "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
         self.pat_str = re.compile(pat_str)
         self.encode_special_tokens = encode_special_tokens
+        self.image_size = image_size
 
         mergeable_ranks = {}
         with open(vocab_file) as f:
@@ -130,109 +134,143 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
         prefix_tokens = [self.convert_tokens_to_ids("[gMASK]"), self.convert_tokens_to_ids("<sop>")]
         return prefix_tokens
 
-    def build_single_message(self, role, metadata, message, tokenize=True):
+    def build_single_message(self, role, metadata, message, tokenize=True, message_prefix=None):
         assert role in ["system", "user", "assistant", "observation"], role
         if tokenize:
             role_tokens = [self.convert_tokens_to_ids(f"<|{role}|>")] + self.tokenizer.encode(f"{metadata}\n",
                                                                                               disallowed_special=())
             message_tokens = self.tokenizer.encode(message, disallowed_special=())
+            if message_prefix is not None:
+                message_tokens = message_prefix + message_tokens
             tokens = role_tokens + message_tokens
             return tokens
         else:
             return str(f"<|{role}|>{metadata}\n{message}")
 
-    # Use Jinja Template in tokenizer_config.json
-    # def apply_chat_template(
-    #         self,
-    #         conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]], "Conversation"],
-    #         add_generation_prompt: bool = False,
-    #         tokenize: bool = True,
-    #         padding: bool = False,
-    #         truncation: bool = False,
-    #         max_length: Optional[int] = None,
-    #         return_tensors: Optional[Union[str, TensorType]] = None,
-    #         return_dict: bool = False,
-    #         tokenizer_kwargs: Optional[Dict[str, Any]] = None,
-    #         add_special_tokens: bool = True,
-    #         **kwargs,
-    # ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]:
-    #
-    #     if return_dict and not tokenize:
-    #         raise ValueError(
-    #             "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict "
-    #             "of tokenizer outputs to return."
-    #         )
-    #
-    #     def handle_single_conversation(conversation):
-    #         input_ids = self.get_prefix_tokens() if add_special_tokens else []
-    #         input_message = "[gMASK]<sop>" if add_special_tokens else ""
-    #         for item in conversation:
-    #             if item.get("tools"):
-    #                 tools = item["tools"]
-    #                 content = "你是一个名为 GhatGLM 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。"
-    #                 content += "\n\n# 可用工具"
-    #                 for tool in tools:
-    #                     if tool["type"] == "function":
-    #                         function = tool["function"]
-    #                         content += f"\n\n## {function['name']}\n\n{json.dumps(function, ensure_ascii=False, indent=4)}"
-    #                         content += "\n在调用上述函数时，请使用 Json 格式表示调用的参数。"
-    #                     elif tool["type"] == "python":
-    #                         content += "\n\n## python\n\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。"
-    #                     elif tool["type"] == "simple_browser":
-    #                         content += "\n\n## simple_browser\n\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效���。\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\n`open_url(url: str)`：打开指定的 URL。\n\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\n\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。"
-    #                     elif tool["type"] == "cogview":
-    #                         content += "\n\n## cogview\n\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。"
-    #                     else:
-    #                         raise NotImplementedError(f"Unknown tool type {tool['type']}")
-    #                 input = self.build_single_message("system", "", content, tokenize=tokenize)
-    #                 if tokenize:
-    #                     input_ids.extend(input)
-    #                 else:
-    #                     input_message += input
-    #             if item["content"]:
-    #                 input = self.build_single_message(
-    #                     item["role"],
-    #                     item.get("metadata", ""),
-    #                     item["content"],
-    #                     tokenize=tokenize
-    #                 )
-    #                 if tokenize:
-    #                     input_ids.extend(input)
-    #                 else:
-    #                     input_message += input
-    #         if add_generation_prompt:
-    #             if tokenize:
-    #                 input_ids.extend([self.convert_tokens_to_ids("<|assistant|>")])
-    #             else:
-    #                 input_message += "<|assistant|>"
-    #         return input_ids if tokenize else input_message
-    #
-    #     # Main logic to handle different conversation formats
-    #     if isinstance(conversation, list) and all(isinstance(i, dict) for i in conversation):
-    #         result = handle_single_conversation(conversation)
-    #     elif isinstance(conversation, list) and all(isinstance(i, list) for i in conversation):
-    #         result = [handle_single_conversation(c) for c in conversation]
-    #     elif hasattr(conversation, "messages"):
-    #         result = handle_single_conversation(conversation.messages)
-    #     else:
-    #         raise ValueError("Invalid conversation format")
-    #
-    #     if tokenize:
-    #         output = self.batch_encode_plus(
-    #             [result] if isinstance(result[0], int) else result,
-    #             padding=padding,
-    #             truncation=truncation,
-    #             max_length=max_length,
-    #             return_tensors=return_tensors,
-    #             is_split_into_words=True,
-    #             add_special_tokens=False
-    #         )
-    #         if return_dict:
-    #             return output
-    #         else:
-    #             return output["input_ids"]
-    #     else:
-    #         return result
+    def apply_chat_template(
+            self,
+            conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]], "Conversation"],
+            add_generation_prompt: bool = False,
+            tokenize: bool = True,
+            padding: bool = False,
+            truncation: bool = False,
+            max_length: Optional[int] = None,
+            return_tensors: Optional[Union[str, TensorType]] = None,
+            return_dict: bool = False,
+            tokenizer_kwargs: Optional[Dict[str, Any]] = None,
+            add_special_tokens: bool = True,
+            **kwargs,
+    ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]:
+
+        if return_dict and not tokenize:
+            raise ValueError(
+                "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict "
+                "of tokenizer outputs to return."
+            )
+
+        def handle_single_conversation(conversation):
+            input_ids = self.get_prefix_tokens() if add_special_tokens else []
+            input_message = "[gMASK]<sop>" if add_special_tokens else ""
+            input_image = None
+            transform = transforms.Compose(
+                [
+                    transforms.Resize(
+                        (self.image_size, self.image_size), interpolation=transforms.InterpolationMode.BICUBIC
+                    ),
+                    transforms.ToTensor(),
+                    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+                ]
+            )
+            for item in conversation:
+                if item.get("tools"):
+                    tools = item["tools"]
+                    content = "你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。"
+                    for tool in tools:
+                        if tool["type"] == "function":
+                            function = tool["function"]
+                            content += f"\n\n## {function['name']}\n\n{json.dumps(function, ensure_ascii=False, indent=4)}"
+                            content += "\n在调用上述函数时，请使用 Json 格式表示调用的参数。"
+                        elif tool["type"] == "python":
+                            content += "\n\n## python\n\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。"
+                        elif tool["type"] == "simple_browser":
+                            content += "\n\n## simple_browser\n\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\n`open_url(url: str)`：打开指定的 URL。\n\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\n\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。"
+                        elif tool["type"] == "cogview":
+                            content += "\n\n## cogview\n\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。"
+                        else:
+                            raise NotImplementedError(f"Unknown tool type {tool['type']}")
+                    input = self.build_single_message("system", "", content, tokenize=tokenize)
+                    if tokenize:
+                        input_ids.extend(input)
+                    else:
+                        input_message += input
+                message = ""
+                message_prefix = None
+                if item.get("image"):
+                    assert input_image is None, "Multiple images are not supported"
+                    input_image = transform(item["image"])
+                    message_prefix = self.convert_tokens_to_ids(
+                        ["<|begin_of_image|>", "<|endoftext|>", "<|end_of_image|>"])
+                if item.get("content"):
+                    message += item["content"]
+                if message or message_prefix:
+                    input = self.build_single_message(
+                        item["role"],
+                        item.get("metadata", ""),
+                        message,
+                        tokenize=tokenize,
+                        message_prefix=message_prefix
+                    )
+                    if tokenize:
+                        input_ids.extend(input)
+                    else:
+                        input_message += input
+            if add_generation_prompt:
+                if tokenize:
+                    input_ids.extend([self.convert_tokens_to_ids("<|assistant|>")])
+                else:
+                    input_message += "<|assistant|>"
+            return {"input": input_ids if tokenize else input_message, "image": input_image}
+
+        # Main logic to handle different conversation formats
+        if isinstance(conversation, list) and all(isinstance(i, dict) for i in conversation):
+            result = handle_single_conversation(conversation)
+            input_ids = result["input"]
+            input_images = [result["image"]]
+        elif isinstance(conversation, list) and all(isinstance(i, list) for i in conversation):
+            results = [handle_single_conversation(c) for c in conversation]
+            input_ids = [item["input"] for item in results]
+            input_images = [item["image"] for item in results]
+        elif hasattr(conversation, "messages"):
+            result = handle_single_conversation(conversation.messages)
+            input_ids = result["input"]
+            input_images = [result["image"]]
+        else:
+            raise ValueError("Invalid conversation format")
+
+        if tokenize:
+            output = self.batch_encode_plus(
+                [input_ids] if isinstance(input_ids[0], int) else input_ids,
+                padding=padding,
+                truncation=truncation,
+                max_length=max_length,
+                return_tensors=return_tensors,
+                is_split_into_words=True,
+                add_special_tokens=False
+            )
+            if return_dict:
+                found_image = False
+                for image in input_images:
+                    if image is not None:
+                        found_image = True
+                        break
+                if found_image:
+                    output["images"] = torch.stack(input_images)
+                return output
+            else:
+                return output["input_ids"]
+        else:
+            return input_ids
+
 
     def build_inputs_with_special_tokens(
             self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None

tokenizer_config.json

@@ -123,12 +123,12 @@
                                "<|user|>", "<|assistant|>", "<|observation|>", "<|begin_of_image|>", "<|end_of_image|>",
                                "<|begin_of_video|>", "<|end_of_video|>"],
   "clean_up_tokenization_spaces": false,
-  "chat_template": "[gMASK]<sop>{% for item in messages %}{% if item['tools'] is defined %}<|system|>\n你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。\n\n# 可用工具{% set tools = item['tools'] %}{% for tool in tools %}{% if tool['type'] == 'function' %}\n\n## {{ tool['function']['name'] }}\n\n{{ tool['function'] | tojson(indent=4) }}\n在调用上述函数时，请使用 Json 格式表示调用的参数。{% elif tool['type'] == 'python' %}\n\n## python\n\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。{% elif tool['type'] == 'simple_browser' %}\n\n## simple_browser\n\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\n`open_url(url: str)`：打开指定的 URL。\n\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\n\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。{% elif tool['type'] == 'cogview' %}\n\n## cogview\n\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。{% endif %}{% endfor %}{% endif %}{% if item['content'] %}<|{{ item['role'] }}|>{{ item['metadata'] }}\n{{ item['content'] }}{% endif %}{% endfor %}{% if add_generation_prompt %}<|assistant|>{% endif %}",
   "do_lower_case": false,
   "eos_token": "<|endoftext|>",
   "pad_token": "<|endoftext|>",
-  "model_max_length": 1024000,
+  "model_max_length": 8192,
   "padding_side": "left",
   "remove_space": false,
-  "tokenizer_class": "ChatGLM4Tokenizer"
+  "tokenizer_class": "ChatGLM4Tokenizer",
+  "image_size": 1120
 }

visual.py

@@ -0,0 +1,180 @@
+import torch
+from torch import nn
+from argparse import Namespace
+import torch.nn.functional as F
+from transformers.activations import ACT2FN
+import math
+from torch.nn import LayerNorm
+
+
+def standard_attention(query_layer, key_layer, value_layer, scaling_attention_score=True):
+    if scaling_attention_score:
+        query_layer = query_layer / math.sqrt(query_layer.shape[-1])
+    attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+    attention_probs = F.softmax(attention_scores, dim=-1)
+
+    context_layer = torch.matmul(attention_probs, value_layer)
+    return context_layer
+
+
+def attention_fn_default(query_layer, key_layer, value_layer, scaling_attention_score=True):
+    if int(torch.__version__.split('.')[0]) >= 2 and scaling_attention_score:
+        # Pytorch 2.0 attention uses very much memory if attention_mask is float, and has NaN bug if attention_mask is None.
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_layer, key_layer, value_layer,
+            attn_mask=None,
+            dropout_p=0.,
+            is_causal=False
+        )
+        return attn_output
+    else:
+        return standard_attention(
+            query_layer, key_layer, value_layer, scaling_attention_score=scaling_attention_score
+        )
+
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.proj = nn.Conv2d(config.in_channels, config.hidden_size, kernel_size=config.patch_size,
+                              stride=config.patch_size)
+        self.cls_embedding = nn.Parameter(torch.zeros(1, config.hidden_size))
+        self.position_embedding = nn.Embedding(config.num_positions, config.hidden_size)
+
+    def forward(self, images: "tensor(B, C, H, W)") -> "tensor(B, L, D)":
+        x = self.proj(images)
+        x = x.flatten(2).transpose(1, 2)
+        cls_token = self.cls_embedding.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_token, x), dim=1)
+        x += self.position_embedding.weight.unsqueeze(0)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.num_heads = config.num_heads
+        head_dim = config.hidden_size // config.num_heads
+        self.scale = head_dim ** -0.5
+        self.query_key_value = nn.Linear(config.hidden_size, config.hidden_size * 3)
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.output_dropout = torch.nn.Dropout(config.dropout_prob)
+
+    def forward(self, x: "tensor(B, L, D)") -> "tensor(B, L, D)":
+        B, L, _ = x.shape
+        qkv = self.query_key_value(x)
+        qkv = qkv.reshape(B, L, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)  # 3, B, H, L, D
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        out = attention_fn_default(
+            q, k, v
+        )
+        output = self.dense(out.transpose(1, 2).reshape(B, L, -1))
+        output = self.output_dropout(output)
+        return output
+
+    def attention(self, q, k, v):
+        attn_weights = torch.matmul(q * self.scale, k.transpose(-2, -1))
+        attn_weights = attn_weights.softmax(dim=-1)
+        output = torch.matmul(attn_weights, v)
+        return output
+
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.activation_fn(x)
+        x = self.fc2(x)
+        return x
+
+
+class TransformerLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.input_layernorm = LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.attention = Attention(config)
+        self.mlp = MLP(config)
+        self.post_attention_layernorm = LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        attention_input = hidden_states
+        attention_output = self.input_layernorm(self.attention(attention_input))
+        hidden_states = attention_input + attention_output
+        mlp_input = hidden_states
+
+        # https://github.com/THUDM/GLM-4/issues/350
+        mlp_output = self.post_attention_layernorm(self.mlp(mlp_input)).to(mlp_input.device)
+        output = mlp_input + mlp_output
+        return output
+
+
+class Transformer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.layers = nn.ModuleList([TransformerLayer(config) for _ in range(config.num_hidden_layers)])
+
+    def forward(self, hidden_states):
+        for layer_module in self.layers:
+            hidden_states = layer_module(hidden_states)
+        return hidden_states
+
+
+class GLU(nn.Module):
+    def __init__(self, config, in_features):
+        super().__init__()
+        self.linear_proj = nn.Linear(in_features, config.hidden_size, bias=False)
+        self.norm1 = nn.LayerNorm(config.hidden_size)
+        self.act1 = nn.GELU()
+        self.act2 = nn.functional.silu
+        self.dense_h_to_4h = nn.Linear(config.hidden_size, config.ffn_hidden_size, bias=False)
+        self.gate_proj = nn.Linear(config.hidden_size, config.ffn_hidden_size, bias=False)
+        self.dense_4h_to_h = nn.Linear(config.ffn_hidden_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        x = self.linear_proj(x)
+        x = self.act1(self.norm1(x))
+        x = self.act2(self.gate_proj(x)) * self.dense_h_to_4h(x)
+        x = self.dense_4h_to_h(x)
+        return x
+
+
+class EVA2CLIPModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        vision_config = Namespace(**config.vision_config)
+        self.patch_embedding = PatchEmbedding(vision_config)
+        self.transformer = Transformer(vision_config)
+        self.linear_proj = GLU(config, in_features=config.hidden_size)
+        self.conv = nn.Conv2d(in_channels=vision_config.hidden_size, out_channels=config.hidden_size, kernel_size=2,
+                              stride=2)
+        self.boi = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+        self.eoi = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+        self.scaling_factor = vision_config.scaling_factor
+
+    def forward(self, images: "tensor(B, C, H, W)") -> "tensor(B, L, D)":
+        x = self.patch_embedding(images)
+        x = self.transformer(x)
+        x = x[:, 1:]
+
+        b, s, h = x.shape
+        grid_size = int(s ** 0.5)
+        x = x.view(b, grid_size, grid_size, h).permute(0, 3, 1, 2)
+        x = self.conv(x)
+
+        x = x.flatten(2).transpose(1, 2)
+        x = self.linear_proj(x)
+
+        # https://github.com/THUDM/GLM-4/issues/350
+        boi = self.boi.expand(x.shape[0], -1, -1).to(x.device)
+        eoi = self.eoi.expand(x.shape[0], -1, -1).to(x.device)
+        x = torch.cat((boi, x, eoi), dim=1)
+        x = x / self.scaling_factor
+        return x