2/n llava e2e init

Author: Gasoonjia

torchchat/cli/convert_hf_checkpoint.py

@@ -7,10 +7,13 @@
 import os
 import re
 import sys
+import glob
 from pathlib import Path
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import torch
+import safetensors.torch
+import shutil
 
 # support running without installing as a package
 wd = Path(__file__).parent.parent
@@ -24,34 +27,34 @@ def _translate_state_dict_for_vision_model(hf_state_dict) -> Dict[str, Any]:
         translated_state_dict = {}
         hf_weight_prefix = "vision_model."
         name_mapping = {
-            f"{hf_weight_prefix}embeddings.class_embedding": "model.encoder.cls_token_embedding.weight",
-            f"{hf_weight_prefix}embeddings.position_embedding.weight": "model.encoder.token_pos_embedding.positional_embedding",
-            f"{hf_weight_prefix}embeddings.patch_embedding.weight": "model.encoder.conv.weight",
-            f"{hf_weight_prefix}pre_layrnorm.weight": "model.encoder.ln_pre.weight",
-            f"{hf_weight_prefix}pre_layrnorm.bias": "model.encoder.ln_pre.bias",
-            f"{hf_weight_prefix}post_layernorm.weight": "model.encoder.ln_post.weight",
-            f"{hf_weight_prefix}post_layernorm.bias": "model.encoder.ln_post.bias",
+            f"{hf_weight_prefix}embeddings.class_embedding": "encoder.cls_token_embedding.weight",
+            f"{hf_weight_prefix}embeddings.position_embedding.weight": "encoder.token_pos_embedding.positional_embedding",
+            f"{hf_weight_prefix}embeddings.patch_embedding.weight": "encoder.conv.weight",
+            f"{hf_weight_prefix}pre_layrnorm.weight": "encoder.ln_pre.weight",
+            f"{hf_weight_prefix}pre_layrnorm.bias": "encoder.ln_pre.bias",
+            f"{hf_weight_prefix}post_layernorm.weight": "encoder.ln_post.weight",
+            f"{hf_weight_prefix}post_layernorm.bias": "encoder.ln_post.bias",
         }
         patterns = [
             (
                 rf"{hf_weight_prefix}encoder\.layers\.([0-9]+)\.self_attn\.(k|q|v)_proj\.(weight|bias)",
-                lambda match: f"model.encoder.layers.{match.group(1)}.attn.{match.group(2)}_proj.{match.group(3)}",
+                lambda match: f"encoder.layers.{match.group(1)}.attn.{match.group(2)}_proj.{match.group(3)}",
             ),
             (
                 rf"{hf_weight_prefix}encoder\.layers\.([0-9]+)\.self_attn\.out_proj\.(weight|bias)",
-                lambda match: f"model.encoder.layers.{match.group(1)}.attn.output_proj.{match.group(2)}",
+                lambda match: f"encoder.layers.{match.group(1)}.attn.output_proj.{match.group(2)}",
             ),
             (
                 rf"{hf_weight_prefix}encoder\.layers\.([0-9]+)\.mlp\.fc(1|2)\.(weight|bias)",
-                lambda match: f"model.encoder.layers.{match.group(1)}.mlp.w{match.group(2)}.{match.group(3)}",
+                lambda match: f"encoder.layers.{match.group(1)}.mlp.w{match.group(2)}.{match.group(3)}",
             ),
             (
                 rf"{hf_weight_prefix}encoder\.layers\.([0-9]+)\.layer_norm1\.(weight|bias)",
-                lambda match: f"model.encoder.layers.{match.group(1)}.sa_norm.{match.group(2)}",
+                lambda match: f"encoder.layers.{match.group(1)}.sa_norm.{match.group(2)}",
             ),
             (
                 rf"{hf_weight_prefix}encoder\.layers\.([0-9]+)\.layer_norm2\.(weight|bias)",
-                lambda match: f"model.encoder.layers.{match.group(1)}.mlp_norm.{match.group(2)}",
+                lambda match: f"encoder.layers.{match.group(1)}.mlp_norm.{match.group(2)}",
             ),
         ]
         for pattern, replacement in patterns:
