Dockerfile

@@ -0,0 +1,34 @@
+FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PATH="/root/miniconda3/bin:${PATH}"
+ARG PATH="/root/miniconda3/bin:${PATH}"
+
+RUN apt-get update && apt-get install -y wget git libglib2.0-0 libsm6 libxext6 libxrender-dev
+
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-py38_4.10.3-Linux-x86_64.sh \
+    && mkdir /root/.conda \
+    && bash Miniconda3-py38_4.10.3-Linux-x86_64.sh -b \
+    && rm -f Miniconda3-py38_4.10.3-Linux-x86_64.sh
+
+RUN conda create -n cvss_env python=3.8 -y
+
+SHELL ["conda", "run", "-n", "cvss_env", "/bin/bash", "-c"]
+
+WORKDIR /app
+
+COPY requirements.txt .
+
+RUN pip install --no-cache-dir -r requirements.txt && \
+    pip install torch==2.3.0+cu121 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html && \
+    pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.3.0+cu121.html
+
+COPY . .
+
+RUN python -c "import nltk; nltk.download('punkt'); nltk.download('stopwords'); nltk.download('wordnet')"
+
+RUN chmod +x /app/entrypoint.sh
+
+EXPOSE 8888 6006
+
+ENTRYPOINT ["/app/entrypoint.sh"]

README.md

@@ -1,14 +0,0 @@
-## :rocket: 背景
-
-近年来，开源软件供应链遭受持续的软件投毒和恶意代码攻击，造成了无法估计的损失。例如，Apache Log4j2远程代码执行漏洞被认为是近10年最严重的漏洞之一，攻击者可以在目标服务器上执行任意代码和嗅探系统信息。网络安全专家认为Log4j 中的远程代码执行漏洞可能需要数月甚至数年时间才能得到妥善解决。受Log4J漏洞影响组件包括Apache的Struts2、Solr、Druid、Flink等，Github上60,644个开源项目发布321,094软件存在风险。因此，当前急需智能化技术辅助降低漏洞风险，提高漏洞工程能力，减少漏洞损失。
-
-
-## :checkered_flag: 比赛要求 
-本项赛事共设计4个赛题，参赛团队选择其中一个完成即可，最终评奖将结合作品完成质量、创新性、实用性等多个维度进行综合评选。
-
-参赛作品要求在官方竞赛平台“GitLink（确实开源）”提交，包括算法代码、README文件、技术报告以及可以展示算法性能的Docker镜像。具体要求参见赛事网站的“参赛指南”。
-
-注：推荐使用开源大模型。
-
-
-

README_V2.md

@@ -0,0 +1,80 @@
+# 漏洞严重等级(CVSS)研判项目
+
+## 文件结构
+```
+./
+├── data/
+│ ├── SIR_test_set.json
+│ ├── SIR_train_set.json
+│ └── SIR_validation_set.json
+├── src/
+│ ├── data_processing/
+│ │ ├── init.py
+│ │ ├── data_loader.py
+│ │ └── feature_engineering.py
+│ ├── models/
+│ │ ├── init.py
+│ │ ├── transformer_cvss.py
+│ │ ├── graph_neural_network.py
+│ │ └── ensemble_model.py
+│ ├── utils/
+│ │ ├── init.py
+│ │ ├── text_preprocessing.py
+│ │ └── visualization.py
+│ └── main.py
+├── notebooks/
+│ ├── exploratory_data_analysis.ipynb
+│ └── model_evaluation.ipynb
+├── tests/
+│ ├── init.py
+│ ├── test_data_processing.py
+│ └── test_models.py
+├── config.yaml
+├── requirements.txt
+└── README_V2.md
+```
+
+## 配置运行环境
+
+   ```
+   pip install -r requirements.txt
+   ```
+
+## 运行代码
+
+1. 数据预处理和特征工程:
+   ```
+   python src/data_processing/data_loader.py
+   python src/data_processing/feature_engineering.py
+   ```
+
+2. 训练模型:
+   ```
+   python src/main.py
+   ```
+
+3. 运行测试:
+   ```
+   python -m unittest discover tests
+   ```
+
+4. 查看结果:
+   - 模型训练和评估结果将在控制台输出。
+   - 可视化结果将保存在 `results/` 目录下。
+
+
+## 得出报告中呈现的结果
+
+1. 运行 `src/main.py` 脚本将生成主要的评估指标,包括 MSE 和 RMSE。
+
+2. 查看 `notebooks/model_evaluation.ipynb` 获取详细的模型性能分析和可视化结果。
+
+3. 检查 `results/` 目录下的图表,包括特征重要性、学习曲线、预测vs实际值等。
+
+4. 最终报告应包括:
+   - 数据集描述和预处理步骤
+   - 模型架构和训练过程
+   - 评估指标(MSE, RMSE)及其解释
+   - 特征重要性分析
+   - 模型性能可视化
+   - 结论和未来改进方向

config.yaml

