app.py

import gc
import json
import os
import threading
import time
from typing import Optional

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# ---------------------------------------------------------------------------
# Model registry
# ---------------------------------------------------------------------------
MODELS = {
    "ISO27001-Expert-1.5B": {
        "base": "Qwen/Qwen2.5-1.5B-Instruct",
        "adapter": "AYI-NEDJIMI/ISO27001-Expert-1.5B",
        "default_prompt": (
            "You are ISO 27001 Expert, a specialized AI assistant for "
            "ISO/IEC 27001 information security management systems. "
            "You help organizations understand, implement, and maintain "
            "ISO 27001 certification, including risk assessment, controls "
            "from Annex A, Statement of Applicability, and audit preparation."
        ),
    },
    "RGPD-Expert-1.5B": {
        "base": "Qwen/Qwen2.5-1.5B-Instruct",
        "adapter": "AYI-NEDJIMI/RGPD-Expert-1.5B",
        "default_prompt": (
            "You are RGPD Expert, a specialized AI assistant for GDPR/RGPD "
            "data protection regulations. You help organizations understand "
            "their obligations under the General Data Protection Regulation, "
            "including data subject rights, Data Protection Impact Assessments, "
            "lawful bases for processing, and breach notification procedures."
        ),
    },
    "CyberSec-Assistant-3B": {
        "base": "Qwen/Qwen2.5-3B-Instruct",
        "adapter": "AYI-NEDJIMI/CyberSec-Assistant-3B",
        "default_prompt": (
            "You are CyberSec Assistant, an expert AI specialized in "
            "cybersecurity, compliance (GDPR, NIS2, DORA, AI Act, ISO 27001), "
            "penetration testing, SOC operations, and AI security."
        ),
    },
}

# ---------------------------------------------------------------------------
# Global model state
# ---------------------------------------------------------------------------
_lock = threading.Lock()
_loaded_model_name = None
_tokenizer = None
_model = None


def load_model(model_name: str):
    """Load or switch to a different model."""
    global _loaded_model_name, _tokenizer, _model

    with _lock:
        if _loaded_model_name == model_name and _model is not None:
            return  # Already loaded

        # Unload previous model
        if _model is not None:
            del _model
            del _tokenizer
            gc.collect()
            torch.cuda.empty_cache()

        # Load new model
        cfg = MODELS[model_name]
        hf_token = os.getenv("HF_TOKEN")

        _tokenizer = AutoTokenizer.from_pretrained(
            cfg["base"],
            trust_remote_code=True,
            token=hf_token,
        )

        base = AutoModelForCausalLM.from_pretrained(
            cfg["base"],
            torch_dtype=torch.float32,
            device_map="cpu",
            trust_remote_code=True,
            token=hf_token,
        )

        _model = PeftModel.from_pretrained(
            base,
            cfg["adapter"],
            torch_dtype=torch.float32,
            token=hf_token,
        )
        _model.eval()

        _loaded_model_name = model_name


def generate_single(
    model_name: str,
    system_prompt: str,
    user_prompt: str,
    temperature: float,
    top_p: float,
    top_k: int,
    max_tokens: int,
    repetition_penalty: float,
) -> tuple[str, dict]:
    """
    Generate a single response with metrics.
    Returns: (response_text, metrics_dict)
    """
    if not user_prompt.strip():
        return "", {}

    # Load model
    try:
        load_model(model_name)
    except Exception as e:
        return f"Error loading model: {e}", {}

    # Build messages
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]

    input_text = _tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = _tokenizer(input_text, return_tensors="pt").to("cpu")
    input_length = inputs.input_ids.shape[1]

    # Generation
    start_time = time.time()

    with torch.no_grad():
        outputs = _model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k if top_k > 0 else None,
            do_sample=temperature > 0,
            repetition_penalty=repetition_penalty,
            pad_token_id=_tokenizer.eos_token_id,
        )

    end_time = time.time()
    elapsed = end_time - start_time

    # Decode
    generated_ids = outputs[0][input_length:]
    response = _tokenizer.decode(generated_ids, skip_special_tokens=True)

    # Metrics
    num_tokens = len(generated_ids)
    tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0

    metrics = {
        "tokens": num_tokens,
        "time_sec": round(elapsed, 2),
        "tokens_per_sec": round(tokens_per_sec, 2),
    }

    return response, metrics


