Update app.py

app.py

@@ -1,87 +1,58 @@
-# app.py - FIXED VERSION
+# app.py - LOAD ON DEMAND
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import json
+import subprocess
+import tempfile
 import os
 
-MODEL_ID = "fdtn-ai/Foundation-Sec-8B"
-
-print("🚀 Loading model...")
-
-# FIX: Download and patch config first
-from huggingface_hub import hf_hub_download
-
-# Download config
-config_path = hf_hub_download(
-    repo_id=MODEL_ID,
-    filename="config.json",
-    local_dir="./cache"
-)
-
-# Read and fix config
-with open(config_path, 'r') as f:
-    config_data = json.load(f)
-
-# Fix rope_scaling for Llama 3
-if 'rope_scaling' in config_data:
-    rope = config_data['rope_scaling']
-    if isinstance(rope, dict):
-        # Convert to standard format
-        rope_scaling = {
-            "type": rope.get("rope_type", "linear"),
-            "factor": rope.get("factor", 1.0)
-        }
-        config_data['rope_scaling'] = rope_scaling
-
-# Save fixed config
-os.makedirs("./fixed_config", exist_ok=True)
-fixed_config_path = "./fixed_config/config.json"
-with open(fixed_config_path, 'w') as f:
-    json.dump(config_data, f)
-
-# Load with fixed config
-from transformers import AutoConfig
-config = AutoConfig.from_pretrained(fixed_config_path)
-
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+def generate(prompt):
+    """Load model on-demand using transformers CLI"""
+    # Create a temporary script
+    script = f"""
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load model
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    config=config,
+    "fdtn-ai/Foundation-Sec-8B",
     torch_dtype=torch.float16,
     device_map="auto",
     trust_remote_code=True
 )
+tokenizer = AutoTokenizer.from_pretrained("fdtn-ai/Foundation-Sec-8B")
 
-print("✅ Model loaded!")
-
-def generate(prompt, max_tokens=200):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+inputs = tokenizer('{prompt}', return_tensors="pt").to(model.device)
+outputs = model.generate(**inputs, max_new_tokens=200)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+"""
     
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=max_tokens,
-        temperature=0.7,
-        do_sample=True
-    )
+    # Write to temp file
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+        f.write(script)
+        script_path = f.name
     
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Create interface
+    try:
+        # Run script
+        result = subprocess.run(
+            ['python', script_path],
+            capture_output=True,
+            text=True,
+            timeout=120
+        )
+        
+        # Cleanup
+        os.unlink(script_path)
+        
+        if result.returncode == 0:
+            return result.stdout.strip()
+        else:
+            return f"Error: {result.stderr}"
+            
+    except subprocess.TimeoutExpired:
+        return "Timeout - Model loading took too long"
+
+# Launch interface
 gr.Interface(
     generate,
-    [
-        gr.Textbox(label="Prompt", lines=3),
-        gr.Slider(50, 500, value=200, label="Max Tokens")
-    ],
+    gr.Textbox(label="Ask about cybersecurity:"),
     gr.Textbox(label="Response", lines=10),
-    title="🔒 Foundation-Sec-8B",
-    examples=[
-        ["Explain cybersecurity:"],
-        ["What is a firewall?"],
-        ["How to create strong passwords?"]
-    ]
+    title="Foundation-Sec-8B (On-demand Loading)"
 ).launch(server_name="0.0.0.0")