feat: Add model parameter management and UI controls

Add support for reading generation_config.json for model parameters
Implement model_loaded_callback for parameter updates
Add UI controls to switch between model and app defaults
Add visual highlighting for modified parameters
Add new API endpoints for parameter management
Improve session handling of model parameters
Update UI styling for parameter controls

Author: LlamaEnjoyer

backend/models.py

@@ -24,6 +24,26 @@
 from backend.config import config_filename
 from backend.util import *
 
+from typing import Callable, Optional, Dict, Any
+
+# Callback type for model parameter updates
+ModelLoadedCallback = Callable[[Dict[str, Any]], None]
+
+# Global callback that will be called when model parameters are loaded/updated
+model_loaded_callback: Optional[ModelLoadedCallback] = None
+
+def set_model_loaded_callback(callback: Optional[ModelLoadedCallback]) -> None:
+    """Set callback to be notified when model parameters are loaded/updated.
+    
+    Args:
+        callback: Function that takes model dict as argument, or None to clear
+    """
+    global model_loaded_callback
+    if callback is not None and not callable(callback):
+        raise TypeError("Model loaded callback must be callable")
+    model_loaded_callback = callback
+
+# Reserve memory for auto-split functionality
 auto_split_reserve_bytes = 512 * 1024**2
 
 models = {}
@@ -158,7 +178,59 @@ def prepare_draft_model(model):
         if "draft_rope_alpha_auto" not in model: model["draft_rope_alpha_auto"] = True
 
 
-def prepare_model(model):
+def prepare_model(model: Dict[str, Any]) -> None:
+    """Prepare model for loading by configuring parameters and resources.
+    
+    Args:
+        model: Dictionary containing model configuration
+        
+    Raises:
+        ValueError: If model directory is invalid
+        JSONDecodeError: If generation_config.json exists but is malformed
+    """
+    # Read generation_config.json if present
+    config_path = os.path.join(expanduser(model["model_directory"]), "generation_config.json")
+    if os.path.exists(config_path):
+        try:
+            with open(config_path, encoding='utf-8') as f:
+                gen_config = json.load(f)
+
+            if not isinstance(gen_config, dict):
+                raise ValueError("generation_config.json must contain a JSON object")
+                
+            print(f"Found generation_config.json: {gen_config}")
+
+            # Map generation config parameters to internal names
+            params_to_check = {
+                "temperature": "temperature",
+                "top_k": "top_k",
+                "top_p": "top_p",
+                "repetition_penalty": "repp"
+            }
+
+            # Store original values for logging
+            orig_values = {k: model.get(k) for k in params_to_check.values()}
+            
+            # Update model with values from generation_config.json
+            for config_name, internal_name in params_to_check.items():
+                if config_name in gen_config:
+                    # Validate parameter types
+                    value = gen_config[config_name]
+                    if not isinstance(value, (int, float)):
+                        print(f"Warning: Invalid type for {config_name} in generation_config.json. Expected number, got {type(value)}")
+                        continue
+                        
+                    model[internal_name] = value
+                    print(f"Setting {internal_name} from {orig_values.get(internal_name)} to {value}")
+
+            # Save updated model config
+            save_models()
+        except json.JSONDecodeError as e:
+            print(f"Error parsing generation_config.json: {e}")
+            print("Using default parameter values")
+        except Exception as e:
+            print(f"Unexpected error reading generation_config.json: {e}")
+            print("Using default parameter values")
 
     prep_config = ExLlamaV2Config()
     prep_config.fasttensors = False
@@ -194,6 +266,14 @@ def prepare_model(model):
     if "gpu_split" not in model: model["gpu_split"] = ""
     if "gpu_split_auto" not in model: model["gpu_split_auto"] = True
 
+    # Log final parameter state
+    print("Final model parameters:", {
+        "temperature": model.get("temperature", 0.8),
+        "top_k": model.get("top_k", 50),
+        "top_p": model.get("top_p", 0.8),
+        "repp": model.get("repp", 1.01)
+    })
+
 
 class ModelContainer:
 
@@ -413,6 +493,16 @@ def load_model(data):
         yield json.dumps(result) + "\n"
         return ""
 
+    # Notify about model load via callback
+    if success and model_loaded_callback is not None:
+        print("Calling model_loaded_callback with params:", {
+            "temperature": model.get("temperature", 0.8),
+            "top_k": model.get("top_k", 50),
+            "top_p": model.get("top_p", 0.8),
+            "repp": model.get("repp", 1.01)
+        })
+        model_loaded_callback(model)
+
     result = { "result": "ok" }
     # print(json.dumps(result) + "\n")
     yield json.dumps(result) + "\n"
