Merge remote-tracking branch 'LlamaEnjoyer/read_sampling_params_from_model_config'

Author: turboderp

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"
@@ -458,6 +458,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
 

static/chat.css

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