Added Llama3VisionTransform in TokenizerArgs and other changes

Author: anirudh

torchchat/cli/builder.py

@@ -252,13 +252,29 @@ class TokenizerArgs:
     is_sentencepiece: bool = False
     is_tiktoken: bool = False
     is_hf_tokenizer: bool = False
+    is_llama_3_2_mm: bool = False
     t: Optional[Any] = None
 
     def __post_init__(self):
+        # special handling for llama-3.2-mm
+        if "llama-3.2-11b-vision" in str(self.tokenizer_path).lower():
+            try:
+                from torchtune.models.llama3_2_vision import llama3_2_vision_transform
+
+                self.t = llama3_2_vision_transform(path=str(self.tokenizer_path))
+                self.is_llama_3_2_mm = True
+                self.is_tiktoken = False
+                self.is_sentencepiece = False
+                self.is_hf_tokenizer = False
+                return
+            except:
+                pass
+
         try:
             from tokenizer.tiktoken import Tokenizer as TiktokenTokenizer
 
             self.t = TiktokenTokenizer(model_path=str(self.tokenizer_path))
+            self.is_llama_3_2_mm = False
             self.is_tiktoken = True
             self.is_sentencepiece = False
             self.is_hf_tokenizer = False
@@ -270,6 +286,7 @@ def __post_init__(self):
             from sentencepiece import SentencePieceProcessor
 
             self.t = SentencePieceProcessor(model_file=str(self.tokenizer_path))
+            self.is_llama_3_2_mm = False
             self.is_tiktoken = False
             self.is_sentencepiece = True
             self.is_hf_tokenizer = False
@@ -281,13 +298,15 @@ def __post_init__(self):
             from tokenizer.hf_tokenizer import HFTokenizer
 
             self.t = HFTokenizer(str(self.tokenizer_path))
+            self.is_llama_3_2_mm = False
             self.is_tiktoken = False
             self.is_sentencepiece = False
             self.is_hf_tokenizer = True
             return
         except:
             pass
 
+        self.is_llama_3_2_mm = False
         self.is_tiktoken = False
         self.is_sentencepiece = False
         self.is_hf_tokenizer = False
@@ -302,20 +321,22 @@ def validate_model(
         if model is None:
             return
 
-        if sum([self.is_tiktoken, self.is_hf_tokenizer, self.is_sentencepiece]) != 1:
+        if sum([self.is_tiktoken, self.is_hf_tokenizer, self.is_sentencepiece, self.is_llama_3_2_mm]) != 1:
             raise RuntimeError(f"no tokenizer was found at {self.tokenizer_path}")
 
         is_tiktoken = self.is_tiktoken
         is_sentencepiece = self.is_sentencepiece
         is_hf_tokenizer = self.is_hf_tokenizer
+        is_llama_3_2_mm = self.is_llama_3_2_mm
+
         use_tiktoken = model.config.use_tiktoken
         use_hf_tokenizer = model.config.use_hf_tokenizer
-        use_sentencepiece = not (use_tiktoken or use_hf_tokenizer)
-
+        use_other_tokenizer = not (use_tiktoken or use_hf_tokenizer)
         if (
             (is_tiktoken and not use_tiktoken) or
             (is_hf_tokenizer and not use_hf_tokenizer) or
-            (is_sentencepiece and not use_sentencepiece)
+            (is_sentencepiece and not use_other_tokenizer) or
+            (is_llama_3_2_mm and not use_other_tokenizer)
         ):
             raise RuntimeError(
                 "model-specified tokenizer ({}) does not match provided tokenizer ({}) for {}".format(

torchchat/model_params/Llama-3.2-11B-Vision.json

@@ -1,6 +1,6 @@
 {
     "model_type": "flamingo",
-    "use_tiktoken": true,
+    "use_tiktoken": false,
     "encoder": {
         "patch_size": 14,
         "num_heads": 16,

torchchat/usages/eval.py

@@ -378,16 +378,13 @@ def _model_multimodal_generate(
 
         # 2. Setup KV cache and masks for bsz 1
         with self.device:
-            if self.model.caches_are_enabled():
-                self.model.reset_caches()
-            else:
-                self.model.setup_caches(
-                    batch_size=1,
-                    dtype=self._dtype,
-                    encoder_max_seq_len=self.model_transform.image_seq_len
-                    * self._max_images_per_sample,
-                    decoder_max_seq_len=self.max_length,
-                )
+            self.model.setup_caches(
+                batch_size=1,
+                dtype=self._dtype,
+                encoder_max_seq_len=self.model_transform.image_seq_len
+                * self._max_images_per_sample,
+                decoder_max_seq_len=self.max_length,
+            )
             causal_mask = torch.tril(
                 torch.ones(
                     size=(self.max_length, self.max_length),
@@ -506,6 +503,8 @@ def multi_model_eval(
     """
     if tasks is None:
         tasks = ["wikitext"]
+    max_seq_length = 4096 if max_seq_length is None else max_seq_length
+    device = utils.get_device(device) if isinstance(device, str) else device
 
     model_eval_wrapper = _VLMEvalWrapper(
         model,
@@ -578,25 +577,30 @@ def main(args) -> None:
         )
         torch._inductor.config.coordinate_descent_tuning = False if device == "cpu" else True
 
-    evaluator = None
-    if modality == "text":
-        evaluator = eval
-    elif modality == "text-image":
-        evaluator = multi_model_eval
-    else:
-        raise ValueError(f"Unsupported modality: {modality}")
-
     with measure_time("Time to run eval: {time:.02f}s."):
-        result = evaluator(
-            model.to(device),
-            model_forward,
-            tokenizer,
-            tasks,
-            limit,
-            max_seq_length,
-            device=builder_args.device,
-            is_pte_model=builder_args.pte_path is not None,
-        )
+        if modality == "text":
+            result = eval(
+                model.to(device),
+                model_forward,
+                tokenizer,
+                tasks,
+                limit,
+                max_seq_length,
+                device=builder_args.device,
+                is_pte_model=builder_args.pte_path is not None,
+            )
+        elif modality == "text-image":
+            result = multi_model_eval(
+                model.to(device),
+                model_forward,
+                tokenizer,
+                tasks,
+                limit,
+                max_seq_length,
+                device=builder_args.device,
+            )
+        else:
+            raise ValueError(f"Unsupported modality: {modality}")
 
     times = torch.tensor(result["times"])
     print(