SRCZ/AI-Security-Generator

1import gradio as gr2from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline3 4# ✅ Light model that works fine on CPU Basic5MODEL = "Qwen/Qwen2.5-3B-Instruct"6 7tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)8mdl = AutoModelForCausalLM.from_pretrained(9    MODEL,10    low_cpu_mem_usage=True,11    trust_remote_code=True12)13 14pipe = pipeline(15    "text-generation",16    model=mdl,17    tokenizer=tok,18    max_new_tokens=400,19    do_sample=False,20    temperature=0.221)22 23SYSTEM_PROMPT = """24You are a senior Gulf-region cybersecurity expert.25Write professional, audit-ready documents in English (concise tone).26Target alignment: SAMA CSF, NCA ECC, ISO 27001, PCI DSS.27 28Each output MUST include:291) Executive Summary302) Scope & Responsibilities (RACI)313) Step-by-Step Procedures324) Controls & KPIs (measurable)335) Appendices (Incident form, IOCs list, Timeline table)34 35Return clean Markdown with headings and tables.36"""37 38def generate(doc_type, industry, region, standards, inputs):39    prompt = f"""{SYSTEM_PROMPT}40Document: {doc_type}41Industry: {industry}42Region: {region}43Regulatory/Standards: {standards}44Customer Inputs: {inputs}45"""46    result = pipe(prompt)[0]["generated_text"]47    return result48 49with gr.Blocks() as demo:50    gr.Markdown("# SOC / IR AI Generator (KSA / GCC)")51    with gr.Row():52        doc = gr.Dropdown(53            ["Incident Report","IR Playbook","Access Control Policy",54             "VPN Policy","Email Security SOP","Vulnerability Mgmt SOP"],55            label="Document"56        )57        ind = gr.Dropdown(58            ["FinTech","Healthcare","Retail","Government","Education","Manufacturing","Telecom","Oil & Gas","E-Commerce"],59            label="Industry"60        )61        reg = gr.Dropdown(["KSA (SAMA/NCA)","GCC (generic)","MENA (generic)"], label="Region")62        std = gr.CheckboxGroup(["SAMA CSF","NCA ECC","ISO 27001","PCI DSS"], label="Standards")63    inp = gr.Textbox(lines=8, label="Inputs (context / tools / teams / SLAs)")64    btn = gr.Button("Generate")65    out = gr.Markdown()66    btn.click(generate, [doc, ind, reg, std, inp], out)67 68print("✅ App loaded successfully")69demo.launch(share=True, debug=True)