Use context length from config

jncraton · Aug 3, 2024 · e1c81e9 · e1c81e9
1 parent e293c98
commit e1c81e9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/languagemodels/models.py b/languagemodels/models.py
@@ -103,7 +103,9 @@ def get_model(model_type, tokenizer_only=False):
                 repo_id=model_config["repo"], filename=model_config["file"]
             )
 
-            model = Llama(model_path=model_path, verbose=False)
+            n_ctx = model_config.get("context_length", 512)
+
+            model = Llama(model_path=model_path, verbose=True, n_ctx=n_ctx, n_batch=64)
             modelcache[model_name] = (None, model)
         else:
             tokenizer = initialize_tokenizer(model_type, model_name)