Added the ability to Modify the Context Length (#210)

2024-02-20 16:22:56 -08:00
parent ce3b261053
commit 9de9a46815
4 changed files with 16 additions and 2 deletions
@@ -11,14 +11,19 @@ class ModelConfig:
        path: str,
        trust_remote_code: bool = True,
        revision: Optional[str] = None,
+        context_length: Optional[int] = None,
    ) -> None:
        self.path = path
        self.trust_remote_code = trust_remote_code
        self.revision = revision
        self.hf_config = get_config(self.path, trust_remote_code, revision)
+        
+        if context_length is not None:
+            self.context_len = context_length
+        else:
+            self.context_len = get_context_length(self.hf_config)

        # Unify the config keys for hf_config
-        self.context_len = get_context_length(self.hf_config)
        self.head_dim = self.hf_config.hidden_size // self.hf_config.num_attention_heads
        self.num_attention_heads = self.hf_config.num_attention_heads
        self.num_key_value_heads = getattr(self.hf_config, "num_key_value_heads", None)