Merge pull request #244 from tushar2407/kbit_bug

fix: fixed bug in loading kbit model

Author: StochasticRomanAgeev

src/xturing/models/causal.py

@@ -399,6 +399,19 @@ def __init__(
 
 
 class CausalLoraKbitModel(CausalLoraModel):
-    def __init__(self, engine: str, weights_path: Optional[str] = None):
+    def __init__(
+            self, 
+            engine: str, 
+            weights_path: Optional[str] = None,
+            model_name: Optional[str] = None,
+            target_modules: Optional[List[str]] = None,
+            **kwargs,
+        ):
         assert_not_cpu_int8()
-        super().__init__(engine, weights_path)
+        super().__init__(
+            engine,
+            weights_path=weights_path,
+            model_name=model_name,
+            target_modules=target_modules,
+            **kwargs,
+        )

src/xturing/models/generic.py

@@ -94,10 +94,12 @@ def __init__(
         model_name: str,
         target_modules: List[str] = ["c_attn"],
         weights_path: Optional[str] = None,
+        **kwargs,
     ):
         super().__init__(
             GenericLoraKbitEngine.config_name,
             weights_path,
             model_name=model_name,
             target_modules=target_modules,
+            **kwargs,
         )