Refactor model_config

Author: MohammedTaherMcW

models/tt_transformers/tt/model_config.py

@@ -575,9 +575,10 @@ def __init__(
             max_prefill_chunk_size_div1024 = int(max_prefill_chunk_size_div1024)
         self.max_prefill_chunk_size = max_prefill_chunk_size_div1024 * 1024
 
-        if (self.base_model_name in ["Llama-3.1-8B", "Llama-3.2-11B", "Mistral-7B"] and self.device_name == "N150") or (
-            self.base_model_name in ["Qwen2.5-7B"] and self.device_name == "N300"
-        ):
+        if (
+            self.base_model_name in ["Llama-3.1-8B", "Llama-3.2-11B", "Mistral-7B", "gemma-3-1b-it"]
+            and self.device_name == "N150"
+        ) or (self.base_model_name in ["Qwen2.5-7B"] and self.device_name == "N300"):
             logger.info(f"Reducing prefill_len_cutoff to 512 for {self.model_name} on {self.device_name}")
             self.prefill_len_cutoff = 512
 
@@ -1396,7 +1397,9 @@ def _set_params_from_dict(self, config, is_hf=False):
         # Try to get text_config, if it doesn't exist everything is text config
         eos_token_id = config.get("eos_token_id", None)
 
-        self.eos_token_id = None if isinstance(eos_token_id, int) else eos_token_id
+        self.eos_token_id = (
+            None if isinstance(eos_token_id, int) else eos_token_id
+        )  # Gemma like models can have a list of eos token ids
 
         self.sliding_window_pattern = config.get("sliding_window_pattern", 1)
 
@@ -2187,7 +2190,7 @@ def reference_embedding(self, reference_model=None):
                 model = self.reference_transformer(wrap=False)
                 layer = model.model.embed_tokens
             else:
-                layer = reference_model.model.embed_tokens
+                layer = reference_model.model.model.embed_tokens
 
             layer._load_state_dict = layer.load_state_dict
             layer.load_state_dict = lambda x: layer._load_state_dict(convert_meta_to_hf(x, self.head_dim))