@@ -430,4 +520,3 @@ def unload_model():
 
     result = { "result": "ok" }
     return result
-

backend/sessions.py

@@ -17,14 +17,22 @@
 )
 
 from backend.config import set_config_dir, global_state, config_filename
-from backend.models import get_loaded_model
+from backend.models import set_model_loaded_callback
 from backend.prompts import prompt_formats
 from backend.util import MultiTimer
+import backend.models as models  # Import as module to avoid circular dependency
 import threading
 
 session_list: dict or None = None
 current_session = None
 
+def handle_model_loaded(model):
+    """Handle model loading - only update new sessions with model params"""
+    pass
+
+# Register callback to handle model loading
+set_model_loaded_callback(handle_model_loaded)
+
 # Cancel
 
 abort_event = threading.Event()
@@ -92,9 +100,14 @@ def delete_session(d_session):
         current_session = None
 
 
-def get_default_session_settings():
-    return \
-    {
+def get_default_session_settings(use_model_params=False):
+    """Get default session settings
+    
+    Args:
+        use_model_params: If True and a model is loaded with custom params,
+                         apply those params instead of defaults
+    """
+    settings = {
         "prompt_format": "Chat-RP",
         "roles": [ "User", "Assistant", "", "", "", "", "", "" ],
         "system_prompt_default": True,
@@ -119,6 +132,23 @@ def get_default_session_settings():
         "temperature_last": False,
         "skew": 0.0,
     }
+    
+    if use_model_params:
+        # If requested, try to use model parameters
+        loaded_model = models.get_loaded_model()
+        if loaded_model is not None:
+            model_dict = loaded_model.model_dict
+            # Only apply if model has custom params defined
+            if any(param in model_dict for param in ["temperature", "top_k", "top_p", "repp"]):
+                settings.update({
+                    "temperature": model_dict.get("temperature", settings["temperature"]),
+                    "top_k": model_dict.get("top_k", settings["top_k"]),
+                    "top_p": model_dict.get("top_p", settings["top_p"]),
+                    "repp": model_dict.get("repp", settings["repp"])
+                })
+                print("Updated settings with model params:", settings)
+    
+    return settings
 
 class Session:
 
@@ -145,7 +175,8 @@ def init_new(self):
         self.session_uuid = str(uuid.uuid4())
         self.history = []
         # self.mode = ""
-        self.settings = get_default_session_settings()
+        # New sessions get app defaults
+        self.settings = get_default_session_settings(use_model_params=False)
 
 
     def to_json(self):
@@ -163,9 +194,13 @@ def from_json(self, j):
         self.session_uuid = j["session_uuid"]
         self.history = j["history"]
         # self.mode = j["mode"]
-        settings = get_default_session_settings()
-        if "settings" in j: settings.update(j["settings"])
-        self.settings = settings
+        
+        # Start with hardcoded defaults (no model params)
+        self.settings = get_default_session_settings(use_model_params=False)
+        
+        # Apply ALL saved settings including sampling params
+        if "settings" in j:
+            self.settings.update(j["settings"])
 
 
     def load(self):
@@ -244,7 +279,7 @@ def create_context(self, prompt_format, max_len, min_len, uptoblock = None, pref
 
     def create_context_instruct(self, prompt_format, max_len, min_len, uptoblock = None, prefix = ""):
 
-        tokenizer = get_loaded_model().tokenizer
+        tokenizer = models.get_loaded_model().tokenizer
         prompts = []
         responses = []
 
@@ -347,7 +382,7 @@ def create_context_instruct(self, prompt_format, max_len, min_len, uptoblock = N
 
     def create_context_raw(self, prompt_format, max_len, min_len, uptoblock = None, prefix=""):
 
-        tokenizer = get_loaded_model().tokenizer
+        tokenizer = models.get_loaded_model().tokenizer
         history_copy = []
         for h in self.history:
             if h["block_uuid"] == uptoblock: break
@@ -413,16 +448,17 @@ def generate(self, data):
         gen_prefix = data.get("prefix", "")
         block_id = data.get("block_id", None)
 
-        if get_loaded_model() is None:
+        if models.get_loaded_model() is None:
             packet = { "result": "fail", "error": "No model loaded." }
             yield json.dumps(packet) + "\n"
             return packet
 
-        model = get_loaded_model().model
-        generator = get_loaded_model().generator
-        tokenizer = get_loaded_model().tokenizer
-        cache = get_loaded_model().cache
-        speculative_mode = get_loaded_model().speculative_mode
+        loaded_model = models.get_loaded_model()
+        model = loaded_model.model
+        generator = loaded_model.generator
+        tokenizer = loaded_model.tokenizer
+        cache = loaded_model.cache
+        speculative_mode = loaded_model.speculative_mode
 
         prompt_format = prompt_formats[self.settings["prompt_format"]]()
 

server.py

@@ -125,7 +125,7 @@ def api_get_default_settings():
     if verbose: print("/api/get_default_settings")
     with api_lock:
         result = { "result": "ok",
-                   "session_settings": get_default_session_settings(),
+                   "session_settings": get_default_session_settings(use_model_params=False),  # Use hardcoded defaults
                    "notepad_settings": get_default_notepad_settings(),
                    "prompt_formats": list_prompt_formats() }
         return json.dumps(result) + "\n"
@@ -439,6 +439,96 @@ def api_cancel_notepad_generate():
         if verbose: print("->", result)
         return result
 
+@app.route("/api/get_model_params")
+def api_get_model_params():
+    global api_lock, verbose
+    if verbose: print("/api/get_model_params")
+    with api_lock:
+        model = get_loaded_model()
+        if model is None:
+            result = { "has_params": False }
+        else:
+            # Check if model has any sampling params defined
+            model_dict = model.model_dict
+            # Track which parameters are defined in the model
+            model_params = {
+                "temperature": "temperature" in model_dict,
+                "top_k": "top_k" in model_dict,
+                "top_p": "top_p" in model_dict,
+                "repp": "repp" in model_dict
+            }
+            has_params = any(model_params.values())
+            result = { 
+                "has_params": has_params,
+                "model_params": model_params
+            }
+        if verbose: print("->", result)
+        return json.dumps(result) + "\n"
+
+@app.route("/api/reset_to_app_defaults", methods=['POST'])
+def api_reset_to_app_defaults():
+    global api_lock, verbose
+    if verbose: print("/api/reset_to_app_defaults")
+    with api_lock:
+        session = get_session()
+        if session is not None:
+            # Get default settings
+            default_settings = get_default_session_settings(use_model_params=False)
+            
+            # Define which parameters are sampling-related
+            sampling_params = [
+                "temperature", "top_k", "top_p", "min_p", "tfs",
+                "mirostat", "mirostat_tau", "mirostat_eta", "typical",
+                "repp", "repr", "repd", "quad_sampling", "temperature_last", "skew"
+            ]
+            
+            # Reset only sampling parameters to defaults
+            updated_params = {}
+            for param in sampling_params:
+                updated_params[param] = default_settings[param]
+                session.settings[param] = default_settings[param]
+            
+            session.save()
+            result = { "result": "ok", "settings": updated_params }
+        else:
+            result = { "result": "fail", "error": "No session loaded" }
+        if verbose: print("->", result)
+        return json.dumps(result) + "\n"
+
+@app.route("/api/apply_model_params", methods=['POST'])
+def api_apply_model_params():
+    global api_lock, verbose
+    if verbose: print("/api/apply_model_params")
+    with api_lock:
+        model = get_loaded_model()
+        session = get_session()
+        if model is not None and session is not None:
+            # Get model's defined parameters
+            model_dict = model.model_dict
+            updated_params = {}
+            
+            # Only update parameters that are defined in the model
+            if "temperature" in model_dict:
+                updated_params["temperature"] = model_dict["temperature"]
+                session.settings["temperature"] = model_dict["temperature"]
+            if "top_k" in model_dict:
+                updated_params["top_k"] = model_dict["top_k"]
+                session.settings["top_k"] = model_dict["top_k"]
+            if "top_p" in model_dict:
+                updated_params["top_p"] = model_dict["top_p"]
+                session.settings["top_p"] = model_dict["top_p"]
+            if "repp" in model_dict:
+                updated_params["repp"] = model_dict["repp"]
+                session.settings["repp"] = model_dict["repp"]
+            
+            session.save()
+            # Only return the sampling parameters that were changed
+            result = { "result": "ok", "settings": updated_params }
+        else:
+            result = { "result": "fail", "error": "No model or session loaded" }
+        if verbose: print("->", result)
+        return json.dumps(result) + "\n"
+
 
 # Prepare torch
 
@@ -467,4 +557,3 @@ def api_cancel_notepad_generate():
     print(f" -- Opening UI in default web browser")
 
 serve(app, host = host, port = port, threads = 8)
-

static/chat.css

@@ -1,4 +1,3 @@
-
 .session-list {
     background-color: var(--background-color-body);
     display: flex;
@@ -492,7 +491,6 @@
     display: block;
 }
 
-
 .save-btn {
     background-color: var(--button-background);
     font-size: var(--font-size-small);
@@ -564,3 +562,7 @@
 .chat-popup action:last-child {
     margin-bottom: 0;
 }
+
+.highlight {
+    font-weight: bold;
+}