add gen param cache

app.py

@@ -6,7 +6,7 @@ from supabase import create_client, Client
 from supabase.client import ClientOptions
 from enum import Enum
 from datasets import get_dataset_infos
-from transformers import AutoConfig
+from transformers import AutoConfig, GenerationConfig
 from huggingface_hub import whoami
 from typing import Optional, Union
 
@@ -17,41 +17,6 @@ from typing import Optional, Union
  - validate max model params
 """
 
-SUPPORTED_MODELS = [
-    "Qwen/Qwen3-4B-Instruct-2507",
-    "Qwen/Qwen3-30B-A3B-Instruct-2507",
-    "meta-llama/Llama-3.2-1B-Instruct",
-    "meta-llama/Llama-3.2-3B-Instruct",
-    "baidu/ERNIE-4.5-21B-A3B-Thinking",
-    "LLM360/K2-Think",
-    "openai/gpt-oss-20b",
-]
-
-
-def verify_pro_status(token: Optional[Union[gr.OAuthToken, str]]) -> bool:
-    """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
-    if not token:
-        return False
-    
-    if isinstance(token, gr.OAuthToken):
-        token_str = token.token
-    elif isinstance(token, str):
-        token_str = token
-    else:
-        return False
-    
-    try:
-        user_info = whoami(token=token_str)
-        return (
-            user_info.get("isPro", False) or
-            any(org.get("isEnterprise", False) for org in user_info.get("orgs", []))
-        )
-    except Exception as e:
-        print(f"Could not verify user's PRO/Enterprise status: {e}")
-        return False
-
-
-
 class GenerationStatus(Enum):
     PENDING = "PENDING"
     RUNNING = "RUNNING"
@@ -64,6 +29,9 @@ MAX_SAMPLES_FREE = 100   # max number of samples for free users
 MAX_TOKENS = 8192
 MAX_MODEL_PARAMS = 20_000_000_000  # 20 billion parameters (for now)
 
+# Cache for model generation parameters
+MODEL_GEN_PARAMS_CACHE = {}
+
 @dataclass
 class GenerationRequest:
     id: str
@@ -89,6 +57,166 @@ class GenerationRequest:
     num_output_examples: int
     private: bool = False
     num_retries: int = 0
