sumitp76/cve-exploitability

1# Self-contained loader for the CVE exploitability model.2# Usage:3#   from huggingface_hub import hf_hub_download4#   import importlib.util, sys5#   spec = importlib.util.spec_from_file_location("hf_model", hf_hub_download("sumitp76/cve-exploitability","hf_model.py"))6#   m = importlib.util.module_from_spec(spec); sys.modules["hf_model"]=m; spec.loader.exec_module(m)7#   pp, net = m.load_model("sumitp76/cve-exploitability")8import re, json9import numpy as np10import torch, torch.nn as nn11 12_WORD = re.compile(r"[A-Za-z0-9_.\-]+")13CAT_COLS = ["severity","AV","AC","PR","UI","S","C","I","A"]14NUM_COLS = ["base_score","v2_base_score","has_v3","year","desc_len"]15 16def tokenize(t): return _WORD.findall(str(t).lower())17def _isnan(x):18    try: return isinstance(x, float) and np.isnan(x)19    except Exception: return False20 21class Preprocessor:22    def __init__(self, max_len=200, max_vocab=20000):23        self.max_len=max_len; self.max_vocab=max_vocab24        self.word2idx={}; self.cat_maps={}; self.cwe_vocab={}25        self.num_mean={}; self.num_std={}26    def transform_text(self, descriptions):27        X = np.zeros((len(descriptions), self.max_len), dtype=np.int64)28        for i,t in enumerate(descriptions):29            for j,w in enumerate(tokenize(t)[:self.max_len]):30                X[i,j] = self.word2idx.get(w,1)31        return X32    def transform_struct(self, df):33        n=len(df)34        def col(name, default):35            return df[name].values if name in df.columns else np.array([default]*n)36        num=np.zeros((n,len(NUM_COLS)),dtype=np.float32)37        desc_len=np.array([len(str(x)) for x in col("description","")],dtype=float)38        srcs={"base_score":col("base_score",np.nan).astype(float),39              "v2_base_score":col("v2_base_score",np.nan).astype(float),40              "has_v3":col("has_v3",0).astype(float),41              "year":col("year",2020).astype(float),42              "desc_len":desc_len}43        for k,name in enumerate(NUM_COLS):44            v=srcs[name].astype(float); v=np.where(np.isnan(v), self.num_mean[name], v)45            num[:,k]=(v-self.num_mean[name])/self.num_std[name]46        blocks=[num]47        for c in CAT_COLS:48            m=self.cat_maps[c]; width=len(m)+2; b=np.zeros((n,width),dtype=np.float32)49            for i,val in enumerate(col(c,None)):50                idx = m.get(val,1) if (val is not None and not _isnan(val)) else 051                b[i,idx]=1.052            blocks.append(b)53        cwe_w=len(self.cwe_vocab)+1; cb=np.zeros((n,cwe_w),dtype=np.float32)54        for i,val in enumerate(col("cwe","UNKNOWN")): cb[i,self.cwe_vocab.get(val,0)]=1.055        blocks.append(cb)56        return np.concatenate(blocks, axis=1)57    @property58    def struct_dim(self):59        return len(NUM_COLS)+sum(len(self.cat_maps[c])+2 for c in CAT_COLS)+len(self.cwe_vocab)+160    @property61    def vocab_size(self): return len(self.word2idx)62    @classmethod63    def from_dict(cls,d):64        p=cls(d["max_len"],d["max_vocab"]); p.word2idx=d["word2idx"]; p.cat_maps=d["cat_maps"]65        p.cwe_vocab=d["cwe_vocab"]; p.num_mean=d["num_mean"]; p.num_std=d["num_std"]; return p66 67class TextCNN(nn.Module):68    def __init__(self, vocab_size, emb_dim=128, kernels=(2,3,4,5), n_filters=96, dropout=0.3, out_dim=192):69        super().__init__()70        self.emb=nn.Embedding(vocab_size, emb_dim, padding_idx=0)71        self.convs=nn.ModuleList(nn.Conv1d(emb_dim,n_filters,k,padding=k//2) for k in kernels)72        self.act=nn.ReLU(); self.drop=nn.Dropout(dropout)73        self.proj=nn.Linear(n_filters*len(kernels), out_dim); self.out_dim=out_dim74    def forward(self,x):75        e=self.emb(x).transpose(1,2)76        feats=[self.act(c(e)).max(dim=2).values for c in self.convs]77        return self.proj(self.drop(torch.cat(feats,dim=1)))78 79class StructuredEncoder(nn.Module):80    def __init__(self, in_dim, hidden=128, out_dim=96, dropout=0.3):81        super().__init__()82        self.net=nn.Sequential(nn.Linear(in_dim,hidden),nn.ReLU(),nn.Dropout(dropout),83                               nn.Linear(hidden,out_dim),nn.ReLU()); self.out_dim=out_dim84    def forward(self,x): return self.net(x)85 86class ExploitabilityNet(nn.Module):87    def __init__(self, vocab_size, struct_dim, dropout=0.3):88        super().__init__()89        self.text=TextCNN(vocab_size,dropout=dropout)90        self.struct=StructuredEncoder(struct_dim,dropout=dropout)91        self.head=nn.Sequential(nn.Linear(self.text.out_dim+self.struct.out_dim,128),92                                nn.ReLU(),nn.Dropout(dropout),nn.Linear(128,1))93    def forward(self, text_ids, struct):94        return self.head(torch.cat([self.text(text_ids), self.struct(struct)],dim=1)).squeeze(-1)95 96def load_model(repo_id, device="cpu"):97    from huggingface_hub import hf_hub_download98    pp = Preprocessor.from_dict(json.load(open(hf_hub_download(repo_id,"preprocessor.json"))))99    net = ExploitabilityNet(pp.vocab_size, pp.struct_dim).to(device)100    net.load_state_dict(torch.load(hf_hub_download(repo_id,"model.pt"), map_location=device))101    net.eval()102    return pp, net103