@@ -0,0 +1,34 @@
+data:
+  train_path: 'data/SIR_train_set.json'
+  test_path: 'data/SIR_test_set.json'
+  validation_path: 'data/SIR_validation_set.json'
+
+model:
+  transformer:
+    bert_model_name: 'bert-base-uncased'
+    num_labels: 3
+    learning_rate: 2e-5
+    batch_size: 32
+    epochs: 10
+  gnn:
+    num_node_features: 300
+    num_classes: 3
+    learning_rate: 0.01
+    batch_size: 32
+    epochs: 10
+  ensemble:
+    n_estimators: 100
+    random_state: 42
+
+feature_engineering:
+  tfidf_max_features: 1000
+  ngram_range: [1, 2]
+
+visualization:
+  figure_size: [10, 6]
+  style: 'seaborn'
+
+evaluation:
+  test_size: 0.2
+  validation_size: 0.1
+  cv_folds: 5

data/label_word_ids.json

@@ -0,0 +1,46 @@
+{
+    "AV": {
+        "network": 2897,
+        "adjacent": 5516,
+        "local": 2334,
+        "physical": 3558
+    },
+    "AC": {
+        "low": 2659,
+        "high": 2152
+    },
+    "PR": {
+        "none": 3904,
+        "low": 2659,
+        "high": 2152
+    },
+    "UI": {
+        "none": 3904,
+        "required": 3223
+    },
+    "S": {
+        "unchanged": 15704,
+        "changed": 2904
+    },
+    "C": {
+        "none": 3904,
+        "low": 2659,
+        "high": 2152
+    },
+    "I": {
+        "none": 3904,
+        "low": 2659,
+        "high": 2152
+    },
+    "A": {
+        "none": 3904,
+        "low": 2659,
+        "high": 2152
+    },
+    "severity": {
+        "low": 2659,
+        "medium": 5396,
+        "high": 2152,
+        "critical": 4187
+    }
+}

data/label_word_ids_CVSS2.json

@@ -0,0 +1,37 @@
+{
+    "AV": {
+        "network": 2897,
+        "adjacent": 5516,
+        "local": 2334
+    },
+    "AC": {
+        "low": 2659,
+        "medium": 5396,
+        "high": 2152
+    },
+    "Au": {
+        "none": 3904,
+        "single": 2309,
+        "multiple": 3674
+    },
+    "C": {
+        "none": 3904,
+        "partial": 7704,
+        "complete": 3143
+    },
+    "I": {
+        "none": 3904,
+        "partial": 7704,
+        "complete": 3143
+    },
+    "A": {
+        "none": 3904,
+        "partial": 7704,
+        "complete": 3143
+    },
+    "severity": {
+        "low": 2659,
+        "medium": 5396,
+        "high": 2152
+    }
+}

requirements.txt

@@ -0,0 +1,10 @@
+numpy==1.21.0
+pandas==1.3.0
+scikit-learn==0.24.2
+torch==1.9.0
+transformers==4.9.2
+matplotlib==3.4.2
+seaborn==0.11.1
+wordcloud==1.8.1
+nltk==3.6.2
+torch-geometric==2.0.1

src/data_processing/__init__.py

@@ -0,0 +1,2 @@
+from .data_loader import DataLoader
+from .feature_engineering import FeatureEngineer

src/data_processing/data_loader.py

@@ -0,0 +1,50 @@
+import json
+import pandas as pd
+import numpy as np
+from typing import Dict, List, Union
+
+class DataLoader:
+    def __init__(self, file_paths: Dict[str, str]):
+        self.file_paths = file_paths
+        self.data = {}
+
+    def load_data(self) -> Dict[str, pd.DataFrame]:
+        for key, path in self.file_paths.items():
+            with open(path, 'r', encoding='utf-8') as file:
+                json_data = json.load(file)
+                df = pd.json_normalize(json_data)
+                self.data[key] = df
+        return self.data
+
+    def preprocess_data(self) -> Dict[str, pd.DataFrame]:
+        for key, df in self.data.items():
+            df['description'] = df['description'].fillna('')
+            df['vectorString'] = df['vectorString'].fillna('')
+            df['severity'] = df['severity'].fillna('UNKNOWN')
+            df['baseScore'] = df['baseScore'].fillna(0)
+            df['impactScore'] = df['impactScore'].fillna(0)
+            df['exploitabilityScore'] = df['exploitabilityScore'].fillna(0)
+            self.data[key] = df
+        return self.data
+
+    def split_data(self, test_size: float = 0.2, validation_size: float = 0.1) -> Dict[str, pd.DataFrame]:
+        train_data = self.data['train']
+        test_data = self.data['test']
+        validation_data = self.data['validation']
+
+        if 'validation' not in self.data:
+            train_val_data = train_data.sample(frac=1, random_state=42)
+            val_size = int(len(train_val_data) * validation_size)
+            validation_data = train_val_data[:val_size]
+            train_data = train_val_data[val_size:]
+
+        return {
+            'train': train_data,
+            'test': test_data,
+            'validation': validation_data
+        }
+
+    def get_feature_target_split(self, data: pd.DataFrame) -> Tuple[pd.DataFrame, pd.Series]:
+        features = data.drop(['baseScore', 'impactScore', 'exploitabilityScore'], axis=1)
+        targets = data[['baseScore', 'impactScore', 'exploitabilityScore']]
+        return features, targets

