Update app.py

app.py

@@ -1,31 +1,76 @@
-# app.py - PUT IN YOUR SPACE
-from transformers import pipeline
+# app.py - FIXED VERSION
 import gradio as gr
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import json
+import os
 
-print("🚀 Loading Foundation-Sec-8B...")
+MODEL_ID = "fdtn-ai/Foundation-Sec-8B"
 
-# Load with optimizations for ZeroGPU
-model = pipeline(
-    "text-generation",
-    model="fdtn-ai/Foundation-Sec-8B",
-    device_map="auto",
+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)
+
+# Load model
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    config=config,
     torch_dtype=torch.float16,
+    device_map="auto",
     trust_remote_code=True
 )
 
 print("✅ Model loaded!")
 
 def generate(prompt, max_tokens=200):
-    result = model(
-        prompt,
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    outputs = model.generate(
+        **inputs,
         max_new_tokens=max_tokens,
         temperature=0.7,
         do_sample=True
     )
-    return result[0]['generated_text']
+    
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Simple interface
+# Create interface
 gr.Interface(
     generate,
     [
@@ -39,4 +84,4 @@ gr.Interface(
         ["What is a firewall?"],
         ["How to create strong passwords?"]
     ]
-).launch()
+).launch(server_name="0.0.0.0")