Refactor comments and clean code

nikileshx · nikileshx · commit 9b4c92af9eae · 2025-08-29T04:23:30.000Z
diff --git a/models/experimental/mistral_24b/tests/pipeline_tests/test_end2end.py b/models/experimental/mistral_24b/tests/pipeline_tests/test_end2end.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
 """Test for Mistral-24B End-to-End Vision-Text Pipeline"""
 
 import torch
@@ -157,7 +160,6 @@ def process_real_vision_inputs(messages, model_args):
     )
 
     image_inputs, video_inputs = process_vision_info(messages)
-    # image_inputs, video_inputs = None, None
 
     encoded = processor(
         text=[text], images=image_inputs, videos=video_inputs, return_tensors="pt", return_dict=True
@@ -284,7 +286,7 @@ def run_generation_exactly_like_test_end2end(
 
     topk_tokens = [model_args.tokenizer.decode([idx.item()]) for idx in topk_indices]
 
-    logger.info("🔍 Top-5 predicted tokens (with probabilities):")
+    logger.info("Top-5 predicted tokens (with probabilities):")
     for i in range(top_k):
         logger.info(f"{i+1}. Token: '{topk_tokens[i]}' (ID={topk_indices[i].item()}), P={topk_probs[i].item():.4f}")
 
@@ -347,17 +349,17 @@ def run_generation_exactly_like_test_end2end(
 
         # Final response (exactly like test_end2end.py)
         response = model_args.tokenizer.decode(all_outputs[0], skip_special_tokens=True)
-        logger.info(f"📝 Each iteration Generated Response:\n{response}")
-        logger.info(f"📝 Each iteration Generated {len(all_outputs[0])} tokens: {all_outputs[0]}")
+        logger.info(f"Each iteration Generated Response:\n{response}")
+        logger.info(f"Each iteration Generated {len(all_outputs[0])} tokens: {all_outputs[0]}")
         chat = parse_chat_output(response)
         display_chat(logger, chat)
 
         logger.info(f" Each iteration Generated {len(results)} tokens successfully")
 
     # Final response (exactly like test_end2end.py)
     response = model_args.tokenizer.decode(all_outputs[0], skip_special_tokens=True)
-    logger.info(f"📝 Final Generated Response:\n{response}")
-    logger.info(f"📝 Generated {len(all_outputs[0])} tokens: {all_outputs[0]}")
+    logger.info(f"Final Generated Response:\n{response}")
+    logger.info(f"Generated {len(all_outputs[0])} tokens: {all_outputs[0]}")
     chat = parse_chat_output(response)
     display_chat(logger, chat)
 
@@ -469,9 +471,6 @@ def test_e2e_vision_text_pipeline(
     # Setup vision prompts and tokenizer
     messages, tokenizer = setup_vision_prompts_and_tokenizer(model_args, instruct)
 
-    # logger.info("Running reference HF vision-text model using messages..... ")
-    # hf_output = run_reference_demo_pipeline(messages)
-
     # Process real vision inputs from images
     processed_inputs = process_real_vision_inputs(messages, model_args)
 
@@ -524,12 +523,12 @@ def test_e2e_vision_text_pipeline(
 
     # Final validation
     if validation_passed and len(results) > 0:
-        logger.info("✅ E2E vision-text pipeline test PASSED!")
+        logger.info("E2E vision-text pipeline test PASSED!")
         logger.info(f"Successfully generated {len(results)} tokens")
 
         # Log generated tokens for debugging
         for i, result in enumerate(results[:5]):
             logger.info(f"Token {i}: {result.token} -> '{result.text}'")
     else:
-        logger.error("❌ E2E pipeline test failed")
+        logger.error("E2E pipeline test failed")
         assert False, f"E2E pipeline failed - generated {len(results)} tokens, validation: {validation_passed}"
diff --git a/models/experimental/mistral_24b/tests/pipeline_tests/test_vision_model.py b/models/experimental/mistral_24b/tests/pipeline_tests/test_vision_model.py
@@ -1,6 +1,10 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 # SPDX-License-Identifier: Apache-2.0
 
+"""
+This file is a unit test for validating the Mistral-24B Vision Model pipeline.
+"""
+
 import os
 import pytest
 import torch
diff --git a/models/experimental/mistral_24b/tests/pipeline_tests/test_vision_tower.py b/models/experimental/mistral_24b/tests/pipeline_tests/test_vision_tower.py
@@ -1,6 +1,10 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 # SPDX-License-Identifier: Apache-2.0
 
+"""
+This file is a unit test for validating the Mistral-24B Vision Tower model.
+"""
+
 import os
 import pytest
 import torch
diff --git a/models/experimental/mistral_24b/tests/test_conv2d.py b/models/experimental/mistral_24b/tests/test_conv2d.py
@@ -1,7 +1,9 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
-
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 # SPDX-License-Identifier: Apache-2.0
 
+"""
+This file is a unit test for validating the Mistral-24B conv2d.
+"""
 import os
 
 import pytest
diff --git a/models/experimental/mistral_24b/tests/test_patch_rot_emb.py b/models/experimental/mistral_24b/tests/test_patch_rot_emb.py
@@ -1,11 +1,13 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
 from loguru import logger
 
 import torch
 import pytest
 import os
 import ttnn
 
-# models/tt_transformers/tt/common.py
 from models.experimental.mistral_24b.tt.vision_rope import VisionRotarySetup as RotarySetup
 
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
diff --git a/models/experimental/mistral_24b/tests/test_pixtral_transformer.py b/models/experimental/mistral_24b/tests/test_pixtral_transformer.py
@@ -1,6 +1,6 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
-
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 # SPDX-License-Identifier: Apache-2.0
+
 import os
 
 import pytest
diff --git a/models/experimental/mistral_24b/tests/test_vision_attention.py b/models/experimental/mistral_24b/tests/test_vision_attention.py
@@ -1,5 +1,4 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
-
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 # SPDX-License-Identifier: Apache-2.0
 
 import os
@@ -81,11 +80,6 @@ def test_vision_attention(mesh_device, seq_len, batch_size):
     cos = torch.ones((1, T, head_dim)).to(torch.bfloat16)
     sin = torch.zeros((1, T, head_dim)).to(torch.bfloat16)
 
-    # attention_input = model_args.prepare_residual_tensor_prefill(
-    #     pt_attention_input,
-    #     force_replicated=True,
-    # )
-
     attention_input = ttnn.from_torch(
         pt_attention_input.unsqueeze(0),
         device=mesh_device,
diff --git a/models/experimental/mistral_24b/tests/test_vision_mlp.py b/models/experimental/mistral_24b/tests/test_vision_mlp.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 
 # SPDX-License-Identifier: Apache-2.0
 
diff --git a/models/experimental/mistral_24b/tests/test_vision_rms.py b/models/experimental/mistral_24b/tests/test_vision_rms.py
@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
 from loguru import logger
 
 import torch
diff --git a/models/experimental/mistral_24b/tt/model.py b/models/experimental/mistral_24b/tt/model.py
@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
 """
 This is the end-to-end pipeline for the Mistral-Small-3.1-24B-Instruct-2503 model.
 
@@ -51,9 +55,7 @@ def prepare_inputs_prefill(self, tokens, start_pos=0, page_table=None, chunk_pag
             layout=ttnn.ROW_MAJOR_LAYOUT,
             mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
         )
-        # self.embed_scale = args.dim**0.5
         tokens_embd = self.embd(tokens)
-        # tokens_embd = ttnn.multiply(tokens_embd, self.embed_scale)
 
         pixel_values = kwargs["processed_inputs"]["pixel_values"]
         input_ids = kwargs["processed_inputs"]["input_ids"]
@@ -73,7 +75,6 @@ def prepare_inputs_prefill(self, tokens, start_pos=0, page_table=None, chunk_pag
             input_ids = torch.nn.functional.pad(
                 input_ids, (0, tokens_embd.shape[1] - input_ids.shape[1]), "constant", 0
             )
-            # image_features = image_features.squeeze(0)
             special_image_mask = (input_ids == 10).unsqueeze(-1)
             special_image_mask = special_image_mask.expand_as(tokens_embd)
             image_features = image_features.to(tokens_embd.device, tokens_embd.dtype)
diff --git a/models/experimental/mistral_24b/tt/pipeline/mistral_vision_tower.py b/models/experimental/mistral_24b/tt/pipeline/mistral_vision_tower.py
@@ -2,6 +2,10 @@
 
 # SPDX-License-Identifier: Apache-2.0
 
+"""
+This file implements the Vision Tower submodule specific for the Mistral-Small-3.1-24B-Instruct-2503 model.
+"""
+
 import ttnn
 from models.common.lightweightmodule import LightweightModule
 from models.experimental.mistral_24b.tt.vision_conv2d import TtMistralConv2dPatch
@@ -139,7 +143,7 @@ def forward(self, input_tensor, image_sizes=None):
         patch_embeds = ttnn.concat(reshaped_patches, dim=0)
 
         # ln_pre RMS Norm
-        mode = "prefill"  # if self.max_seq_len <= 32 else "prefill"
+        mode = "prefill"
         patch_embeds = self.ln_pre(patch_embeds, mode=mode)
 
         # # positional embeddings
diff --git a/models/experimental/mistral_24b/tt/pipeline/vision_model.py b/models/experimental/mistral_24b/tt/pipeline/vision_model.py
@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
 """
 This is the end-to-end architecture of the Mistral-24B vision model.
 
diff --git a/models/experimental/mistral_24b/tt/vision_attention.py b/models/experimental/mistral_24b/tt/vision_attention.py
@@ -1,7 +1,10 @@
 # SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
 
 # SPDX-License-Identifier: Apache-2.0
+"""
+This file implements the vision attention submodule specific for the Mistral-Small-3.1-24B-Instruct-2503 model.
 
+"""
 import torch
 
 import ttnn
@@ -139,7 +142,6 @@ def pad_head_dim(weight, heads_out=True):
             dtype=self.dtype,
             memory_config=ttnn.DRAM_MEMORY_CONFIG,
             layout=ttnn.TILE_LAYOUT,
-            # cache_file_name=cache_name("wqkv_sharded"),
         )
 
         self.wo = ttnn.as_tensor(
@@ -153,7 +155,6 @@ def pad_head_dim(weight, heads_out=True):
             memory_config=ttnn.DRAM_MEMORY_CONFIG,
             dtype=self.dtype,
             layout=ttnn.TILE_LAYOUT,
-            # cache_file_name=cache_name("wo_sharded"),
         )
 
         self.scale = self.head_dim**-0.5
diff --git a/models/experimental/mistral_24b/tt/vision_conv2d.py b/models/experimental/mistral_24b/tt/vision_conv2d.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 
 # SPDX-License-Identifier: Apache-2.0
 
@@ -10,7 +10,6 @@
 
 class TtMistralConv2dPatch(LightweightModule):
     """Conv2D Patching layer.
-    Column parallel over unfolded input.
     Arguments:
         in_channels: Input channels.
         out_channels: Output channels.
@@ -61,10 +60,6 @@ def __init__(
         weight = state_dict[f"{state_dict_prefix}weight"]
         if weight.ndim == 4:
             weight = weight.reshape(out_channels, -1).T
-        # pad_len = nearest_32(weight.shape[-1]) - weight.shape[-1]
-        # padding = torch.zeros(self.out_channels, pad_len, dtype=weight.dtype)
-        # padded_weight = torch.cat([weight, padding], dim=-1)
-        # padded_weight = padded_weight.permute(1, 0).reshape(1, 1, -1, self.out_channels)
 
         self._linear_weight = ttnn.as_tensor(
             weight,
@@ -87,11 +82,6 @@ def forward(self, x: torch.Tensor):
         x = self._unfold(x)
         x = x.permute(0, 2, 1)
 
-        # Need to pad the last dimension of x to be a multiple of a tile
-        # pad_len = nearest_32(x.shape[-1]) - x.shape[-1]
-        # padding = torch.zeros((x.shape[0], x.shape[1], pad_len), dtype=x.dtype, device=x.device)
-        # x = torch.cat([x, padding], dim=-1)
-
         x = ttnn.as_tensor(
             x,
             dtype=ttnn.bfloat16,
diff --git a/models/experimental/mistral_24b/tt/vision_mlp.py b/models/experimental/mistral_24b/tt/vision_mlp.py
@@ -1,7 +1,10 @@
 # SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 
 # SPDX-License-Identifier: Apache-2.0
+"""
+This file implements the Vision FeedForward submodule specific for the Mistral-Small-3.1-24B-Instruct-2503 model.
 
+"""
 import torch
 
 import ttnn
@@ -75,9 +78,6 @@ def forward(self, x: ttnn.Tensor) -> ttnn.Tensor:
             w2 -> down_proj
         """
 
-        # if x.shape[-2] >= self.args.prefill_len_cutoff and mode != "decode":
-        #     x = ttnn.reshape(x, [1, x.shape[-2] // self.args.prefill_len_cutoff, self.args.prefill_len_cutoff, -1])
-
         # Linear with SILU activation
         w1_out = ttnn.linear(
             x,
diff --git a/models/experimental/mistral_24b/tt/vision_mmp.py b/models/experimental/mistral_24b/tt/vision_mmp.py
@@ -8,6 +8,10 @@
 import ttnn
 from ttnn import ConcatMeshToTensor
 
+"""
+This file implements the Vision pixtral image submodule specific for the Mistral-Small-3.1-24B-Instruct-2503 model.
+"""
+
 
 class TTMistral3PatchMerger(LightweightModule):
     def __init__(
@@ -25,11 +29,6 @@ def __init__(
         self.spatial_merge_size = 2  # TODO Handle in Model_config spatial_merge_size
         self.patch_size = args.vision_patch_size
         self.args = args
-        # self.patch_size = ttnn.from_torch(
-        #     torch.tensor(args.vision_patch_size, dtype=torch.int32),
-        #     device=mesh_device,
-        #     dtype=ttnn.int32
-        # )
 
         def get_weight(name):
             return torch.transpose(state_dict[f"{state_dict_prefix}{name}.weight"], -2, -1)
@@ -51,7 +50,6 @@ def as_tensor(name, dtype, is_bias=False):
                 mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
                 layout=ttnn.TILE_LAYOUT,
                 memory_config=ttnn.DRAM_MEMORY_CONFIG,
-                # cache_file_name=cache_name(name),
             )
 
         self.merging_weights = as_tensor("merging_layer", dtype)
@@ -147,7 +145,6 @@ def as_tensor(name, dtype, is_bias=False):
                 mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
                 layout=ttnn.TILE_LAYOUT,
                 memory_config=ttnn.DRAM_MEMORY_CONFIG,
-                # cache_file_name=cache_name(name),
             )
 
         self.linear_1_weight = as_tensor("linear_1", dtype)
diff --git a/models/experimental/mistral_24b/tt/vision_pixtral_image_block.py b/models/experimental/mistral_24b/tt/vision_pixtral_image_block.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 
 # SPDX-License-Identifier: Apache-2.0
 
@@ -9,6 +9,10 @@
 from models.experimental.mistral_24b.tt.vision_attention import TtMistralImageAttention as TtLlamaImageAttention
 from models.experimental.mistral_24b.tt.vision_mlp import MistralTTVisionMLP as MLP
 
+"""
+This file implements the Pixtral_image_block submodule specific for the Mistral-Small-3.1-24B-Instruct-2503 model.
+"""
+
 
 class TtPixtralImageTransformerBlock(LightweightModule):
     def __init__(
diff --git a/models/experimental/mistral_24b/tt/vision_pixtral_transformer.py b/models/experimental/mistral_24b/tt/vision_pixtral_transformer.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 
 # SPDX-License-Identifier: Apache-2.0
 
diff --git a/models/experimental/mistral_24b/tt/vision_rope.py b/models/experimental/mistral_24b/tt/vision_rope.py
@@ -1,17 +1,17 @@
-# SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 
 # SPDX-License-Identifier: Apache-2.0
 
 import torch
 
 import ttnn
 from models.common.lightweightmodule import LightweightModule
-from models.tt_transformers.tt.common import precompute_vision_freqs
+from models.tt_transformers.tt.common import precompute_mistral_vision_freqs
 from ttnn import ReplicateTensorToMesh
 
 
 def compute_gather_cos_sin(dhead, max_patches_per_side, theta, scale_factor, orig_context_len, position_ids):
-    cos, sin = precompute_vision_freqs(dhead, max_patches_per_side, theta, scale_factor, orig_context_len)
+    cos, sin = precompute_mistral_vision_freqs(dhead, max_patches_per_side, theta, scale_factor, orig_context_len)
     return cos, sin
 
 
diff --git a/models/tt_transformers/tt/common.py b/models/tt_transformers/tt/common.py
@@ -307,7 +307,7 @@ def apply_scaling_vision(freqs: torch.Tensor, scale_factor: float, orig_context_
     return freqs / scale_factor
 
 
-def precompute_vision_freqs(
+def precompute_mistral_vision_freqs(
     dim: int, max_patches_per_side: int, theta: float, scale_factor=None, orig_context_len=None
 ):
     # Compute base frequencies
diff --git a/models/tt_transformers/tt/model_config.py b/models/tt_transformers/tt/model_config.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.`
	`1`	`+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.`
`2`	`2`
`3`	`3`	`# SPDX-License-Identifier: Apache-2.0`
`4`	`4`