@@ -82,18 +85,18 @@ def _translate_state_dict_for_vision_model(hf_state_dict) -> Dict[str, Any]:
 
     def _translate_state_dict_for_text_model(hf_state_dict) -> Dict[str, Any]:
         key_map = {
-            r"model.layers.([0-9]+).self_attn.q_proj.": r"model.decoder.layers.\1.attention.wq.",
-            r"model.layers.([0-9]+).self_attn.k_proj.": r"model.decoder.layers.\1.attention.wk.",
-            r"model.layers.([0-9]+).self_attn.v_proj.": r"model.decoder.layers.\1.attention.wv.",
-            r"model.layers.([0-9]+).self_attn.o_proj.": r"model.decoder.layers.\1.attention.wo.",
-            r"model.layers.([0-9]+).input_layernorm.": r"model.decoder.layers.\1.attention_norm.",
-            r"model.layers.([0-9]+).mlp.gate_proj.": r"model.decoder.layers.\1.feed_forward.w1.",
-            r"model.layers.([0-9]+).mlp.down_proj.": r"model.decoder.layers.\1.feed_forward.w2.",
-            r"model.layers.([0-9]+).mlp.up_proj.": r"model.decoder.layers.\1.feed_forward.w3.",
-            r"model.layers.([0-9]+).post_attention_layernorm.": r"model.decoder.layers.\1.ffn_norm.",
-            r"model.norm.": r"model.decoder.norm.",
+            r"model.layers.([0-9]+).self_attn.q_proj.": r"decoder.layers.\1.attention.wq.",
+            r"model.layers.([0-9]+).self_attn.k_proj.": r"decoder.layers.\1.attention.wk.",
+            r"model.layers.([0-9]+).self_attn.v_proj.": r"decoder.layers.\1.attention.wv.",
+            r"model.layers.([0-9]+).self_attn.o_proj.": r"decoder.layers.\1.attention.wo.",
+            r"model.layers.([0-9]+).input_layernorm.": r"decoder.layers.\1.attention_norm.",
+            r"model.layers.([0-9]+).mlp.gate_proj.": r"decoder.layers.\1.feed_forward.w1.",
+            r"model.layers.([0-9]+).mlp.down_proj.": r"decoder.layers.\1.feed_forward.w2.",
+            r"model.layers.([0-9]+).mlp.up_proj.": r"decoder.layers.\1.feed_forward.w3.",
+            r"model.layers.([0-9]+).post_attention_layernorm.": r"decoder.layers.\1.ffn_norm.",
+            r"model.norm.": r"decoder.norm.",
             # r"model.embed_tokens.": r"tok_embeddings.", # load separately
-            r"lm_head.": r"model.decoder.output.",
+            r"lm_head.": r"decoder.output.",
         }
         new_state_dict = {}
         def get_new_key(old_key: str) -> str:
@@ -109,7 +112,7 @@ def get_new_key(old_key: str) -> str:
     def _translate_state_dict_for_mm_projector_model(hf_state_dict) -> Dict[str, Any]:
         new_state_dict = {}
         for old_key in hf_state_dict.keys():
-            new_key = "model.mm_projector." + old_key
+            new_key = "mm_projector." + old_key
             new_state_dict[new_key] = hf_state_dict[old_key]
         return new_state_dict
 
@@ -127,13 +130,65 @@ def split_checkpoint(llava_ckpt):
         return language_model_ckpt, multi_modal_ckpt, vision_tower_ckpt
     language_model_ckpt, multi_modal_ckpt, vision_tower_ckpt = split_checkpoint(llava_ckpt)
     remapped_state_dict = {
-        "model.tok_embeddings.weight": language_model_ckpt.pop("model.embed_tokens.weight"),
+        "tok_embeddings.weight": language_model_ckpt.pop("model.embed_tokens.weight"),
     }
     remapped_state_dict.update(_translate_state_dict_for_text_model(language_model_ckpt))
     remapped_state_dict.update(_translate_state_dict_for_vision_model(vision_tower_ckpt))
     remapped_state_dict.update(_translate_state_dict_for_mm_projector_model(multi_modal_ckpt))
     return remapped_state_dict
 
