add model gen params

app.py

@@ -70,54 +70,69 @@ SUPPORTED_MODELS = [
 
 
 def fetch_model_generation_params(model_name: str) -> dict:
-    """Fetch generation parameters from model's generation config on the hub"""
+    """Fetch generation parameters and model config from the hub"""
     default_params = {
         "max_tokens": 1024,
         "temperature": 0.7,
         "top_k": 50,
-        "top_p": 0.95
+        "top_p": 0.95,
+        "max_position_embeddings": 2048,
+        "recommended_max_tokens": 1024
     }
 
     try:
-        print(f"Attempting to fetch generation config for: {model_name}")
+        print(f"Attempting to fetch configs for: {model_name}")
+
+        # Always try to load the model config first for max_position_embeddings
+        model_config = None
+        max_position_embeddings = default_params["max_position_embeddings"]
+
+        try:
+            output_dataset_token = os.getenv("OUTPUT_DATASET_TOKEN")
+            model_config = AutoConfig.from_pretrained(model_name, force_download=False, token=output_dataset_token)
+            max_position_embeddings = getattr(model_config, 'max_position_embeddings', default_params["max_position_embeddings"])
+            print(f"Loaded AutoConfig for {model_name}, max_position_embeddings: {max_position_embeddings}")
+        except Exception as e:
+            print(f"Failed to load AutoConfig for {model_name}: {e}")
+
+        # Calculate recommended max tokens (conservative estimate)
+        # Leave some room for the prompt, so use ~75% of max_position_embeddings
+        recommended_max_tokens = min(int(max_position_embeddings * 0.75), MAX_TOKENS)
+        recommended_max_tokens = max(256, recommended_max_tokens)  # Ensure minimum
 
         # Try to load the generation config
+        gen_config = None
         try:
-            gen_config = GenerationConfig.from_pretrained(model_name, force_download=False)
+            gen_config = GenerationConfig.from_pretrained(model_name, force_download=False, token=output_dataset_token)
             print(f"Successfully loaded generation config for {model_name}")
-            print(f"Config attributes: {dir(gen_config)}")
         except Exception as e:
             print(f"Failed to load GenerationConfig for {model_name}: {e}")
-            # Try loading from model config instead
-            try:
-                from transformers import AutoConfig
-                model_config = AutoConfig.from_pretrained(model_name, force_download=False)
-                print(f"Loaded AutoConfig for {model_name} instead")
-
-                # Use some reasonable defaults based on model type
-                if "qwen" in model_name.lower():
-                    return {"max_tokens": 2048, "temperature": 0.7, "top_k": 50, "top_p": 0.8}
-                elif "llama" in model_name.lower():
-                    return {"max_tokens": 2048, "temperature": 0.6, "top_k": 40, "top_p": 0.9}
-                elif "ernie" in model_name.lower():
-                    return {"max_tokens": 1024, "temperature": 0.7, "top_k": 50, "top_p": 0.95}
-                else:
-                    return default_params
-
-            except Exception as e2:
-                print(f"Failed to load any config for {model_name}: {e2}")
-                return default_params
-
-        # Extract relevant parameters with fallbacks to defaults
-        params = {
-            "max_tokens": getattr(gen_config, 'max_new_tokens', None) or getattr(gen_config, 'max_length', default_params["max_tokens"]),
-            "temperature": getattr(gen_config, 'temperature', default_params["temperature"]),
-            "top_k": getattr(gen_config, 'top_k', default_params["top_k"]),
-            "top_p": getattr(gen_config, 'top_p', default_params["top_p"])
-        }
+
+        # Extract parameters from generation config or use model-specific defaults
+        if gen_config:
+            params = {
+                "max_tokens": getattr(gen_config, 'max_new_tokens', None) or getattr(gen_config, 'max_length', recommended_max_tokens),
+                "temperature": getattr(gen_config, 'temperature', default_params["temperature"]),
+                "top_k": getattr(gen_config, 'top_k', default_params["top_k"]),
+                "top_p": getattr(gen_config, 'top_p', default_params["top_p"]),
+                "max_position_embeddings": max_position_embeddings,
+                "recommended_max_tokens": recommended_max_tokens
+            }
+        else:
+            # Use model-specific defaults based on model name
+            if "qwen" in model_name.lower():
+                params = {"max_tokens": recommended_max_tokens, "temperature": 0.7, "top_k": 50, "top_p": 0.8, "max_position_embeddings": max_position_embeddings, "recommended_max_tokens": recommended_max_tokens}
+            elif "llama" in model_name.lower():
+                params = {"max_tokens": recommended_max_tokens, "temperature": 0.6, "top_k": 40, "top_p": 0.9, "max_position_embeddings": max_position_embeddings, "recommended_max_tokens": recommended_max_tokens}
+            elif "ernie" in model_name.lower():
+                params = {"max_tokens": min(recommended_max_tokens, 1024), "temperature": 0.7, "top_k": 50, "top_p": 0.95, "max_position_embeddings": max_position_embeddings, "recommended_max_tokens": recommended_max_tokens}
+            else:
+                params = dict(default_params)
+                params["max_position_embeddings"] = max_position_embeddings
+                params["recommended_max_tokens"] = recommended_max_tokens
 
         # Ensure parameters are within valid ranges
-        params["max_tokens"] = max(256, min(params["max_tokens"], MAX_TOKENS))
+        params["max_tokens"] = max(256, min(params["max_tokens"], MAX_TOKENS, params["recommended_max_tokens"]))
         params["temperature"] = max(0.0, min(params["temperature"], 2.0))
         params["top_k"] = max(5, min(params["top_k"], 100))
         params["top_p"] = max(0.0, min(params["top_p"], 1.0))
@@ -126,7 +141,7 @@ def fetch_model_generation_params(model_name: str) -> dict:
         return params
 
     except Exception as e:
-        print(f"Could not fetch generation config for {model_name}: {e}")
+        print(f"Could not fetch configs for {model_name}: {e}")
         return default_params
 
 
@@ -146,8 +161,12 @@ def update_generation_params(model_name: str):
     if model_name in MODEL_GEN_PARAMS_CACHE:
         params = MODEL_GEN_PARAMS_CACHE[model_name]
         print(f"Found cached params for {model_name}: {params}")
+
+        # Set the max_tokens slider maximum to the model's recommended max
+        max_tokens_limit = min(params.get("recommended_max_tokens", MAX_TOKENS), MAX_TOKENS)
+
         return (
-            gr.update(value=params["max_tokens"]),    # max_tokens
+            gr.update(value=params["max_tokens"], maximum=max_tokens_limit),  # max_tokens with dynamic maximum
             gr.update(value=params["temperature"]),   # temperature
             gr.update(value=params["top_k"]),         # top_k
             gr.update(value=params["top_p"])          # top_p
@@ -156,7 +175,7 @@ def update_generation_params(model_name: str):
         # Fallback to defaults if model not in cache
         print(f"Model {model_name} not found in cache, using defaults")
         return (
-            gr.update(value=1024),   # max_tokens
+            gr.update(value=1024, maximum=MAX_TOKENS),   # max_tokens
             gr.update(value=0.7),    # temperature
             gr.update(value=50),     # top_k
             gr.update(value=0.95)    # top_p
@@ -183,7 +202,9 @@ def cache_all_model_params():
                 "max_tokens": 1024,
                 "temperature": 0.7,
                 "top_k": 50,
-                "top_p": 0.95
+                "top_p": 0.95,
+                "max_position_embeddings": 2048,
+                "recommended_max_tokens": 1024
             }
             MODEL_GEN_PARAMS_CACHE[model_name] = default_params
             print(f"Using default params for {model_name}: {default_params}")