# ---------------------------------------------------------------------------
# UI handlers
# ---------------------------------------------------------------------------
def generate_response(
    model_name: str,
    system_prompt: str,
    user_prompt: str,
    temperature: float,
    top_p: float,
    top_k: int,
    max_tokens: int,
    repetition_penalty: float,
):
    """Handler for single-turn Q&A."""
    response, metrics = generate_single(
        model_name, system_prompt, user_prompt,
        temperature, top_p, top_k, max_tokens, repetition_penalty
    )

    metrics_text = ""
    if metrics:
        metrics_text = (
            f"**Generation Metrics:**\n"
            f"- Tokens: {metrics['tokens']}\n"
            f"- Time: {metrics['time_sec']}s\n"
            f"- Speed: {metrics['tokens_per_sec']} tokens/sec"
        )

    return response, metrics_text


def export_json(model_name: str, system_prompt: str, user_prompt: str, response: str, metrics_text: str):
    """Export conversation as JSON."""
    data = {
        "model": model_name,
        "system_prompt": system_prompt,
        "user_prompt": user_prompt,
        "response": response,
        "metrics": metrics_text,
    }
    return json.dumps(data, indent=2, ensure_ascii=False)


def generate_comparison(
    model_name: str,
    system_prompt: str,
    user_prompt: str,
    # Config A
    temp_a: float, top_p_a: float, top_k_a: int, max_tok_a: int, rep_pen_a: float,
    # Config B
    temp_b: float, top_p_b: float, top_k_b: int, max_tok_b: int, rep_pen_b: float,
):
    """Generate side-by-side comparison with different parameter sets."""

    response_a, metrics_a = generate_single(
        model_name, system_prompt, user_prompt,
        temp_a, top_p_a, top_k_a, max_tok_a, rep_pen_a
    )

    response_b, metrics_b = generate_single(
        model_name, system_prompt, user_prompt,
        temp_b, top_p_b, top_k_b, max_tok_b, rep_pen_b
    )

    metrics_text_a = ""
    if metrics_a:
        metrics_text_a = (
            f"**Config A Metrics:**\n"
            f"- Tokens: {metrics_a['tokens']}\n"
            f"- Time: {metrics_a['time_sec']}s\n"
            f"- Speed: {metrics_a['tokens_per_sec']} tok/s"
        )

    metrics_text_b = ""
    if metrics_b:
        metrics_text_b = (
            f"**Config B Metrics:**\n"
            f"- Tokens: {metrics_b['tokens']}\n"
            f"- Time: {metrics_b['time_sec']}s\n"
            f"- Speed: {metrics_b['tokens_per_sec']} tok/s"
        )

    return response_a, metrics_text_a, response_b, metrics_text_b


def update_system_prompt(model_name: str):
    """Update system prompt textbox when model changes."""
    return MODELS[model_name]["default_prompt"]


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
DESCRIPTION = """\
## Model Playground

Experiment with **3 fine-tuned cybersecurity models** using customizable parameters.

**Features:**
- Single-turn Q&A (no chat history)
- Adjustable generation parameters (temperature, top-p, top-k, max tokens, repetition penalty)
- Real-time generation metrics (tokens/sec, total time, token count)
- Export conversations as JSON
- Side-by-side comparison mode with 2 different parameter configurations
- Dark theme optimized for readability

**Models:**
- **ISO27001-Expert-1.5B**: ISO/IEC 27001 ISMS specialist
- **RGPD-Expert-1.5B**: GDPR/RGPD compliance expert
- **CyberSec-Assistant-3B**: General cybersecurity assistant
"""

theme = gr.themes.Monochrome(
    primary_hue="red",
    secondary_hue="purple",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
).set(
    body_background_fill="#0a0a0a",
    body_background_fill_dark="#0a0a0a",
    block_background_fill="#1a1a1a",
    block_background_fill_dark="#1a1a1a",
    input_background_fill="#262626",
    input_background_fill_dark="#262626",
    button_primary_background_fill="#dc2626",
    button_primary_background_fill_dark="#dc2626",
)

