new factory func to produce Model from modelargs

Author: Gasoonjia

torchchat/generate.py

@@ -303,7 +303,6 @@ def sample(
         temperature: float = 1.0,
         top_k: Optional[int] = None,
     ):  
-        print(f"logits {logits}")
         if temperature == 0 and not need_probs:
             _, idx_next = torch.topk(logits[0, -1], k=1, dim=-1)
             return (idx_next, None)
@@ -323,21 +322,15 @@ def prefill(
         # logging.debug(f"x: {x}, input_pos: {input_pos}")
         width = x.size(1)
         assert input_pos.size(0) == width
-        print("x: ", x, "input_pos: ", input_pos, "width: ", width)
-        print("sequential_prefill: ", sequential_prefill)
 
         if sequential_prefill:
             for i in range(width):
-                print("i: ", i)
                 x_sliced, ip_sliced = x[:, i].view(-1, 1), input_pos[i].view(-1)
-                print("x_sliced: ", x_sliced, "ip_sliced: ", ip_sliced)
                 # logging.debug(f"<sliced> x: {x_sliced}, input_pos: {ip_sliced}")
                 logits = model(x_sliced, ip_sliced)  # (x[:, i], input_pos[i])
-                print("logits: ", logits)
         else:
             # input_pos: [B, S]
             logits = model(x, input_pos)
-            print("logits: ", logits)
             # print(f"logits {logits.shape}")
 
         # print(f"x: {x},\n  input_pos: {input_pos}\n")

torchchat/model.py

@@ -310,11 +310,11 @@ def build_model(self) -> nn.Module:
 
     @abstractmethod
     def forward(self, *args, **kwargs):
-        pass
+        raise NotImplementedError("forward method is not implemented")
 
     @abstractmethod
     def setup_caches(self, *args, **kwargs):
-        pass
+        raise NotImplementedError("setup_caches method is not implemented")
 
     @classmethod
     def _get_model_instance(cls, config: ModelArgs):
@@ -323,6 +323,10 @@ def _get_model_instance(cls, config: ModelArgs):
             raise ValueError("Unsupported model type:", str(config.model_type))
         return model_class(config)
 
+    @classmethod
+    def from_model_args(cls, config: ModelArgs):
+        return cls._get_model_instance(config)
+
     @classmethod
     def from_name(cls, name: str):
         return cls._get_model_instance(ModelArgs.from_name(name))

torchchat/utils/gguf_loader.py

@@ -558,7 +558,7 @@ def load_model(gguf_file: str) -> torch.nn.Module:
     # metadata.get(f"{arch}.rope.dimension_count", None)
 
     with torch.device("meta"):
-        model = Model(model_args)
+        model = Model.from_model_args(model_args)
     return model