src/data_processing/feature_engineering.py

@@ -0,0 +1,44 @@
+import pandas as pd
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.preprocessing import LabelEncoder
+from typing import Dict, List, Tuple
+
+class FeatureEngineer:
+    def __init__(self):
+        self.tfidf_vectorizer = TfidfVectorizer(max_features=1000, stop_words='english')
+        self.label_encoder = LabelEncoder()
+
+    def engineer_features(self, data: Dict[str, pd.DataFrame]) -> Dict[str, Tuple[pd.DataFrame, pd.DataFrame]]:
+        engineered_data = {}
+        for key, df in data.items():
+            features = self._extract_text_features(df['description'])
+            features = pd.concat([features, self._encode_categorical(df['severity'])], axis=1)
+            features = pd.concat([features, self._parse_vector_string(df['vectorString'])], axis=1)
+            targets = df[['baseScore', 'impactScore', 'exploitabilityScore']]
+            engineered_data[key] = (features, targets)
+        return engineered_data
+
+    def _extract_text_features(self, text_series: pd.Series) -> pd.DataFrame:
+        if not hasattr(self, 'tfidf_matrix'):
+            self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(text_series)
+        else:
+            self.tfidf_matrix = self.tfidf_vectorizer.transform(text_series)
+        return pd.DataFrame(self.tfidf_matrix.toarray(), columns=self.tfidf_vectorizer.get_feature_names_out())
+
+    def _encode_categorical(self, cat_series: pd.Series) -> pd.DataFrame:
+        if not hasattr(self, 'label_encoder_classes'):
+            encoded = self.label_encoder.fit_transform(cat_series)
+            self.label_encoder_classes = self.label_encoder.classes_
+        else:
+            encoded = self.label_encoder.transform(cat_series)
+        return pd.DataFrame({'severity_encoded': encoded})
+
+    def _parse_vector_string(self, vector_series: pd.Series) -> pd.DataFrame:
+        vector_df = vector_series.str.extract(r'AV:(?P<AV>\w)/AC:(?P<AC>\w)/PR:(?P<PR>\w)/UI:(?P<UI>\w)/S:(?P<S>\w)/C:(?P<C>\w)/I:(?P<I>\w)/A:(?P<A>\w)')
+        for col in vector_df.columns:
+            vector_df[col] = vector_df[col].map({'N': 0, 'L': 1, 'H': 2, 'C': 3})
+        return vector_df.fillna(0)
+
+    def inverse_transform_targets(self, predictions: np.ndarray) -> pd.DataFrame:
+        return pd.DataFrame(predictions, columns=['baseScore', 'impactScore', 'exploitabilityScore'])

src/main.py

@@ -0,0 +1,50 @@
+import pandas as pd
+import numpy as np
+from sklearn.model_selection import train_test_split
+from data_processing.data_loader import DataLoader
+from data_processing.feature_engineering import FeatureEngineer
+from models.transformer_cvss import TransformerCVSS
+from models.graph_neural_network import GraphNeuralNetwork
+from models.ensemble_model import EnsembleModel
+from utils.text_preprocessing import TextPreprocessor
+from utils.visualization import Visualizer
+
+def main():
+    data_loader = DataLoader({
+        'train': 'data/SIR_train_set.json',
+        'test': 'data/SIR_test_set.json',
+        'validation': 'data/SIR_validation_set.json'
+    })
+    data = data_loader.load_data()
+    data = data_loader.preprocess_data()
+
+    feature_engineer = FeatureEngineer()
+    engineered_data = feature_engineer.engineer_features(data)
+
+    transformer_model = TransformerCVSS()
+    gnn_model = GraphNeuralNetwork(num_node_features=300, num_classes=3)
+
+    X_train, y_train = engineered_data['train']
+    X_val, y_val = engineered_data['validation']
+    X_test, y_test = engineered_data['test']
+
+    transformer_model.train_model(X_train, X_val)
+    gnn_model.train_model(X_train, X_val)
+
+    ensemble_model = EnsembleModel(transformer_model, gnn_model)
+    ensemble_model.train(X_train, y_train)
+
+    test_predictions = ensemble_model.predict(X_test)
+    test_evaluation = ensemble_model.evaluate(X_test, y_test)
+
+    print("Test Evaluation Results:")
+    print(f"MSE: {test_evaluation['MSE']:.4f}")
+    print(f"RMSE: {test_evaluation['RMSE']:.4f}")
+
+    visualizer = Visualizer()
+    visualizer.plot_score_distribution(data['train'], 'baseScore', 'Base Score Distribution')
+    visualizer.plot_feature_importance(ensemble_model.feature_importance(), X_train.columns)
+    visualizer.plot_prediction_vs_actual(y_test.values, test_predictions)
+
+if __name__ == "__main__":
+    main()

src/models/__init__.py

@@ -0,0 +1,3 @@
+from .transformer_cvss import TransformerCVSS
+from .graph_neural_network import GraphNeuralNetwork
+from .ensemble_model import EnsembleModel

src/models/ensemble_model.py

