make text_transformer_args a real attribute

Author: Gasoonjia

torchchat/model.py

@@ -304,6 +304,7 @@ def __init__(self, config: ModelArgs) -> None:
         super().__init__()
         self.config = config
         self.model = self.build_model()
+        self.text_transformer_args = None
 
     def build_model(self) -> nn.Module:
         """
@@ -331,11 +332,6 @@ def forward(self, *args, **kwargs):
     @abstractmethod
     def setup_caches(self, *args, **kwargs):
         raise NotImplementedError("setup_caches method is not implemented")
-    
-    @property
-    @abstractmethod
-    def text_transformer_args(self):
-        raise NotImplementedError("no text_transformer_args is created")
 
     @classmethod
     def _get_model_instance(cls, config: ModelArgs):
@@ -371,15 +367,15 @@ def from_gguf(cls, gguf_path: str, **kwargs):
 
 
 class TextOnlyModel(Model):
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__(config)
+        self.text_transformer_args = self.model.config
+
     def forward(self, tokens: Tensor, input_pos: Optional[Tensor] = None) -> Tensor:
         return self.model(tokens, input_pos)
 
     def setup_caches(self, max_batch_size, max_seq_length):
         self.model.setup_caches(max_batch_size, max_seq_length)
-    
-    @property
-    def text_transformer_args(self):
-        return self.model.model.config
 
 
 class Llama31Model(Model):
@@ -391,11 +387,6 @@ def setup_caches(self, max_batch_size, dtype):
 
     def reset_caches(self):
         self.model.reset_caches()
-    
-    @property
-    def text_transformer_args(self):
-        # TODO: add support for llama3_1
-        return None
 
 
 class FlamingoModel(Model):
@@ -416,11 +407,7 @@ def setup_caches(self, max_batch_size, dtype):
 
     def reset_caches(self):
         self.model.reset_caches()
-    
-    @property
-    def text_transformer_args(self):
-        # TODO: add support for flamingo
-        return None
+
 
 
 MODEL_TYPE_TO_CLASS = {
@@ -813,7 +800,7 @@ def __init__(self, config, path) -> None:
             self.config = config
             self.model_ = exec_lib._load_for_executorch(str(path))
 
-            self.text_transformer_config = TransformerArgs.from_params(self.config.transformer_args["text"])
+            self.text_transformer_args = TransformerArgs.from_params(self.config.transformer_args["text"])
 
         def forward(self, x, input_pos):
             # model_.forward expects inputs to be wrapped in a tuple

torchchat/utils/gguf_loader.py

@@ -544,15 +544,13 @@ def load_model(gguf_file: str) -> torch.nn.Module:
     model_args = ModelArgs(
         {
             "text": {
-                "config": {
-                    "dim": metadata[f"{arch}.embedding_length"],
-                    "n_layers": metadata[f"{arch}.block_count"],
-                    "n_heads": metadata[f"{arch}.attention.head_count"],
-                    "n_local_heads": metadata[f"{arch}.attention.head_count_kv"],
-                    "vocab_size": len(metadata["tokenizer.ggml.tokens"]),
-                    "norm_eps": metadata[f"{arch}.attention.layer_norm_rms_epsilon"],
-                    "hidden_dim": metadata[f"{arch}.feed_forward_length"],
-                }
+                "dim": metadata[f"{arch}.embedding_length"],
+                "n_layers": metadata[f"{arch}.block_count"],
+                "n_heads": metadata[f"{arch}.attention.head_count"],
+                "n_local_heads": metadata[f"{arch}.attention.head_count_kv"],
+                "vocab_size": len(metadata["tokenizer.ggml.tokens"]),
+                "norm_eps": metadata[f"{arch}.attention.layer_norm_rms_epsilon"],
+                "hidden_dim": metadata[f"{arch}.feed_forward_length"],
             }
         }
     )