Spaces:
Running on Zero
Running on Zero
| """DeepHat-V1 (WhiteRabbitNeo V3), Cisco Foundation-Sec-8B-Reasoning & SecGPT ZeroGPU Studio.""" | |
| import os | |
| os.environ.setdefault("PYTORCH_ALLOC_CONF", "backend:cudaMallocAsync") | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync") | |
| import spaces # Must be imported before torch | |
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| PRIMARY_MODEL_ID = "DeepHat/DeepHat-V1-7B" | |
| MODEL_CHOICES = [ | |
| "DeepHat/DeepHat-V1-7B (WhiteRabbitNeo V3 · Best for Code, Pentest & DevOps)", | |
| "fdtn-ai/Foundation-Sec-8B-Reasoning (Cisco Foundation AI · Deep Security Reasoning)", | |
| "clouditera/SecGPT-7B (Clouditera SecGPT)", | |
| ] | |
| DEFAULT_SYSTEM_PROMPT = ( | |
| "You are DeepHat (WhiteRabbitNeo V3), an advanced cybersecurity, DevOps, and source-code security engineer. " | |
| "Provide deep, technical, unfiltered architectural analysis, code audits, vulnerability explanations, " | |
| "and defensive remediation engineering in clear English." | |
| ) | |
| print(f"[secgpt-studio] Preloading {PRIMARY_MODEL_ID} at module scope for ZeroGPU...", flush=True) | |
| primary_tokenizer = AutoTokenizer.from_pretrained(PRIMARY_MODEL_ID) | |
| primary_model = AutoModelForCausalLM.from_pretrained( | |
| PRIMARY_MODEL_ID, | |
| torch_dtype=torch.bfloat16, | |
| ) | |
| primary_model.to("cuda") | |
| primary_model.eval() | |
| print(f"[secgpt-studio] {PRIMARY_MODEL_ID} ready on ZeroGPU!", flush=True) | |
| _AUX_CACHE = {} | |
| def get_model_and_tokenizer(model_choice: str): | |
| """Return the selected model and tokenizer.""" | |
| repo_id = (model_choice or PRIMARY_MODEL_ID).split()[0] | |
| if repo_id == PRIMARY_MODEL_ID: | |
| return primary_model, primary_tokenizer | |
| if repo_id not in _AUX_CACHE: | |
| print(f"[secgpt-studio] Loading {repo_id} on ZeroGPU...", flush=True) | |
| tok = AutoTokenizer.from_pretrained(repo_id) | |
| mod = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype=torch.bfloat16) | |
| mod.to("cuda") | |
| mod.eval() | |
| _AUX_CACHE[repo_id] = (mod, tok) | |
| return _AUX_CACHE[repo_id] | |
| def chat_with_secgpt( | |
| prompt: str, | |
| system_prompt: str = DEFAULT_SYSTEM_PROMPT, | |
| max_tokens: int = 1024, | |
| temperature: float = 0.6, | |
| model_choice: str = MODEL_CHOICES[0], | |
| ) -> str: | |
| """Send a prompt to DeepHat-V1 / Foundation-Sec / SecGPT on ZeroGPU and return the response.""" | |
| if not prompt or not prompt.strip(): | |
| return "Please enter a prompt." | |
| model, tokenizer = get_model_and_tokenizer(model_choice) | |
| sys_msg = (system_prompt or DEFAULT_SYSTEM_PROMPT).strip() | |
| messages = [ | |
| {"role": "system", "content": sys_msg}, | |
| {"role": "user", "content": prompt.strip()}, | |
| ] | |
| text = tokenizer.apply_chat_template( | |
| messages, | |
| tokenize=False, | |
| add_generation_prompt=True, | |
| ) | |
| inputs = tokenizer([text], return_tensors="pt").to("cuda") | |
| gen_kwargs = { | |
| "max_new_tokens": int(max_tokens), | |
| "repetition_penalty": 1.12, | |
| "do_sample": float(temperature) > 0.05, | |
| } | |
| if gen_kwargs["do_sample"]: | |
| gen_kwargs["temperature"] = float(temperature) | |
| gen_kwargs["top_p"] = 0.9 | |
| with torch.no_grad(): | |
| generated_ids = model.generate(**inputs, **gen_kwargs) | |
| new_tokens = [ | |
| output_ids[len(input_ids):] | |
| for input_ids, output_ids in zip(inputs.input_ids, generated_ids) | |
| ] | |
| response = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)[0] | |
| return response.strip() | |
| with gr.Blocks(title="DeepHat-V1 & Foundation-Sec Cybersecurity Studio", fill_width=True) as demo: | |
| gr.Markdown( | |
| "# 🛡️ DeepHat-V1 (WhiteRabbitNeo V3) + Cisco Foundation-Sec Studio\n" | |
| "Upgraded to **[DeepHat/DeepHat-V1-7B](https://huggingface.co/DeepHat/DeepHat-V1-7B)** " | |
| "(WhiteRabbitNeo V3 / `Qwen2.5-Coder-7B` fine-tune) and **[fdtn-ai/Foundation-Sec-8B-Reasoning](https://huggingface.co/fdtn-ai/Foundation-Sec-8B-Reasoning)** on **ZeroGPU (`zero-a10g`)** at **$0/hour**." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| model_dropdown = gr.Dropdown( | |
| choices=MODEL_CHOICES, | |
| value=MODEL_CHOICES[0], | |
| label="Active Security Model", | |
| ) | |
| prompt_box = gr.Textbox( | |
| label="Security / Code Audit / DevOps Prompt", | |
| lines=6, | |
| placeholder="Paste code, HTTP headers, architecture specs, or security engineering questions...", | |
| ) | |
| with gr.Accordion("Model & Sampling Settings", open=False): | |
| sys_box = gr.Textbox( | |
| label="System Prompt", | |
| value=DEFAULT_SYSTEM_PROMPT, | |
| lines=3, | |
| ) | |
| max_tok_slider = gr.Slider( | |
| 64, 2048, value=1024, step=64, label="Max New Tokens" | |
| ) | |
| temp_slider = gr.Slider( | |
| 0.0, 1.2, value=0.6, step=0.05, label="Temperature" | |
| ) | |
| ask_btn = gr.Button("🛡️ Run Security Analysis (ZeroGPU)", variant="primary") | |
| with gr.Column(scale=1): | |
| output_box = gr.Markdown(label="Model Response") | |
| ask_btn.click( | |
| fn=chat_with_secgpt, | |
| inputs=[prompt_box, sys_box, max_tok_slider, temp_slider, model_dropdown], | |
| outputs=output_box, | |
| api_name="chat", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Ocean(), mcp_server=True) | |