+    
+SUPPORTED_MODELS = [
+    "Qwen/Qwen3-4B-Instruct-2507",
+    "Qwen/Qwen3-30B-A3B-Instruct-2507",
+    "meta-llama/Llama-3.2-1B-Instruct",
+    "meta-llama/Llama-3.2-3B-Instruct",
+    "baidu/ERNIE-4.5-21B-A3B-Thinking",
+    "LLM360/K2-Think",
+    "openai/gpt-oss-20b",
+]
+
+
+def fetch_model_generation_params(model_name: str) -> dict:
+    """Fetch generation parameters from model's generation config on the hub"""
+    default_params = {
+        "max_tokens": 1024,
+        "temperature": 0.7,
+        "top_k": 50,
+        "top_p": 0.95
+    }
+
+    try:
+        print(f"Attempting to fetch generation config for: {model_name}")
+
+        # Try to load the generation config
+        try:
+            gen_config = GenerationConfig.from_pretrained(model_name, force_download=False)
+            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"])
+        }
+
+        # Ensure parameters are within valid ranges
+        params["max_tokens"] = max(256, min(params["max_tokens"], 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))
+
+        print(f"Final params for {model_name}: {params}")
+        return params
+
+    except Exception as e:
+        print(f"Could not fetch generation config for {model_name}: {e}")
+        return default_params
+
+
+def update_generation_params(model_name: str):
+    """Update generation parameters based on selected model"""
+    global MODEL_GEN_PARAMS_CACHE
+
+    print(f"Updating generation parameters for model: {model_name}")
+    print(f"Cache is empty: {len(MODEL_GEN_PARAMS_CACHE) == 0}")
+    print(f"Current cache keys: {list(MODEL_GEN_PARAMS_CACHE.keys())}")
+
+    # If cache is empty, try to populate it now
+    if len(MODEL_GEN_PARAMS_CACHE) == 0:
+        print("Cache is empty, attempting to populate now...")
+        cache_all_model_params()
+
+    if model_name in MODEL_GEN_PARAMS_CACHE:
+        params = MODEL_GEN_PARAMS_CACHE[model_name]
+        print(f"Found cached params for {model_name}: {params}")
+        return (
+            gr.update(value=params["max_tokens"]),    # max_tokens
+            gr.update(value=params["temperature"]),   # temperature
+            gr.update(value=params["top_k"]),         # top_k
+            gr.update(value=params["top_p"])          # top_p
+        )
+    else:
+        # 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=0.7),    # temperature
+            gr.update(value=50),     # top_k
+            gr.update(value=0.95)    # top_p
+        )
+
+
+def cache_all_model_params():
+    """Cache generation parameters for all supported models at startup"""
+    global MODEL_GEN_PARAMS_CACHE
+
+    print(f"Starting to cache parameters for {len(SUPPORTED_MODELS)} models...")
+    print(f"Supported models: {SUPPORTED_MODELS}")
+
+    for model_name in SUPPORTED_MODELS:
+        try:
+            print(f"Processing model: {model_name}")
+            params = fetch_model_generation_params(model_name)
+            MODEL_GEN_PARAMS_CACHE[model_name] = params
+            print(f"Successfully cached params for {model_name}: {params}")
+        except Exception as e:
+            print(f"Exception while caching params for {model_name}: {e}")
+            # Use default parameters if caching fails
+            default_params = {
+                "max_tokens": 1024,
+                "temperature": 0.7,
+                "top_k": 50,
+                "top_p": 0.95
+            }
+            MODEL_GEN_PARAMS_CACHE[model_name] = default_params
+            print(f"Using default params for {model_name}: {default_params}")
+
+    print(f"Caching complete. Final cache contents:")
+    for model, params in MODEL_GEN_PARAMS_CACHE.items():
+        print(f"  {model}: {params}")
+    print(f"Cache size: {len(MODEL_GEN_PARAMS_CACHE)} models")
+
+def verify_pro_status(token: Optional[Union[gr.OAuthToken, str]]) -> bool:
+    """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
+    if not token:
+        return False
+    
+    if isinstance(token, gr.OAuthToken):
+        token_str = token.token
+    elif isinstance(token, str):
+        token_str = token
+    else:
+        return False
+    
+    try:
+        user_info = whoami(token=token_str)
+        return (
+            user_info.get("isPro", False) or
+            any(org.get("isEnterprise", False) for org in user_info.get("orgs", []))
+        )
+    except Exception as e:
+        print(f"Could not verify user's PRO/Enterprise status: {e}")
+        return False
+
+
+
 
 def validate_request(request: GenerationRequest, oauth_token: Optional[Union[gr.OAuthToken, str]] = None) -> GenerationRequest:
     # checks that the request is valid
@@ -401,6 +529,11 @@ def get_generation_stats_safe():
 
 
 def main():
+    # Cache model generation parameters at startup
+    print("Caching model generation parameters...")
+    cache_all_model_params()
+    print("Model parameter caching complete.")
+
     with gr.Blocks(title="Synthetic Data Generation") as demo:
         gr.HTML("<h3 style='text-align:center'>Generate synthetic data with AI models. Free to use! Sign up for PRO benefits (10k samples vs 100). <a href='http://huggingface.co/subscribe/pro?source=synthetic-data-universe' target='_blank'>Upgrade to PRO</a></h3>", elem_id="sub_title")
         
@@ -659,6 +792,13 @@ def main():
                 outputs=[input_dataset_config, input_dataset_split, prompt_column, output_dataset_name, num_output_samples, load_info_status]
             )
 
+            # Wire up model change to update generation parameters
+            model_name_or_path.change(
+                update_generation_params,
+                inputs=[model_name_or_path],
+                outputs=[max_tokens, temperature, top_k, top_p]
+            )
+
             submit_btn.click(
                 submit_request,
                 inputs=[input_dataset_name, input_dataset_split, input_dataset_config, output_dataset_name, prompt_column, model_name_or_path,