From 51ff962f140fc1622a2dc61d1a60828da6a5c6ec Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Wed, 22 Nov 2023 14:15:00 +0100
Subject: [PATCH] Add autosplit, temperature_last

---
 exllama.py | 36 +++++++++++++++++++++++-------------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/exllama.py b/exllama.py
index db07fb5..928a876 100644
--- a/exllama.py
+++ b/exllama.py
@@ -34,6 +34,7 @@ class Loader:
         self.generator = None
         self.gpu_split = None
         self.cache_8bit = False
+        self.cache_lazy = False
 
     def process(self, model_dir, gpu_split, cache_8bit, max_seq_len):
         self.unload()
@@ -41,7 +42,9 @@ class Loader:
         self.config.model_dir = model_dir
         self.config.prepare()
 
-        if gpu_split:
+        if gpu_split == "auto":
+            self.cache_lazy = True
+        elif gpu_split:
             self.gpu_split = [float(a) for a in gpu_split.split(",")]
 
         if max_seq_len:
@@ -57,12 +60,16 @@ class Loader:
             return
 
         self.base = ExLlamaV2(self.config)
-        self.base.load(gpu_split=self.gpu_split)
 
         if self.cache_8bit:
-            self.cache = ExLlamaV2Cache_8bit(self.base)
+            self.cache = ExLlamaV2Cache_8bit(self.base, lazy=self.cache_lazy)
         else:
-            self.cache = ExLlamaV2Cache(self.base)
+            self.cache = ExLlamaV2Cache(self.base, lazy=self.cache_lazy)
+
+        if self.cache_lazy:
+            self.base.load_autosplit(self.cache)
+        else:
+            self.base.load(gpu_split=self.gpu_split)
 
         self.tokenizer = ExLlamaV2Tokenizer(self.config)
         self.generator = ExLlamaV2StreamingGenerator(self.base, self.cache, self.tokenizer)
@@ -87,11 +94,12 @@ class Generator:
                 "single_line": ("BOOLEAN", {"default": False}),
                 "max_tokens": ("INT", {"default": 128, "max": 2**16}),
                 "temperature": ("FLOAT", {"default": 1, "max": 2, "step": 0.01}),
+                "temperature_last": ("BOOLEAN", {"default": True}),
                 "min_p": ("FLOAT", {"default": 0.1, "max": 1, "step": 0.01}),
-                "top_k": ("INT", {"max": 200}),
                 "top_p": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}),
-                "typical": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}),
-                "penalty": ("FLOAT", {"default": 1, "min": 1, "max": 2, "step": 0.01}),
+                "typical_p": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}),
+                "top_k": ("INT", {"max": 200}),
+                "rep_penalty": ("FLOAT", {"default": 1, "min": 1, "max": 2, "step": 0.01}),
                 "seed": ("INT", {"max": 2**64 - 1}),
                 "text": ("STRING", {"multiline": True}),
             },
@@ -113,11 +121,12 @@ class Generator:
         single_line,
         max_tokens,
         temperature,
+        temperature_last,
         min_p,
-        top_k,
         top_p,
-        typical,
-        penalty,
+        typical_p,
+        top_k,
+        rep_penalty,
         seed,
         text,
         info=None,
@@ -144,11 +153,12 @@ class Generator:
 
         settings = ExLlamaV2Sampler.Settings()
         settings.temperature = temperature
+        settings.temperature_last = temperature_last
         settings.min_p = min_p
-        settings.top_k = top_k
         settings.top_p = top_p
-        settings.typical = typical
-        settings.token_repetition_penalty = penalty
+        settings.typical = typical_p
+        settings.top_k = top_k
+        settings.token_repetition_penalty = rep_penalty
 
         model.generator.begin_stream(input, settings, token_healing=True)
         progress = ProgressBar(max_tokens)