@@ -0,0 +1,49 @@
+import numpy as np
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.model_selection import cross_val_score
+from .transformer_cvss import TransformerCVSS
+from .graph_neural_network import GraphNeuralNetwork
+
+class EnsembleModel:
+    def __init__(self, transformer_model, gnn_model):
+        self.transformer_model = transformer_model
+        self.gnn_model = gnn_model
+        self.rf_model = RandomForestRegressor(n_estimators=100, random_state=42)
+
+    def train(self, X_train, y_train):
+        transformer_preds = np.array([self.transformer_model.predict(text) for text in X_train['description']])
+        gnn_preds = np.array([self.gnn_model.predict(text) for text in X_train['description']])
+        combined_features = np.hstack([transformer_preds, gnn_preds, X_train.drop('description', axis=1).values])
+        self.rf_model.fit(combined_features, y_train)
+
+    def predict(self, X_test):
+        transformer_preds = np.array([self.transformer_model.predict(text) for text in X_test['description']])
+        gnn_preds = np.array([self.gnn_model.predict(text) for text in X_test['description']])
+        combined_features = np.hstack([transformer_preds, gnn_preds, X_test.drop('description', axis=1).values])
+        return self.rf_model.predict(combined_features)
+
+    def evaluate(self, X_val, y_val):
+        predictions = self.predict(X_val)
+        mse = np.mean((predictions - y_val) ** 2)
+        rmse = np.sqrt(mse)
+        return {'MSE': mse, 'RMSE': rmse}
+
+    def cross_validate(self, X, y, cv=5):
+        transformer_preds = np.array([self.transformer_model.predict(text) for text in X['description']])
+        gnn_preds = np.array([self.gnn_model.predict(text) for text in X['description']])
+        combined_features = np.hstack([transformer_preds, gnn_preds, X.drop('description', axis=1).values])
+        
+        mse_scores = cross_val_score(self.rf_model, combined_features, y, cv=cv, scoring='neg_mean_squared_error')
+        rmse_scores = np.sqrt(-mse_scores)
+        
+        return {
+            'MSE': -mse_scores.mean(),
+            'RMSE': rmse_scores.mean(),
+            'MSE_std': mse_scores.std(),
+            'RMSE_std': rmse_scores.std()
+        }
+
+    def feature_importance(self):
+        feature_importance = self.rf_model.feature_importances_
+        feature_names = [f'transformer_{i}' for i in range(3)] + [f'gnn_{i}' for i in range(3)] + list(X.drop('description', axis=1).columns)
+        return dict(zip(feature_names, feature_importance))