with gr.Blocks(theme=theme, title="Model Playground") as demo:

    gr.Markdown("# Model Playground")
    gr.Markdown(DESCRIPTION)

    with gr.Tabs():

        # ===================================================================
        # Tab 1: Single-Turn Q&A
        # ===================================================================
        with gr.Tab("Single-Turn Q&A"):
            with gr.Row():
                with gr.Column(scale=2):
                    model_select = gr.Dropdown(
                        choices=list(MODELS.keys()),
                        value="ISO27001-Expert-1.5B",
                        label="Select Model",
                    )
                with gr.Column(scale=3):
                    system_prompt_box = gr.Textbox(
                        value=MODELS["ISO27001-Expert-1.5B"]["default_prompt"],
                        label="System Prompt (Editable)",
                        lines=4,
                    )

            user_prompt_box = gr.Textbox(
                label="Your Question",
                placeholder="Enter your question here...",
                lines=3,
            )

            with gr.Accordion("Generation Parameters", open=True):
                with gr.Row():
                    temperature_slider = gr.Slider(
                        minimum=0, maximum=2, value=0.7, step=0.1,
                        label="Temperature",
                        info="Higher = more creative, lower = more deterministic"
                    )
                    top_p_slider = gr.Slider(
                        minimum=0, maximum=1, value=0.9, step=0.05,
                        label="Top-p (nucleus sampling)",
                    )
                    top_k_slider = gr.Slider(
                        minimum=0, maximum=100, value=50, step=5,
                        label="Top-k (0 = disabled)",
                    )

                with gr.Row():
                    max_tokens_slider = gr.Slider(
                        minimum=128, maximum=2048, value=512, step=128,
                        label="Max Tokens",
                    )
                    repetition_penalty_slider = gr.Slider(
                        minimum=1.0, maximum=2.0, value=1.1, step=0.1,
                        label="Repetition Penalty",
                    )

            generate_btn = gr.Button("Generate Response", variant="primary", size="lg")

            with gr.Row():
                with gr.Column(scale=3):
                    response_box = gr.Textbox(
                        label="Response",
                        lines=15,
                        interactive=False,
                    )
                with gr.Column(scale=1):
                    metrics_box = gr.Markdown(label="Metrics")

            with gr.Row():
                export_btn = gr.Button("Export as JSON")
                json_output = gr.Textbox(label="JSON Export", lines=10, visible=False)

            # Wire up events
            model_select.change(
                fn=update_system_prompt,
                inputs=[model_select],
                outputs=[system_prompt_box],
            )

            generate_btn.click(
                fn=generate_response,
                inputs=[
                    model_select, system_prompt_box, user_prompt_box,
                    temperature_slider, top_p_slider, top_k_slider,
                    max_tokens_slider, repetition_penalty_slider,
                ],
                outputs=[response_box, metrics_box],
            )

            export_btn.click(
                fn=export_json,
                inputs=[model_select, system_prompt_box, user_prompt_box, response_box, metrics_box],
                outputs=[json_output],
            ).then(
                fn=lambda: gr.update(visible=True),
                outputs=[json_output],
            )

            gr.Examples(
                examples=[
                    ["What are the mandatory clauses of ISO 27001:2022?"],
                    ["What are the 6 lawful bases for processing under GDPR?"],
                    ["Explain the MITRE ATT&CK framework."],
                    ["What are the main requirements of the NIS2 directive?"],
                ],
                inputs=user_prompt_box,
            )

        # ===================================================================
        # Tab 2: Side-by-Side Comparison
        # ===================================================================
        with gr.Tab("Side-by-Side Comparison"):
            gr.Markdown("### Compare responses from the same model with 2 different parameter configurations")

            with gr.Row():
                with gr.Column(scale=2):
                    model_select_comp = gr.Dropdown(
                        choices=list(MODELS.keys()),
                        value="ISO27001-Expert-1.5B",
                        label="Select Model",
                    )
                with gr.Column(scale=3):
                    system_prompt_comp = gr.Textbox(
                        value=MODELS["ISO27001-Expert-1.5B"]["default_prompt"],
                        label="System Prompt (Editable)",
                        lines=4,
                    )

            user_prompt_comp = gr.Textbox(
                label="Your Question",
                placeholder="Enter your question here...",
                lines=3,
            )

            with gr.Row():
                # Config A
                with gr.Column():
                    gr.Markdown("#### Configuration A")
                    temp_a = gr.Slider(0, 2, value=0.3, step=0.1, label="Temperature")
                    top_p_a = gr.Slider(0, 1, value=0.9, step=0.05, label="Top-p")
                    top_k_a = gr.Slider(0, 100, value=50, step=5, label="Top-k")
                    max_tok_a = gr.Slider(128, 2048, value=512, step=128, label="Max Tokens")
                    rep_pen_a = gr.Slider(1.0, 2.0, value=1.1, step=0.1, label="Rep. Penalty")

                # Config B
                with gr.Column():
                    gr.Markdown("#### Configuration B")
                    temp_b = gr.Slider(0, 2, value=1.2, step=0.1, label="Temperature")
                    top_p_b = gr.Slider(0, 1, value=0.95, step=0.05, label="Top-p")
                    top_k_b = gr.Slider(0, 100, value=40, step=5, label="Top-k")
                    max_tok_b = gr.Slider(128, 2048, value=512, step=128, label="Max Tokens")
                    rep_pen_b = gr.Slider(1.0, 2.0, value=1.2, step=0.1, label="Rep. Penalty")

            compare_btn = gr.Button("Generate Comparison", variant="primary", size="lg")

            with gr.Row():
                with gr.Column():
                    response_a = gr.Textbox(label="Response A", lines=12, interactive=False)
                    metrics_a = gr.Markdown()
                with gr.Column():
                    response_b = gr.Textbox(label="Response B", lines=12, interactive=False)
                    metrics_b = gr.Markdown()

            # Wire up events
            model_select_comp.change(
                fn=update_system_prompt,
                inputs=[model_select_comp],
                outputs=[system_prompt_comp],
            )

            compare_btn.click(
                fn=generate_comparison,
                inputs=[
                    model_select_comp, system_prompt_comp, user_prompt_comp,
                    temp_a, top_p_a, top_k_a, max_tok_a, rep_pen_a,
                    temp_b, top_p_b, top_k_b, max_tok_b, rep_pen_b,
                ],
                outputs=[response_a, metrics_a, response_b, metrics_b],
            )

            gr.Examples(
                examples=[
                    ["What is a Data Protection Impact Assessment?"],
                    ["Explain the concept of Zero Trust security."],
                    ["What are the penalties for GDPR non-compliance?"],
                ],
                inputs=user_prompt_comp,
            )

    # Footer
    gr.HTML("""
    <div style="text-align:center; margin-top:2rem; padding-top:1rem; border-top:1px solid #333; color:#888; font-size:0.85rem;">
      <p>Built by <a href="https://huggingface.co/AYI-NEDJIMI" style="color:#dc2626;">Ayi NEDJIMI</a>
      | Models: <a href="https://huggingface.co/AYI-NEDJIMI/ISO27001-Expert-1.5B" style="color:#dc2626;">ISO27001</a>,
      <a href="https://huggingface.co/AYI-NEDJIMI/RGPD-Expert-1.5B" style="color:#dc2626;">RGPD</a>,
      <a href="https://huggingface.co/AYI-NEDJIMI/CyberSec-Assistant-3B" style="color:#dc2626;">CyberSec-3B</a>
      | <a href="https://huggingface.co/collections/AYI-NEDJIMI/cybersec-ai-portfolio-datasets-models-and-spaces-699224074a478ec0feeac493" style="color:#dc2626;">Portfolio</a></p>
      <p style="font-size:0.75rem; color:#666;">Fine-tuned with QLoRA on Qwen 2.5 | Model Playground</p>
    </div>
    """)

if __name__ == "__main__":
    demo.launch()