Expand model card with Mermaid architecture, training, and inference diagrams

README.md

@@ -39,6 +39,140 @@ Fine-tuned [`microsoft/graphcodebert-base`](https://huggingface.co/microsoft/gra
 
 Part of [SolidityGuard](https://github.com/tanaymitra54/solidity_guard_razorpay) — used as a first-pass detector alongside Slither and an LLM auditor.
 
+## Architecture overview
+
+How this model sits inside the SolidityGuard audit stack:
+
+```mermaid
+flowchart TB
+    subgraph Input
+        SOL["Solidity source (.sol)"]
+    end
+
+    subgraph Detectors["First-pass detectors"]
+        SL["Slither / patterns"]
+        GCB["Graph CodeBERT<br/>this model"]
+    end
+
+    subgraph LLM["LLM agents"]
+        SCAN["Scanner"]
+        AN["Analyzer"]
+        EX["Exploit Gen"]
+        FX["Fix Suggester"]
+    end
+
+    subgraph Out["Output"]
+        REP["Audit findings<br/>+ severity + confidence"]
+    end
+
+    SOL --> SL
+    SOL --> GCB
+    SOL --> SCAN
+    SL --> SCAN
+    GCB --> SCAN
+    SCAN --> AN
+    AN --> EX
+    AN --> FX
+    SCAN --> REP
+    AN --> REP
+    EX --> REP
+    FX --> REP
+```
+
+## Model internals
+
+Sequence-classification head on Graph CodeBERT (RoBERTa-style encoder):
+
+```mermaid
+flowchart LR
+    A["Solidity source<br/>string"] --> B["Tokenizer<br/>max_length=512"]
+    B --> C["input_ids<br/>attention_mask"]
+    C --> D["Graph CodeBERT<br/>encoder<br/>~125M params"]
+    D --> E["[CLS] pooled<br/>representation"]
+    E --> F["Linear classifier<br/>12 logits"]
+    F --> G["Softmax"]
+    G --> H["label + confidence"]
+```
+
+## Training pipeline
+
+End-to-end fine-tune path used for this checkpoint:
+
+```mermaid
+flowchart TD
+    W["SmartBugs-Wild<br/>~47k contracts"] --> P["Parse + cap<br/>WILD_LIMIT=5000"]
+    R["SmartBugs-Results<br/>results_wild.json"] --> V["Tool consensus<br/>≥2 tools agree"]
+    P --> M["Labeled samples"]
+    V --> M
+    M --> S["Split 70 / 15 / 15<br/>train / val / test"]
+    S --> AUG["Light augmentation"]
+    AUG --> FT["Fine-tune<br/>Graph CodeBERT-base"]
+    FT --> ES["Early stopping<br/>on val macro-F1"]
+    ES --> BEST["checkpoints/best<br/>→ this Hub repo"]
+```
+
+## Data labeling logic
+
+```mermaid
+flowchart LR
+    C["Contract address"] --> T1["Tool A categories"]
+    C --> T2["Tool B categories"]
+    C --> T3["Tool N categories"]
+    T1 --> VOTE["Vote per category"]
+    T2 --> VOTE
+    T3 --> VOTE
+    VOTE --> Q{"votes ≥ 2?"}
+    Q -->|yes| LBL["Mapped vuln label<br/>e.g. reentrancy"]
+    Q -->|no| SAFE["safe"]
+```
+
+## Inference sequence
+
+```mermaid
+sequenceDiagram
+    participant U as Caller / API
+    participant D as GraphCodeBERTDetector
+    participant H as Hugging Face Hub
+    participant M as Model (GPU/CPU)
+
+    U->>D: scan(source_code)
+    alt first load
+        D->>H: from_pretrained(repo_id)
+        H-->>D: weights + config + tokenizer
+        D->>M: eval mode
+    end
+    D->>M: tokenize + forward
+    M-->>D: logits → softmax
+    alt label != safe AND conf ≥ threshold
+        D-->>U: Finding(issue_type, severity, confidence)
+    else
+        D-->>U: [] (no finding)
+    end
+```
+
+## Label taxonomy
+
+```mermaid
+mindmap
+  root((primary_vuln))
+    Safe
+      safe
+    Critical-leaning
+      reentrancy
+      access_control
+      tx_origin_auth
+      integer_overflow
+      unsafe_delegatecall
+    Medium / other
+      weak_randomness
+      unbounded_loop
+      other
+    Style / gas
+      redundant_storage
+      gas_optimization
+      best_practice
+```
+
 ## Intended use
 
 - Input: Solidity source code (string)
@@ -47,7 +181,20 @@ Part of [SolidityGuard](https://github.com/tanaymitra54/solidity_guard_razorpay)
 
 ## Labels
 
-`safe`, `reentrancy`, `access_control`, `tx_origin_auth`, `integer_overflow`, `unsafe_delegatecall`, `weak_randomness`, `unbounded_loop`, `redundant_storage`, `gas_optimization`, `best_practice`, `other`
+| ID | Label | Typical severity hint |
+|----|-------|------------------------|
+| 0 | `safe` | — |
+| 1 | `reentrancy` | Critical |
+| 2 | `access_control` | Critical |
+| 3 | `tx_origin_auth` | Critical |
+| 4 | `integer_overflow` | Critical |
+| 5 | `unsafe_delegatecall` | Critical |
+| 6 | `weak_randomness` | Medium |
+| 7 | `unbounded_loop` | Medium |
+| 8 | `redundant_storage` | Low |
+| 9 | `gas_optimization` | Low |
+| 10 | `best_practice` | Low |
+| 11 | `other` | Medium |
 
 ## Training data
 
@@ -66,6 +213,15 @@ Part of [SolidityGuard](https://github.com/tanaymitra54/solidity_guard_razorpay)
 | F1 `other` | 0.340 |
 | F1 `access_control` | 0.200 |
 
+```mermaid
+%%{init: {"theme": "neutral"}}%%
+xychart-beta
+    title "Per-class F1 on held-out test"
+    x-axis [safe, integer_overflow, reentrancy, other, access_control]
+    y-axis "F1" 0 --> 1
+    bar [0.694, 0.634, 0.461, 0.340, 0.200]
+```
+
 Labels are noisy (static-analysis consensus), so scores are moderate by design.
 
 ## Quick start
@@ -104,11 +260,29 @@ print(model.config.id2label[pred], float(probs[pred]))
 export GRAPHCODEBERT_PATH=tanaymitra01/graphcodebert-vulnerability-detector
 ```
 
+## Deploy / access map
+
+```mermaid
+flowchart TB
+    CKPT["Local checkpoint<br/>training/checkpoints/best"]
+    GH["GitHub LFS<br/>tanaymitra54/solidity_guard_razorpay"]
+    HF["Hugging Face Hub<br/>tanaymitra01/graphcodebert-vulnerability-detector"]
+
+    CKPT -->|git lfs push| GH
+    CKPT -->|hf upload| HF
+
+    HF --> COLAB["Colab / scripts"]
+    HF --> SPACE["HF Space / API host"]
+    HF --> LOCAL["Any machine<br/>from_pretrained"]
+    GH --> DEV["Dev clones"]
+```
+
 ## Files
 
 - `model.safetensors` — weights
 - `config.json` — RobertaForSequenceClassification config + label maps
 - `label_map.json` — label list / id maps used in training
+- `README.md` — this model card
 
 ## Limitations