amirimmd/ExBERT-Classifier-Inference

1import streamlit as st2import torch3import torch.nn as nn4from transformers import AutoModel, AutoTokenizer5import shap6import numpy as np7import matplotlib.pyplot as plt8from huggingface_hub import hf_hub_download9 10try:11    from safetensors.torch import load_file as load_safetensors12except ImportError:13    load_safetensors = None14 15 16# --- 1. Model Architecture Definition (ExBERT) ---17 18class ExBERT_Classifier(nn.Module):19    def __init__(self, phase1_model_path, num_labels=1, use_multi_layer=True, dropout_rate=0.2):20        super(ExBERT_Classifier, self).__init__()21        self.bert = AutoModel.from_pretrained(phase1_model_path, output_hidden_states=True)22        self.lstm = nn.LSTM(23            input_size=self.bert.config.hidden_size, hidden_size=256,24            num_layers=1, batch_first=True, bidirectional=False25        )26        self.layer_norm = nn.LayerNorm(256)27        self.dropout = nn.Dropout(dropout_rate)28        self.classifier = nn.Linear(256, num_labels)29 30    def forward(self, input_ids, attention_mask, labels=None, **kwargs):31        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)32        hidden_states = outputs.hidden_states[-4:]33        pooled_outputs = [h[:, 0, :] for h in hidden_states]34        stacked_output = torch.stack(pooled_outputs, dim=1)35        lstm_output, _ = self.lstm(stacked_output)36        lstm_pooled = lstm_output[:, -1, :]37        normalized_output = self.layer_norm(lstm_pooled)38        dropped_output = self.dropout(normalized_output)39        logits = self.classifier(dropped_output)40        return logits41 42 43# --- 2. Page Configuration ---44 45st.set_page_config(46    page_title="VulnSight AI Core",47    page_icon="🛡️",48    layout="wide",49    initial_sidebar_state="expanded"50)51 52st.markdown("""53<style>54    .reportview-container { background: #050505; color: #e0e0e0; }55    .sidebar .sidebar-content { background: #0a0a0a; }56    h1, h2, h3 { color: #00f0ff !important; font-family: 'Courier New', monospace; }57    .stButton>button { border: 1px solid #00f0ff; color: #00f0ff; background: transparent; transition: all 0.3s; }58    .stButton>button:hover { background: rgba(0, 240, 255, 0.1); box-shadow: 0 0 10px rgba(0, 240, 255, 0.5); }59    .stTextArea textarea { background-color: #111; color: #fff; border: 1px solid #333; }60    div[data-testid="stMarkdownContainer"] p { font-family: 'Consolas', monospace; }61    .metric-box { border: 1px solid #333; padding: 10px; border-radius: 5px; text-align: center; background: #111; }62    .metric-label { font-size: 0.8em; color: #888; }63    .metric-value { font-size: 1.5em; font-weight: bold; color: #fff; }64</style>65""", unsafe_allow_html=True)66 67st.title("🛡️ VulnSight: AI Inference Core")68st.markdown("Intelligent Vulnerability Analysis System based on **ExBERT** Architecture and **SHAP** Interpretability.")69 70 71# --- 3. Session State Initialization ---72 73if 'prediction_result' not in st.session_state:74    st.session_state.prediction_result = None75if 'shap_html' not in st.session_state:76    st.session_state.shap_html = None77 78 79# --- 4. Model Selection ---80 81USER_NAME = "amirimmd"82 83MODEL_MAPPINGS = {84    "ExBERT Phase 2 (Baseline)": f"{USER_NAME}/ExBERT-Phase2-Best",85    "ExBERT Phase 3 (Method 1)": f"{USER_NAME}/ExBERT-Phase3-Method1",86    "ExBERT Phase 3 (Method 3)": f"{USER_NAME}/ExBERT-Phase3-Method3"87}88 89with st.sidebar:90    st.header("⚙️ Model Configuration")91    selected_model_name = st.selectbox("Select Model Version:", list(MODEL_MAPPINGS.keys()))92    selected_repo_id = MODEL_MAPPINGS[selected_model_name]93    st.info(f"Loaded Repository:\n`{selected_repo_id}`")94    use_shap = st.checkbox("Active XAI Engine (SHAP)", value=False, help="Calculation may take a few seconds depending on text length.")95 96 97# --- 5. Model Loading Function ---98 99@st.cache_resource100def load_model(repo_id):101    try:102        tokenizer = AutoTokenizer.from_pretrained(repo_id)103        model = ExBERT_Classifier(phase1_model_path=repo_id)104        try:105            weights_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")106            if load_safetensors:107                state_dict = load_safetensors(weights_path)108            else:109                return None, "Safetensors library not installed."110        except Exception:111            weights_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")112            state_dict = torch.load(weights_path, map_location=torch.device('cpu'))113        model.load_state_dict(state_dict)114        model.eval()115        return tokenizer, model116    except Exception as e:117        return None, str(e)118 119 120# --- 6. Main Logic ---121 122tokenizer, model = load_model(selected_repo_id)123 124if isinstance(model, str) or model is None:125    st.error(f"❌ Error loading model: {model}")126    st.warning("Please check if the Hugging Face repository exists and is Public.")127else:128    col1, col2 = st.columns([2, 1])129 130    with col1:131        text_input = st.text_area(132            "Vulnerability Description (CVE):",133            height=150,134            placeholder="Example: Buffer overflow in libssh2 before 1.11.0 allows remote attackers to execute arbitrary code..."135        )136 137    with col2:138        st.markdown("### Analysis Options")139        st.write("Click below to analyze the text using the selected neural network.")140        analyze_btn = st.button("🚀 Analyze Vulnerability", use_container_width=True)141 142    # --- Prediction ---143    if analyze_btn and text_input:144        with st.spinner(f"🧠 Processing with {selected_model_name}..."):145            inputs = tokenizer(text_input, return_tensors="pt", padding=True, truncation=True, max_length=256)146            with torch.no_grad():147                logits = model(**inputs)148                prob = torch.sigmoid(logits).item()149 150            if prob > 0.5:151                prediction_label = "EXPLOITABLE (High Risk)"152                confidence_score = prob * 100153                result_color = "red"154                status_icon = "🔴"155            else:156                prediction_label = "NOT EXPLOITABLE (Low Risk)"157                confidence_score = (1 - prob) * 100158                result_color = "green"159                status_icon = "🟢"160 161            # ذخیره در session_state — باقی می‌ماند تا analysis جدید162            st.session_state.prediction_result = {163                'label': prediction_label,164                'confidence': confidence_score,165                'prob': prob,166                'color': result_color,167                'icon': status_icon,168                'model': selected_model_name,169            }170            st.session_state.shap_html = None  # SHAP قبلی پاک می‌شود171 172        # --- SHAP Analysis ---173        if use_shap:174            with st.spinner("Calculating feature importance..."):175                try:176                    def f(texts):177                        tv = torch.tensor([tokenizer.encode(v, padding='max_length', max_length=256, truncation=True) for v in texts])178                        outputs = model(tv, attention_mask=(tv != 0).type(torch.int64))179                        return torch.sigmoid(outputs).detach().numpy()180 181                    explainer = shap.Explainer(f, tokenizer)182                    shap_values = explainer([text_input])183                    st.session_state.shap_html = shap.plots.text(shap_values, display=False)184                except Exception as e:185                    st.warning(f"Could not generate SHAP explanation: {e}")186 187    # --- Display Results (همیشه از session_state خوانده می‌شود) ---188    if st.session_state.prediction_result:189        result = st.session_state.prediction_result190        st.divider()191        st.markdown(192            f"<h2 style='text-align: center; color: {result['color']};'>{result['icon']} {result['label']}</h2>",193            unsafe_allow_html=True194        )195 196        m_col1, m_col2, m_col3 = st.columns(3)197        with m_col1:198            st.markdown(f"<div class='metric-box'><div class='metric-label'>Model Used</div><div class='metric-value' style='font-size: 1em;'>{result['model']}</div></div>", unsafe_allow_html=True)199        with m_col2:200            st.markdown(f"<div class='metric-box'><div class='metric-label'>Confidence Score</div><div class='metric-value' style='color: {result['color']};'>{result['confidence']:.2f}%</div></div>", unsafe_allow_html=True)201        with m_col3:202            st.markdown(f"<div class='metric-box'><div class='metric-label'>Raw Probability</div><div class='metric-value' style='font-size: 1em;'>{result['prob']:.4f}</div></div>", unsafe_allow_html=True)203 204        if st.session_state.shap_html:205            st.divider()206            st.subheader("XAI Interpretation (SHAP)")207            st.caption("Red highlights increase exploitability risk, Blue decreases it.")208            st.components.v1.html(st.session_state.shap_html, height=350, scrolling=True)