src/models/graph_neural_network.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import GCNConv, global_mean_pool
+from torch_geometric.data import Data, Batch
+import numpy as np
+
+class GraphNeuralNetwork(nn.Module):
+    def __init__(self, num_node_features, num_classes):
+        super(GraphNeuralNetwork, self).__init__()
+        self.conv1 = GCNConv(num_node_features, 64)
+        self.conv2 = GCNConv(64, 64)
+        self.fc = nn.Linear(64, num_classes)
+
+    def forward(self, data):
+        x, edge_index, batch = data.x, data.edge_index, data.batch
+        x = F.relu(self.conv1(x, edge_index))
+        x = F.relu(self.conv2(x, edge_index))
+        x = global_mean_pool(x, batch)
+        x = self.fc(x)
+        return x
+
+    def create_graph_from_text(self, text):
+        words = text.split()
+        num_nodes = len(words)
+        edge_index = []
+        for i in range(num_nodes):
+            for j in range(max(0, i-5), min(num_nodes, i+6)):
+                if i != j:
+                    edge_index.append([i, j])
+        edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
+        x = torch.randn(num_nodes, 300)
+        return Data(x=x, edge_index=edge_index)
+
+    def predict(self, text):
+        self.eval()
+        with torch.no_grad():
+            graph = self.create_graph_from_text(text)
+            data = Batch.from_data_list([graph])
+            output = self(data)
+        return output.squeeze().numpy()
+
+    def train_model(self, train_data, val_data, epochs=10, batch_size=32, learning_rate=0.01):
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.to(device)
+        optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)
+        criterion = nn.MSELoss()
+
+        for epoch in range(epochs):
+            self.train()
+            total_loss = 0
+            for i in range(0, len(train_data), batch_size):
+                batch_texts = train_data['description'][i:i+batch_size].tolist()
+                batch_scores = train_data[['baseScore', 'impactScore', 'exploitabilityScore']][i:i+batch_size].values
+                batch_graphs = [self.create_graph_from_text(text) for text in batch_texts]
+                batch_data = Batch.from_data_list(batch_graphs).to(device)
+                labels = torch.tensor(batch_scores, dtype=torch.float).to(device)
+
+                optimizer.zero_grad()
+                outputs = self(batch_data)
+                loss = criterion(outputs, labels)
+                loss.backward()
+                optimizer.step()
+                total_loss += loss.item()
+
+            avg_train_loss = total_loss / (len(train_data) // batch_size)
+            val_loss = self.evaluate(val_data)
+            print(f'Epoch {epoch+1}/{epochs}, Train Loss: {avg_train_loss:.4f}, Val Loss: {val_loss:.4f}')
+
+    def evaluate(self, val_data):
+        self.eval()
+        criterion = nn.MSELoss()
+        device = next(self.parameters()).device
+        total_loss = 0
+        with torch.no_grad():
+            for i in range(len(val_data)):
+                text = val_data['description'][i]
+                scores = val_data[['baseScore', 'impactScore', 'exploitabilityScore']].iloc[i].values
+                graph = self.create_graph_from_text(text)
+                data = Batch.from_data_list([graph]).to(device)
+                labels = torch.tensor(scores, dtype=torch.float).to(device)
+
+                outputs = self(data)
+                loss = criterion(outputs.squeeze(), labels)
+                total_loss += loss.item()
+        return total_loss / len(val_data)

src/models/transformer_cvss.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from transformers import BertModel, BertTokenizer
+import numpy as np
+
+class TransformerCVSS(nn.Module):
+    def __init__(self, bert_model_name='bert-base-uncased', num_labels=3):
+        super(TransformerCVSS, self).__init__()
+        self.bert = BertModel.from_pretrained(bert_model_name)
+        self.dropout = nn.Dropout(0.1)
+        self.fc = nn.Linear(self.bert.config.hidden_size, num_labels)
+        self.tokenizer = BertTokenizer.from_pretrained(bert_model_name)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs[1]
+        x = self.dropout(pooled_output)
+        x = self.fc(x)
+        return x
+
+    def predict(self, text):
+        self.eval()
+        with torch.no_grad():
+            inputs = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)
+            outputs = self(inputs['input_ids'], inputs['attention_mask'])
+        return outputs.squeeze().numpy()
+
+    def train_model(self, train_data, val_data, epochs=10, batch_size=32, learning_rate=2e-5):
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.to(device)
+        optimizer = optim.AdamW(self.parameters(), lr=learning_rate)
+        criterion = nn.MSELoss()
+
+        for epoch in range(epochs):
+            self.train()
+            total_loss = 0
+            for i in range(0, len(train_data), batch_size):
+                batch_texts = train_data['description'][i:i+batch_size].tolist()
+                batch_scores = train_data[['baseScore', 'impactScore', 'exploitabilityScore']][i:i+batch_size].values
+                inputs = self.tokenizer(batch_texts, return_tensors='pt', padding=True, truncation=True, max_length=512)
+                input_ids = inputs['input_ids'].to(device)
+                attention_mask = inputs['attention_mask'].to(device)
+                labels = torch.tensor(batch_scores, dtype=torch.float).to(device)
+
+                optimizer.zero_grad()
+                outputs = self(input_ids, attention_mask)
+                loss = criterion(outputs, labels)
+                loss.backward()
+                optimizer.step()
+                total_loss += loss.item()
+
+            avg_train_loss = total_loss / (len(train_data) // batch_size)
+            val_loss = self.evaluate(val_data)
+            print(f'Epoch {epoch+1}/{epochs}, Train Loss: {avg_train_loss:.4f}, Val Loss: {val_loss:.4f}')
+
+    def evaluate(self, val_data):
+        self.eval()
+        criterion = nn.MSELoss()
+        device = next(self.parameters()).device
+        total_loss = 0
+        with torch.no_grad():
+            for i in range(len(val_data)):
+                text = val_data['description'][i]
+                scores = val_data[['baseScore', 'impactScore', 'exploitabilityScore']].iloc[i].values
+                inputs = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)
+                input_ids = inputs['input_ids'].to(device)
+                attention_mask = inputs['attention_mask'].to(device)
+                labels = torch.tensor(scores, dtype=torch.float).to(device)
+
+                outputs = self(input_ids, attention_mask)
+                loss = criterion(outputs.squeeze(), labels)
+                total_loss += loss.item()
+        return total_loss / len(val_data)

src/utils/__init__.py

@@ -0,0 +1,2 @@
+from .text_preprocessing import TextPreprocessor
+from .visualization import Visualizer

src/utils/text_preprocessing.py

