From 58c6001be91eeb09aeeb315661aab26e9a0c6fcf Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Thu, 16 Nov 2023 07:09:40 -0800
Subject: [PATCH] Add missing exllamav2 samplers

---
 modules/exllamav2.py | 5 +++++
 modules/loaders.py   | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/modules/exllamav2.py b/modules/exllamav2.py
index e2bcfd1b..cb142af5 100644
--- a/modules/exllamav2.py
+++ b/modules/exllamav2.py
@@ -96,7 +96,12 @@ class Exllamav2Model:
         settings.temperature = state['temperature']
         settings.top_k = state['top_k']
         settings.top_p = state['top_p']
+        settings.min_p = state['min_p']
+        settings.tfs = state['tfs']
         settings.typical = state['typical_p']
+        settings.mirostat = state['mirostat_mode'] == 2
+        settings.mirostat_tau = state['mirostat_tau']
+        settings.mirostat_eta = state['mirostat_eta']
         settings.token_repetition_penalty = state['repetition_penalty']
         settings.token_repetition_range = -1 if state['repetition_penalty_range'] <= 0 else state['repetition_penalty_range']
         if state['ban_eos_token']:
diff --git a/modules/loaders.py b/modules/loaders.py
index 455ef96a..32d99c6f 100644
--- a/modules/loaders.py
+++ b/modules/loaders.py
@@ -235,11 +235,16 @@ loaders_samplers = {
     'ExLlamav2': {
         'temperature',
         'top_p',
+        'min_p',
         'top_k',
         'typical_p',
+        'tfs',
         'repetition_penalty',
         'repetition_penalty_range',
         'seed',
+        'mirostat_mode',
+        'mirostat_tau',
+        'mirostat_eta',
         'ban_eos_token',
         'add_bos_token',
         'custom_token_bans',