add n_batch support for llama.cpp (#1115)

2023-04-24 02:46:18 -04:00 · 2023-04-24 02:46:18 -04:00 · 78d1977ebf
commit 78d1977ebf
parent 2f6e2ddeac
3 changed files with 4 additions and 1 deletions
--- a/modules/llamacpp_model_alternative.py
+++ b/modules/llamacpp_model_alternative.py
@ -24,7 +24,8 @@ class LlamaCppModel:
            'model_path': str(path),
            'n_ctx': 2048,
            'seed': 0,
-            'n_threads': shared.args.threads or None
+            'n_threads': shared.args.threads or None,
+            'n_batch': shared.args.n_batch
        }
        self.model = Llama(**params)
        self.model.set_cache(LlamaCache)