+    
+@torch.inference_mode
+def convert_llava_checkpoint(    
+    *,
+    model_dir: Optional[Path] = None,
+) -> None:
+    
+    """
+    Process safetensor files from a specific directory structure and save the remapped model.
+    
+    Args:
+        model_dir (str): Base directory containing the model subdirectories.
+    """
+
+    def _get_llava_files_with_pattern(pattern):
+        pattern = os.path.join(model_dir, f"models--llava-hf--llava-1.5-7b-hf/snapshots/*/{pattern}")
+        return glob.glob(pattern)
+
+    # get all safetensor files in the model directory
+    safetensor_files = _get_llava_files_with_pattern("*.safetensors")
+    
+    if not safetensor_files:
+        raise ValueError("No safetensor files found.")
+    
+    merged_weights = {}
+    
+    # Merge safetensor files into a whole
+    for file in safetensor_files:
+        # Load weights from the current file
+        part_weights = safetensors.torch.load_file(file)
+        
+        # Iterate over each weight in the current file
+        for key, value in part_weights.items():
+            if key in merged_weights:
+                # If the key already exists, concatenate tensors
+                merged_weights[key] = torch.cat((merged_weights[key], value), dim=0)
+            else:
+                # If the key does not exist, add it to the dictionary
+                merged_weights[key] = value
+    
+    # Remap the checkpoint and save it as pth
+    remapped_weights = remap_llava_checkpoint(merged_weights)
+    model_path = model_dir / "model.pth"
+    torch.save(remapped_weights, model_path)
+
+    # copy tokenizer
+    tokenizer_files = _get_llava_files_with_pattern("tokenizer.model")
+    assert len(tokenizer_files) == 1, "Should get only one tokenizer file, but got {}".format(tokenizer_files)
+
+    tokenizer_path = model_dir / "tokenizer.model"
+    shutil.copy(tokenizer_files[0], tokenizer_path)
+
 
 @torch.inference_mode()
 def convert_text_only_hf_checkpoint(
@@ -245,18 +300,18 @@ def permute(w, n_heads):
 
 
 @torch.inference_mode()
-def convert_text_only_hf_checkpoint(
+def convert_hf_checkpoint(
     *,
     model_dir: Optional[Path] = None,
     model_name: Optional[str] = None,
     remove_bin_files: bool = False,
 ):
-    if model_name == "llava-1.5":
-        print("Converting LLaVA 1.5 checkpoint.")
-        print(os.listdir(model_dir))
-        exit(0)
+    print(model_name)
+    print("***********************")
+    if "llava" in model_name:
+        convert_llava_checkpoint(model_dir=model_dir)
     else:
-        convert_text_only_hf_checkpoint(model_dir, model_name, remove_bin_files)
+        convert_text_only_hf_checkpoint(model_dir=model_dir, model_name=model_name, remove_bin_files=remove_bin_files)
 
 
 if __name__ == "__main__":

torchchat/cli/download.py

@@ -28,34 +28,14 @@ def _download_hf_snapshot(
     # Download and store the HF model artifacts.
     print(f"Downloading {model_config.name} from HuggingFace...", file=sys.stderr)
     try:
-
-        import huggingface_hub
-        # 定义模型名称和版本
-        model_name = "llava-hf/llava-1.5-7b-hf"
-        # 下载模型checkpoint
-        repo_id = model_name
-        revision = "main"  # 默认分支
-        # 强制重新下载
-        snapshot_dir = huggingface_hub.snapshot_download(
-            repo_id=repo_id,
-            revision=revision,
+        snapshot_download(
+            model_config.distribution_path,
             cache_dir=artifact_dir,
-            force_download=True,
+            local_dir_use_symlinks=False,
+            token=hf_token,
+            ignore_patterns=None if "llava" in model_config.name else "*safetensors*",
         )
-        print(f"模型下载完成，保存在 {snapshot_dir} 目录下")
-
-
-        # snapshot_download(
-        #     model_config.distribution_path,
-        #     cache_dir=artifact_dir,
-        #     local_dir_use_symlinks=False,
-        #     token=hf_token,
-        #     ignore_patterns="*safetensors*",
-        # )
-        print("*****************")
-        print(os.listdir(artifact_dir))
-        shutil.copytree(artifact_dir, "/home/gasoonjia/download/hahaha")
-        exit(0)
+
     except HTTPError as e:
         if e.response.status_code == 401:  # Missing HuggingFace CLI login.
             print(
@@ -99,8 +79,8 @@ def download_and_convert(
     # location once the download and conversion is complete. This
     # allows recovery in the event that the download or conversion
     # fails unexpectedly.
-    # temp_dir = models_dir / "downloads" / model_config.name
-    temp_dir = Path("/home/gasoonjia") / "downloads" / model_config.name
+    temp_dir = models_dir / "downloads" / model_config.name
+    # temp_dir = Path("/home/gasoonjia") / "downloads" / model_config.name
 
     if os.path.isdir(temp_dir):
         shutil.rmtree(temp_dir)

torchchat/generate.py

@@ -36,6 +36,7 @@
 from torchchat.model import Model, ModelType
 from torchchat.utils.build_utils import device_sync, set_precision
 from torchchat.utils.device_info import get_device_info
+from torchchat.utils.preprocessors import llava_image_preprocess
 
 # torchtune model definition dependencies
 from torchtune.data import Message
@@ -622,6 +623,13 @@ def generate(
             sequential_prefill=sequential_prefill,
             **sampling_kwargs,
         )
+
+        # For llava, we need to extract next pos id from prefill result
+        if self.model.config.model_type == ModelType.Llava:
+            next_token, context_len = next_token
+        else:
+            next_token, context_len = next_token, T
+
         if is_speculative:
             self.prefill(
                 draft_model,
@@ -636,7 +644,7 @@ def generate(
         # max_new_tokens <= 2 means we are effectively not calling decode_n_tokens().
         callback(next_token.clone().view(-1), done_generating=max_new_tokens <= 2)
 
-        input_pos = torch.tensor([start_pos + T], device=device, dtype=torch.int)
+        input_pos = torch.tensor([start_pos + context_len], device=device, dtype=torch.int)
         accept_counts = [0] * (
             speculate_k + 1
         )  # creates array of [0, 0, 0, ...] that is speculate_k + 1 long
@@ -729,31 +737,54 @@ def chat(
         print("Builder Args:")
         print(self.builder_args)
 
-        exit(0)
-
         if generator_args.image_prompts is not None:
             print("Image prompts", generator_args.image_prompts)
-
             # Support for just the first image prompt for now
             images = [Image.open(generator_args.image_prompts[0])]
-            messages = [
-                Message(
-                    role="user",
-                    content=[
-                        {"type": "image", "content": images[0]},
-                        {"type": "text", "content": generator_args.prompt},
-                    ],
-                    eot=True,
-                ),
-                Message(role="assistant", content=""),
-            ]
 
-            transform = flamingo_transform(str(self.tokenizer_args.tokenizer_path))
-            data = transform({"messages": messages}, inference=True)
-            batch = padded_collate([data], self.builder_args.device)
-            batch.pop("mask")
-            encoded = batch["tokens"]
+            assert len(images) == 1, "Only one image prompt is supported for now"
+
+            #TODO: updated encoded variable for multi-modality models to include image tokens.
+            if self.model.config.model_type == ModelType.Flamingo:
+                messages = [
+                    Message(
+                        role="user",
+                        content=[
+                            {"type": "image", "content": images[0]},
+                            {"type": "text", "content": generator_args.prompt},
+                        ],
+                        eot=True,
+                    ),
+                    Message(role="assistant", content=""),
+                ]
 
+                transform = flamingo_transform(str(self.tokenizer_args.tokenizer_path))
+                data = transform({"messages": messages}, inference=True)
+                batch = padded_collate([data], self.builder_args.device)
+                batch.pop("mask")
+                encoded = batch["tokens"]
+            elif self.model.config.model_type == ModelType.Llava:
+                #TODO: double check the tokenizer.
+                def find_subtensor(tensor, target):
+                    target_len = len(target)
+                    for i in range(len(tensor) - target_len + 1):
+                        if torch.all(tensor[i:i+target_len] == target):
+                            return i
+                    return -1
+
+                input_ids = self.encode_tokens(generator_args.prompt, bos=True, device=self.builder_args.device)
+                image_token_indices = self.encode_tokens("<image>", device=self.builder_args.device)[1:]
+                index = find_subtensor(input_ids, image_token_indices)
+
+                batch = {
+                    "tokens": input_ids[:index].unsqueeze(0),
+                    "encoder_input": llava_image_preprocess(images[0], device=self.builder_args.device),
+                    "post_tokens": input_ids[index + len(image_token_indices) :].unsqueeze(0),
+                }
+                print("BATTTTTTTCHCHHHHHHHHH")
+                print(batch)
+                encoded = torch.cat([batch["tokens"].view(1, -1), batch["post_tokens"].view(1, -1)], dim=-1).view(-1)
+                
         else:
             encoded = self.encode_tokens(
                 generator_args.prompt, bos=True, device=self.builder_args.device