@@ -0,0 +1,71 @@
+import re
+import nltk
+from nltk.corpus import stopwords
+from nltk.tokenize import word_tokenize
+from nltk.stem import WordNetLemmatizer
+import string
+
+class TextPreprocessor:
+    def __init__(self):
+        nltk.download('punkt', quiet=True)
+        nltk.download('stopwords', quiet=True)
+        nltk.download('wordnet', quiet=True)
+        self.stop_words = set(stopwords.words('english'))
+        self.lemmatizer = WordNetLemmatizer()
+
+    def preprocess(self, text):
+        text = self._remove_html_tags(text)
+        text = self._remove_urls(text)
+        text = self._remove_special_characters(text)
+        text = self._convert_to_lowercase(text)
+        tokens = self._tokenize(text)
+        tokens = self._remove_stopwords(tokens)
+        tokens = self._lemmatize(tokens)
+        return ' '.join(tokens)
+
+    def _remove_html_tags(self, text):
+        return re.sub(r'<.*?>', '', text)
+
+    def _remove_urls(self, text):
+        return re.sub(r'http\S+|www.\S+', '', text)
+
+    def _remove_special_characters(self, text):
+        return re.sub(r'[^\w\s]', '', text)
+
+    def _convert_to_lowercase(self, text):
+        return text.lower()
+
+    def _tokenize(self, text):
+        return word_tokenize(text)
+
+    def _remove_stopwords(self, tokens):
+        return [token for token in tokens if token not in self.stop_words]
+
+    def _lemmatize(self, tokens):
+        return [self.lemmatizer.lemmatize(token) for token in tokens]
+
+    def extract_keywords(self, text, top_n=10):
+        processed_text = self.preprocess(text)
+        words = processed_text.split()
+        word_freq = {}
+        for word in words:
+            if word in word_freq:
+                word_freq[word] += 1
+            else:
+                word_freq[word] = 1
+        sorted_words = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)
+        return [word for word, freq in sorted_words[:top_n]]
+
+    def calculate_text_complexity(self, text):
+        words = text.split()
+        sentences = text.split('.')
+        avg_word_length = sum(len(word) for word in words) / len(words)
+        avg_sentence_length = len(words) / len(sentences)
+        unique_words = len(set(words))
+        complexity_score = (avg_word_length * 0.5) + (avg_sentence_length * 0.5) + (unique_words * 0.1)
+        return complexity_score
+
+    def generate_ngrams(self, text, n=2):
+        words = text.split()
+        ngrams = zip(*[words[i:] for i in range(n)])
+        return [' '.join(ngram) for ngram in ngrams]

src/utils/visualization.py

@@ -0,0 +1,92 @@
+import matplotlib.pyplot as plt
+import seaborn as sns
+import pandas as pd
+import numpy as np
+from wordcloud import WordCloud
+from sklearn.manifold import TSNE
+from sklearn.decomposition import PCA
+
+class Visualizer:
+    def __init__(self):
+        plt.style.use('seaborn')
+
+    def plot_score_distribution(self, data, score_column, title):
+        plt.figure(figsize=(10, 6))
+        sns.histplot(data=data, x=score_column, kde=True)
+        plt.title(title)
+        plt.xlabel('Score')
+        plt.ylabel('Frequency')
+        plt.show()
+
+    def plot_correlation_heatmap(self, data):
+        plt.figure(figsize=(12, 10))
+        sns.heatmap(data.corr(), annot=True, cmap='coolwarm', linewidths=0.5)
+        plt.title('Correlation Heatmap')
+        plt.show()
+
+    def plot_feature_importance(self, feature_importance, feature_names):
+        importance_df = pd.DataFrame({'feature': feature_names, 'importance': feature_importance})
+        importance_df = importance_df.sort_values('importance', ascending=False)
+        plt.figure(figsize=(12, 8))
+        sns.barplot(x='importance', y='feature', data=importance_df)
+        plt.title('Feature Importance')
+        plt.xlabel('Importance')
+        plt.ylabel('Feature')
+        plt.show()
+
+    def plot_learning_curve(self, train_scores, val_scores, metric_name):
+        plt.figure(figsize=(10, 6))
+        plt.plot(range(1, len(train_scores) + 1), train_scores, label='Train')
+        plt.plot(range(1, len(val_scores) + 1), val_scores, label='Validation')
+        plt.title(f'Learning Curve - {metric_name}')
+        plt.xlabel('Epoch')
+        plt.ylabel(metric_name)
+        plt.legend()
+        plt.show()
+
+    def plot_confusion_matrix(self, cm, classes):
+        plt.figure(figsize=(10, 8))
+        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=classes, yticklabels=classes)
+        plt.title('Confusion Matrix')
+        plt.xlabel('Predicted')
+        plt.ylabel('Actual')
+        plt.show()
+
+    def plot_wordcloud(self, text_data):
+        wordcloud = WordCloud(width=800, height=400, background_color='white').generate(' '.join(text_data))
+        plt.figure(figsize=(10, 5))
+        plt.imshow(wordcloud, interpolation='bilinear')
+        plt.axis('off')
+        plt.title('Word Cloud')
+        plt.show()
+
+    def plot_tsne_clusters(self, features, labels):
+        tsne = TSNE(n_components=2, random_state=42)
+        tsne_results = tsne.fit_transform(features)
+        plt.figure(figsize=(10, 8))
+        scatter = plt.scatter(tsne_results[:, 0], tsne_results[:, 1], c=labels, cmap='viridis')
+        plt.colorbar(scatter)
+        plt.title('t-SNE Visualization of Clusters')
+        plt.xlabel('t-SNE 1')
+        plt.ylabel('t-SNE 2')
+        plt.show()
+
+    def plot_pca_explained_variance(self, features):
+        pca = PCA()
+        pca.fit(features)
+        cumulative_variance_ratio = np.cumsum(pca.explained_variance_ratio_)
+        plt.figure(figsize=(10, 6))
+        plt.plot(range(1, len(cumulative_variance_ratio) + 1), cumulative_variance_ratio, 'bo-')
+        plt.title('Cumulative Explained Variance Ratio')
+        plt.xlabel('Number of Components')
+        plt.ylabel('Cumulative Explained Variance Ratio')
+        plt.show()
+
+    def plot_prediction_vs_actual(self, y_true, y_pred):
+        plt.figure(figsize=(10, 6))
+        plt.scatter(y_true, y_pred, alpha=0.5)
+        plt.plot([y_true.min(), y_true.max()], [y_true.min(), y_true.max()], 'r--', lw=2)
+        plt.xlabel('Actual Values')
+        plt.ylabel('Predicted Values')
+        plt.title('Prediction vs Actual')
+        plt.show()

tests/__init__.py

@@ -0,0 +1,2 @@
+from .test_data_processing import TestDataProcessing
+from .test_models import TestModels

tests/test_data_processing.py

@@ -0,0 +1,103 @@
+import unittest
+import pandas as pd
+import numpy as np
+from src.data_processing.data_loader import DataLoader
+from src.data_processing.feature_engineering import FeatureEngineer
+
+class TestDataProcessing(unittest.TestCase):
+    def setUp(self):
+        self.data_loader = DataLoader({
+            'train': 'data/SIR_train_set.json',
+            'test': 'data/SIR_test_set.json',
+            'validation': 'data/SIR_validation_set.json'
+        })
+        self.feature_engineer = FeatureEngineer()
+
+    def test_data_loader(self):
+        data = self.data_loader.load_data()
+        self.assertIsInstance(data, dict)
+        self.assertIn('train', data)
+        self.assertIn('test', data)
+        self.assertIn('validation', data)
+        self.assertIsInstance(data['train'], pd.DataFrame)
+        self.assertIsInstance(data['test'], pd.DataFrame)
+        self.assertIsInstance(data['validation'], pd.DataFrame)
+
+    def test_data_preprocessing(self):
+        data = self.data_loader.load_data()
+        preprocessed_data = self.data_loader.preprocess_data()
+        self.assertEqual(len(data), len(preprocessed_data))
+        for key in data:
+            self.assertFalse(preprocessed_data[key]['description'].isnull().any())
+            self.assertFalse(preprocessed_data[key]['vectorString'].isnull().any())
+            self.assertFalse(preprocessed_data[key]['severity'].isnull().any())
+            self.assertFalse(preprocessed_data[key]['baseScore'].isnull().any())
+            self.assertFalse(preprocessed_data[key]['impactScore'].isnull().any())
+            self.assertFalse(preprocessed_data[key]['exploitabilityScore'].isnull().any())
+
+    def test_feature_engineering(self):
+        data = self.data_loader.load_data()
+        preprocessed_data = self.data_loader.preprocess_data()
+        engineered_data = self.feature_engineer.engineer_features(preprocessed_data)
+        self.assertIsInstance(engineered_data, dict)
+        for key in engineered_data:
+            features, targets = engineered_data[key]
+            self.assertIsInstance(features, pd.DataFrame)
+            self.assertIsInstance(targets, pd.DataFrame)
+            self.assertEqual(len(features), len(targets))
+            self.assertGreater(features.shape[1], 0)
+            self.assertEqual(targets.shape[1], 3)
+
+    def test_data_split(self):
+        data = self.data_loader.load_data()
+        preprocessed_data = self.data_loader.preprocess_data()
+        split_data = self.data_loader.split_data()
+        self.assertIn('train', split_data)
+        self.assertIn('test', split_data)
+        self.assertIn('validation', split_data)
+        self.assertGreater(len(split_data['train']), len(split_data['validation']))
+        self.assertGreater(len(split_data['train']), len(split_data['test']))
+
+    def test_feature_target_split(self):
+        data = self.data_loader.load_data()
+        preprocessed_data = self.data_loader.preprocess_data()
+        split_data = self.data_loader.split_data()
+        features, targets = self.data_loader.get_feature_target_split(split_data['train'])
+        self.assertIsInstance(features, pd.DataFrame)
+        self.assertIsInstance(targets, pd.Series)
+        self.assertEqual(len(features), len(targets))
+        self.assertNotIn('baseScore', features.columns)
+        self.assertNotIn('impactScore', features.columns)
+        self.assertNotIn('exploitabilityScore', features.columns)
+
+    def test_tfidf_vectorization(self):
+        data = self.data_loader.load_data()
+        preprocessed_data = self.data_loader.preprocess_data()
+        engineered_data = self.feature_engineer.engineer_features(preprocessed_data)
+        for key in engineered_data:
+            features, _ = engineered_data[key]
+            self.assertIn('tfidf_0', features.columns)
+            self.assertGreaterEqual(features.filter(regex='^tfidf_').shape[1], 100)
+
+    def test_categorical_encoding(self):
+        data = self.data_loader.load_data()
+        preprocessed_data = self.data_loader.preprocess_data()
+        engineered_data = self.feature_engineer.engineer_features(preprocessed_data)
+        for key in engineered_data:
+            features, _ = engineered_data[key]
+            self.assertIn('severity_encoded', features.columns)
+            self.assertEqual(features['severity_encoded'].dtype, np.int64)
+
+    def test_vector_string_parsing(self):
+        data = self.data_loader.load_data()
+        preprocessed_data = self.data_loader.preprocess_data()
+        engineered_data = self.feature_engineer.engineer_features(preprocessed_data)
+        for key in engineered_data:
+            features, _ = engineered_data[key]
+            vector_columns = ['AV', 'AC', 'PR', 'UI', 'S', 'C', 'I', 'A']
+            for col in vector_columns:
+                self.assertIn(col, features.columns)
+                self.assertEqual(features[col].dtype, np.float64)
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_models.py

@@ -0,0 +1,86 @@
+import unittest
+import torch
+import pandas as pd
+import numpy as np
+from src.models.transformer_cvss import TransformerCVSS
+from src.models.graph_neural_network import GraphNeuralNetwork
+from src.models.ensemble_model import EnsembleModel
+
+class TestModels(unittest.TestCase):
+    def setUp(self):
+        self.transformer_model = TransformerCVSS()
+        self.gnn_model = GraphNeuralNetwork(num_node_features=300, num_classes=3)
+        self.ensemble_model = EnsembleModel(self.transformer_model, self.gnn_model)
+        self.sample_data = pd.DataFrame({
+            'description': ['This is a sample vulnerability description'] * 10,
+            'baseScore': np.random.rand(10),
+            'impactScore': np.random.rand(10),
+            'exploitabilityScore': np.random.rand(10)
+        })
+
+    def test_transformer_model(self):
+        self.assertIsInstance(self.transformer_model, torch.nn.Module)
+        output = self.transformer_model.predict(self.sample_data['description'][0])
+        self.assertIsInstance(output, np.ndarray)
+        self.assertEqual(output.shape, (3,))
+
+    def test_gnn_model(self):
+        self.assertIsInstance(self.gnn_model, torch.nn.Module)
+        output = self.gnn_model.predict(self.sample_data['description'][0])
+        self.assertIsInstance(output, np.ndarray)
+        self.assertEqual(output.shape, (3,))
+
+    def test_ensemble_model(self):
+        self.ensemble_model.train(self.sample_data, self.sample_data[['baseScore', 'impactScore', 'exploitabilityScore']])
+        predictions = self.ensemble_model.predict(self.sample_data)
+        self.assertIsInstance(predictions, np.ndarray)
+        self.assertEqual(predictions.shape, (10, 3))
+
+    def test_model_evaluation(self):
+        self.ensemble_model.train(self.sample_data, self.sample_data[['baseScore', 'impactScore', 'exploitabilityScore']])
+        evaluation = self.ensemble_model.evaluate(self.sample_data, self.sample_data[['baseScore', 'impactScore', 'exploitabilityScore']])
+        self.assertIn('MSE', evaluation)
+        self.assertIn('RMSE', evaluation)
+        self.assertIsInstance(evaluation['MSE'], float)
+        self.assertIsInstance(evaluation['RMSE'], float)
+
+    def test_cross_validation(self):
+        cv_results = self.ensemble_model.cross_validate(self.sample_data, self.sample_data[['baseScore', 'impactScore', 'exploitabilityScore']])
+        self.assertIn('MSE', cv_results)
+        self.assertIn('RMSE', cv_results)
+        self.assertIn('MSE_std', cv_results)
+        self.assertIn('RMSE_std', cv_results)
+        self.assertIsInstance(cv_results['MSE'], float)
+        self.assertIsInstance(cv_results['RMSE'], float)
+        self.assertIsInstance(cv_results['MSE_std'], float)
+        self.assertIsInstance(cv_results['RMSE_std'], float)
+
+    def test_transformer_model_training(self):
+        self.transformer_model.train_model(self.sample_data, self.sample_data)
+        self.assertTrue(hasattr(self.transformer_model, 'bert'))
+        self.assertTrue(hasattr(self.transformer_model, 'fc'))
+
+    def test_gnn_model_training(self):
+        self.gnn_model.train_model(self.sample_data, self.sample_data)
+        self.assertTrue(hasattr(self.gnn_model, 'conv1'))
+        self.assertTrue(hasattr(self.gnn_model, 'conv2'))
+        self.assertTrue(hasattr(self.gnn_model, 'fc'))
+
+    def test_ensemble_model_feature_importance(self):
+        self.ensemble_model.train(self.sample_data, self.sample_data[['baseScore', 'impactScore', 'exploitabilityScore']])
+        feature_importance = self.ensemble_model.feature_importance()
+        self.assertIsInstance(feature_importance, dict)
+        self.assertGreater(len(feature_importance), 0)
+        for feature, importance in feature_importance.items():
+            self.assertIsInstance(importance, float)
+            self.assertGreaterEqual(importance, 0)
+            self.assertLessEqual(importance, 1)
+
+    def test_model_prediction_range(self):
+        self.ensemble_model.train(self.sample_data, self.sample_data[['baseScore', 'impactScore', 'exploitabilityScore']])
+        predictions = self.ensemble_model.predict(self.sample_data)
+        self.assertTrue(np.all(predictions >= 0))
+        self.assertTrue(np.all(predictions <= 10))
+
+if __name__ == '__main__':
+    unittest.main()