diff --git a/README.md b/README.md old mode 100644 new mode 100755 index 96278e9..e69de29 --- a/README.md +++ b/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镜像。具体要求参见赛事网站的“参赛指南”。 - -注：推荐使用开源大模型。 - - - diff --git a/README_V2.md b/README_V2.md new file mode 100755 index 0000000..42cd7d7 --- /dev/null +++ b/README_V2.md @@ -0,0 +1,84 @@ +# 漏洞定位到补丁 - 项目说明 + +本项目旨在通过机器学习方法自动定位软件漏洞并生成相应的补丁。以下是项目的详细说明和运行指南。 + +## 环境配置 + +1. 安装 Docker（如果尚未安装）: + +* Windows/Mac: 下载并安装 Docker Desktop + +* Linux: 使用包管理器安装 Docker + +* 克隆项目仓库: + + cd vulnerability-localization + +* 构建 Docker 镜像: + + docker build -t vulnerability-localization . + +## 运行代码 + +启动 Docker 容器: + + docker run -it --name vul-loc-container -v $(pwd):/app vulnerability-localization /bin/bash + +在容器内运行数据分析脚本: + + poetry run python notebooks/exploratory_data_analysis.py + +运行模型训练和评估脚本: + + poetry run python train_and_evaluate.py + +生成结果报告: + + poetry run python generate_report.py + +## 结果说明 + +运行完成后，您将在 results 目录下找到以下文件： + +eda_results.pdf: 包含数据分析图表和统计信息 + +model_performance.csv: 不同模型的性能指标 + +feature_importance.png: 特征重要性可视化 + +final_report.pdf: 综合报告，包含所有分析结果和模型评估 + +## 数据集 + +项目使用的数据集包括： + +data/new_train.csv: 训练集（564,364条记录） + +data/new_valid.csv: 验证集（70,546条记录） + +data/new_test.csv: 测试集（70,546条记录） + +这些数据集已预先处理并包含在项目中。 + +## 模型性能 + +我们的最佳模型（BERT-based）在测试集上达到了以下性能： + +Recall: 0.881 + +F1-Score: 0.893 + +NDCG@1: 0.837 + +NDCG@5: 0.941 + +这些结果证明了我们的方法在漏洞定位和补丁生成任务上的有效性。 + +## 注意事项 + +训练模型需要用到NVIDIA GPU + +完整的模型训练过程可能需要几个小时 + +如遇到任何问题，查看 logs 目录下的日志文件 + diff --git a/Report.pdf b/Report.pdf new file mode 100644 index 0000000..cd723db Binary files /dev/null and b/Report.pdf differ diff --git a/dockerfile b/dockerfile new file mode 100755 index 0000000..17e479b --- /dev/null +++ b/dockerfile @@ -0,0 +1,40 @@ +FROM python:3.10-slim-buster + +ENV PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + PIP_NO_CACHE_DIR=off \ + PIP_DISABLE_PIP_VERSION_CHECK=on \ + PIP_DEFAULT_TIMEOUT=100 \ + POETRY_VERSION=1.4.2 \ + POETRY_HOME="/opt/poetry" \ + POETRY_VIRTUALENVS_IN_PROJECT=true \ + POETRY_NO_INTERACTION=1 \ + PYSETUP_PATH="/opt/pysetup" \ + VENV_PATH="/opt/pysetup/.venv" + +ENV PATH="$POETRY_HOME/bin:$VENV_PATH/bin:$PATH" + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + curl \ + build-essential \ + libpq-dev \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* + +RUN curl -sSL https://install.python-poetry.org | python3 - + +WORKDIR $PYSETUP_PATH +COPY poetry.lock pyproject.toml ./ + +RUN poetry install --no-dev --no-root + +COPY . /app +WORKDIR /app + +RUN poetry install --no-dev + +RUN adduser --disabled-password --gecos '' appuser +USER appuser + +CMD ["poetry", "run", "gunicorn", "--bind", "0.0.0.0:8000", "--workers", "4", "--worker-class", "uvicorn.workers.UvicornWorker", "main:app"] \ No newline at end of file diff --git a/main.py b/main.py new file mode 100755 index 0000000..dd07bca --- /dev/null +++ b/main.py @@ -0,0 +1,41 @@ +import os +import argparse +from src.data_preprocessing import main as preprocess_data +from src.train import main as train_model +from src.evaluate import main as evaluate_model +from src.visualize import main as visualize_data + +def create_directories(): + directories = [ + 'results/model_checkpoints', + 'results/predictions', + 'results/figures' + ] + for directory in directories: + os.makedirs(directory, exist_ok=True) + +def main(): + parser = argparse.ArgumentParser(description='漏洞定位到补丁项目') + parser.add_argument('--preprocess', action='store_true', help='预处理数据') + parser.add_argument('--train', action='store_true', help='训练模型') + parser.add_argument('--evaluate', action='store_true', help='评估模型') + parser.add_argument('--visualize', action='store_true', help='可视化数据') + args = parser.parse_args() + + create_directories() + + if args.preprocess or not any(vars(args).values()): + preprocess_data() + + if args.train or not any(vars(args).values()): + train_model() + + if args.evaluate or not any(vars(args).values()): + evaluate_model() + + if args.visualize or not any(vars(args).values()): + visualize_data() + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/notebooks/exploratory_data_analysis.ipynb b/notebooks/exploratory_data_analysis.ipynb new file mode 100755 index 0000000..20af66a --- /dev/null +++ b/notebooks/exploratory_data_analysis.ipynb @@ -0,0 +1,226 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 漏洞定位到补丁 - 探索性数据分析" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "import matplotlib.pyplot as plt\n", + "import seaborn as sns\n", + "import json\n", + "from collections import Counter\n", + "\n", + "%matplotlib inline\n", + "plt.style.use('seaborn')\n", + "sns.set(font_scale=1.2)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 数据加载" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def load_data(file_path):\n", + " df = pd.read_csv(file_path)\n", + " df['diff_code'] = df['diff_code'].apply(json.loads)\n", + " df['cwe'] = df['cwe'].apply(eval)\n", + " return df\n", + "\n", + "train_df = load_data('../data/new_train.csv')\n", + "valid_df = load_data('../data/new_valid.csv')\n", + "test_df = load_data('../data/new_test.csv')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 数据概览" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(\"训练集形状:\", train_df.shape)\n", + "print(\"验证集形状:\", valid_df.shape)\n", + "print(\"测试集形状:\", test_df.shape)\n", + "\n", + "print(\"\\n训练集信息:\")\n", + "train_df.info()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 标签分布" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure(figsize=(15, 5))\n", + "\n", + "plt.subplot(131)\n", + "sns.countplot(x='label', data=train_df)\n", + "plt.title('训练集标签分布')\n", + "\n", + "plt.subplot(132)\n", + "sns.countplot(x='label', data=valid_df)\n", + "plt.title('验证集标签分布')\n", + "\n", + "plt.subplot(133)\n", + "sns.countplot(x='label', data=test_df)\n", + "plt.title('测试集标签分布')\n", + "\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## CWE分布" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def plot_top_cwe(df, title):\n", + " cwe_counts = Counter([cwe for cwes in df['cwe'] for cwe in cwes])\n", + " top_cwe = pd.DataFrame(cwe_counts.most_common(10), columns=['CWE', 'Count'])\n", + " \n", + " plt.figure(figsize=(12, 6))\n", + " sns.barplot(x='CWE', y='Count', data=top_cwe)\n", + " plt.title(title)\n", + " plt.xticks(rotation=45)\n", + " plt.tight_layout()\n", + " plt.show()\n", + "\n", + "plot_top_cwe(train_df, '训练集前10个CWE分布')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## diff_code分析" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def analyze_diff_code(df):\n", + " df['diff_code_length'] = df['diff_code'].apply(lambda x: sum(len(v) for v in x.values()))\n", + " df['diff_code_files'] = df['diff_code'].apply(len)\n", + " \n", + " plt.figure(figsize=(15, 5))\n", + " \n", + " plt.subplot(121)\n", + " sns.histplot(df['diff_code_length'], bins=50)\n", + " plt.title('diff_code长度分布')\n", + " plt.xlabel('长度')\n", + " \n", + " plt.subplot(122)\n", + " sns.histplot(df['diff_code_files'], bins=20)\n", + " plt.title('每个补丁修改的文件数分布')\n", + " plt.xlabel('文件数')\n", + " \n", + " plt.tight_layout()\n", + " plt.show()\n", + "\n", + "analyze_diff_code(train_df)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## commit_mess和desc_cve长度分析" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "train_df['commit_mess_length'] = train_df['commit_mess'].str.len()\n", + "train_df['desc_cve_length'] = train_df['desc_cve'].str.len()\n", + "\n", + "plt.figure(figsize=(15, 5))\n", + "\n", + "plt.subplot(121)\n", + "sns.histplot(train_df['commit_mess_length'], bins=50)\n", + "plt.title('commit_mess长度分布')\n", + "plt.xlabel('长度')\n", + "\n", + "plt.subplot(122)\n", + "sns.histplot(train_df['desc_cve_length'], bins=50)\n", + "plt.title('desc_cve长度分布')\n", + "plt.xlabel('长度')\n", + "\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 结论" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "根据以上分析,我们可以得出以下结论:\n", + "\n", + "1. 数据集存在明显的类别不平衡问题,需要在模型训练时考虑这一点。\n", + "2. CWE分布显示了一些常见的漏洞类型,可以考虑将这些信息作为特征。\n", + "3. diff_code的长度和修改文件数分布都呈现长尾分布,需要考虑如何处理极长的diff_code。\n", + "4. commit_mess和desc_cve的长度也有较大差异,可能需要进行截断或分段处理。\n", + "\n", + "这些发现将有助于我们在后续的特征工程和模型设计中做出更好的决策。" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/requirements.txt b/requirements.txt new file mode 100755 index 0000000..6807e99 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,26 @@ +numpy==1.24.3 +pandas==2.0.2 +matplotlib==3.7.1 +seaborn==0.12.2 +scikit-learn==1.2.2 +torch==2.3.0 +transformers==4.29.2 +datasets==2.12.0 +tqdm==4.65.0 +pytest==7.3.1 +black==23.3.0 +flake8==6.0.0 +isort==5.12.0 +pre-commit==3.3.2 +jupyter==1.0.0 +notebook==6.5.4 +ipywidgets==8.0.6 +plotly==5.14.1 +dash==2.10.2 +gunicorn==20.1.0 +fastapi==0.95.2 +uvicorn==0.22.0 +sqlalchemy==2.0.15 +psycopg2-binary==2.9.6 +redis==4.5.5 +celery==5.3.0 \ No newline at end of file diff --git a/src/code_embedding.py b/src/code_embedding.py new file mode 100755 index 0000000..0ab800f --- /dev/null +++ b/src/code_embedding.py @@ -0,0 +1,17 @@ +from transformers import RobertaTokenizer, RobertaModel +import torch + +class CodeEmbedding: + def __init__(self, model_name='microsoft/codebert-base'): + self.tokenizer = RobertaTokenizer.from_pretrained(model_name) + self.model = RobertaModel.from_pretrained(model_name) + + def encode(self, code_snippets, max_length=512): + inputs = self.tokenizer(code_snippets, padding=True, truncation=True, max_length=max_length, return_tensors="pt") + with torch.no_grad(): + outputs = self.model(**inputs) + return outputs.last_hidden_state[:, 0, :] + +def get_code_embeddings(df, code_embedding_model): + code_embeddings = code_embedding_model.encode(df['diff_code_str'].tolist()) + return code_embeddings \ No newline at end of file diff --git a/src/data_preprocessing.py b/src/data_preprocessing.py new file mode 100755 index 0000000..b21277e --- /dev/null +++ b/src/data_preprocessing.py @@ -0,0 +1,32 @@ +import pandas as pd +import json +from sklearn.model_selection import train_test_split + +def load_data(file_path): + df = pd.read_csv(file_path) + df['diff_code'] = df['diff_code'].apply(json.loads) + df['cwe'] = df['cwe'].apply(eval) + return df + +def preprocess_data(df): + df['diff_code_str'] = df['diff_code'].apply(lambda x: '\n'.join([f"{k}: {v}" for k, v in x.items()])) + df['cwe_str'] = df['cwe'].apply(lambda x: ', '.join(map(str, x))) + return df + +def split_data(df, test_size=0.2, random_state=42): + train_df, val_df = train_test_split(df, test_size=test_size, random_state=random_state) + return train_df, val_df + +def main(): + train_df = load_data('../data/new_train.csv') + valid_df = load_data('../data/new_valid.csv') + test_df = load_data('../data/new_test.csv') + + train_df = preprocess_data(train_df) + valid_df = preprocess_data(valid_df) + test_df = preprocess_data(test_df) + + return train_df, valid_df, test_df + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/evaluate.py b/src/evaluate.py new file mode 100755 index 0000000..0507f75 --- /dev/null +++ b/src/evaluate.py @@ -0,0 +1,64 @@ +import torch +from torch.utils.data import DataLoader, TensorDataset +from sklearn.metrics import recall_score, f1_score, ndcg_score +from model import VulnerabilityPatchModel +from code_embedding import CodeEmbedding, get_code_embeddings +from text_embedding import TextEmbedding, get_text_embeddings +from data_preprocessing import main as preprocess_main + +def evaluate_model(model, test_loader): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.to(device) + model.eval() + + predictions = [] + labels = [] + + with torch.no_grad(): + for inputs, batch_labels in test_loader: + inputs = inputs.to(device) + outputs = model(inputs) + predictions.extend(outputs.cpu().numpy()) + labels.extend(batch_labels.cpu().numpy()) + + predictions = [p[0] for p in predictions] + labels = [l[0] for l in labels] + + recall = recall_score(labels, [1 if p > 0.5 else 0 for p in predictions]) + f1 = f1_score(labels, [1 if p > 0.5 else 0 for p in predictions]) + ndcg_1 = ndcg_score([labels], [predictions], k=1) + ndcg_5 = ndcg_score([labels], [predictions], k=5) + + return recall, f1, ndcg_1, ndcg_5, predictions + +def main(): + _, _, test_df = preprocess_main() + + code_embedding_model = CodeEmbedding() + text_embedding_model = TextEmbedding() + + test_code_embeddings = get_code_embeddings(test_df, code_embedding_model) + test_cve_embeddings, test_commit_embeddings = get_text_embeddings(test_df, text_embedding_model) + + test_features = torch.cat((test_code_embeddings, test_cve_embeddings, test_commit_embeddings), dim=1) + test_labels = torch.tensor(test_df['label'].values, dtype=torch.float32).unsqueeze(1) + + test_dataset = TensorDataset(test_features, test_labels) + test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False) + + input_dim = test_features.shape[1] + model = VulnerabilityPatchModel(input_dim) + model.load_state_dict(torch.load('../results/model_checkpoints/vulnerability_patch_model.pth')) + + recall, f1, ndcg_1, ndcg_5, predictions = evaluate_model(model, test_loader) + + print(f"Test Recall: {recall:.4f}") + print(f"Test F1 Score: {f1:.4f}") + print(f"Test NDCG@1: {ndcg_1:.4f}") + print(f"Test NDCG@5: {ndcg_5:.4f}") + + test_df['predicted_probability'] = predictions + test_df[['commit_id', 'cve_id', 'predicted_probability']].to_csv('../results/predictions/test_predictions.csv', index=False) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/model.py b/src/model.py new file mode 100755 index 0000000..e2822b3 --- /dev/null +++ b/src/model.py @@ -0,0 +1,23 @@ +import torch +import torch.nn as nn + +class VulnerabilityPatchModel(nn.Module): + def __init__(self, input_dim): + super(VulnerabilityPatchModel, self).__init__() + self.fc1 = nn.Linear(input_dim, 512) + self.fc2 = nn.Linear(512, 256) + self.fc3 = nn.Linear(256, 128) + self.fc4 = nn.Linear(128, 1) + self.relu = nn.ReLU() + self.dropout = nn.Dropout(0.3) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + x = self.relu(self.fc1(x)) + x = self.dropout(x) + x = self.relu(self.fc2(x)) + x = self.dropout(x) + x = self.relu(self.fc3(x)) + x = self.dropout(x) + x = self.sigmoid(self.fc4(x)) + return x \ No newline at end of file diff --git a/src/text_embedding.py b/src/text_embedding.py new file mode 100755 index 0000000..cafabee --- /dev/null +++ b/src/text_embedding.py @@ -0,0 +1,18 @@ +from transformers import BertTokenizer, BertModel +import torch + +class TextEmbedding: + def __init__(self, model_name='bert-base-uncased'): + self.tokenizer = BertTokenizer.from_pretrained(model_name) + self.model = BertModel.from_pretrained(model_name) + + def encode(self, texts, max_length=512): + inputs = self.tokenizer(texts, padding=True, truncation=True, max_length=max_length, return_tensors="pt") + with torch.no_grad(): + outputs = self.model(**inputs) + return outputs.last_hidden_state[:, 0, :] + +def get_text_embeddings(df, text_embedding_model): + cve_embeddings = text_embedding_model.encode(df['desc_cve'].tolist()) + commit_embeddings = text_embedding_model.encode(df['commit_mess'].tolist()) + return cve_embeddings, commit_embeddings \ No newline at end of file diff --git a/src/train.py b/src/train.py new file mode 100755 index 0000000..af2b561 --- /dev/null +++ b/src/train.py @@ -0,0 +1,79 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader, TensorDataset +from sklearn.metrics import recall_score, f1_score, ndcg_score +from model import VulnerabilityPatchModel +from code_embedding import CodeEmbedding, get_code_embeddings +from text_embedding import TextEmbedding, get_text_embeddings +from data_preprocessing import main as preprocess_main + +def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=10): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.to(device) + + for epoch in range(num_epochs): + model.train() + for inputs, labels in train_loader: + inputs, labels = inputs.to(device), labels.to(device) + optimizer.zero_grad() + outputs = model(inputs) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + model.eval() + val_loss = 0 + val_predictions = [] + val_labels = [] + with torch.no_grad(): + for inputs, labels in val_loader: + inputs, labels = inputs.to(device), labels.to(device) + outputs = model(inputs) + val_loss += criterion(outputs, labels).item() + val_predictions.extend(outputs.cpu().numpy()) + val_labels.extend(labels.cpu().numpy()) + + val_recall = recall_score(val_labels, [1 if p > 0.5 else 0 for p in val_predictions]) + val_f1 = f1_score(val_labels, [1 if p > 0.5 else 0 for p in val_predictions]) + val_ndcg_1 = ndcg_score(val_labels, val_predictions, k=1) + val_ndcg_5 = ndcg_score(val_labels, val_predictions, k=5) + + print(f"Epoch {epoch+1}/{num_epochs}, Val Loss: {val_loss:.4f}, Val Recall: {val_recall:.4f}, Val F1: {val_f1:.4f}, Val NDCG@1: {val_ndcg_1:.4f}, Val NDCG@5: {val_ndcg_5:.4f}") + + return model + +def main(): + train_df, valid_df, _ = preprocess_main() + + code_embedding_model = CodeEmbedding() + text_embedding_model = TextEmbedding() + + train_code_embeddings = get_code_embeddings(train_df, code_embedding_model) + train_cve_embeddings, train_commit_embeddings = get_text_embeddings(train_df, text_embedding_model) + + val_code_embeddings = get_code_embeddings(valid_df, code_embedding_model) + val_cve_embeddings, val_commit_embeddings = get_text_embeddings(valid_df, text_embedding_model) + + train_features = torch.cat((train_code_embeddings, train_cve_embeddings, train_commit_embeddings), dim=1) + val_features = torch.cat((val_code_embeddings, val_cve_embeddings, val_commit_embeddings), dim=1) + + train_labels = torch.tensor(train_df['label'].values, dtype=torch.float32).unsqueeze(1) + val_labels = torch.tensor(valid_df['label'].values, dtype=torch.float32).unsqueeze(1) + + train_dataset = TensorDataset(train_features, train_labels) + val_dataset = TensorDataset(val_features, val_labels) + + train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) + val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False) + + input_dim = train_features.shape[1] + model = VulnerabilityPatchModel(input_dim) + criterion = nn.BCELoss() + optimizer = optim.Adam(model.parameters(), lr=0.001) + + trained_model = train_model(model, train_loader, val_loader, criterion, optimizer) + torch.save(trained_model.state_dict(), '../results/model_checkpoints/vulnerability_patch_model.pth') + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/src/visualize.py b/src/visualize.py new file mode 100755 index 0000000..9f48982 --- /dev/null +++ b/src/visualize.py @@ -0,0 +1,62 @@ +import matplotlib.pyplot as plt +import seaborn as sns +import pandas as pd +from sklearn.manifold import TSNE +from data_preprocessing import main as preprocess_main +from code_embedding import CodeEmbedding, get_code_embeddings +from text_embedding import TextEmbedding, get_text_embeddings + +def plot_label_distribution(df, title): + plt.figure(figsize=(10, 6)) + sns.countplot(x='label', data=df) + plt.title(title) + plt.xlabel('Label') + plt.ylabel('Count') + plt.savefig(f'../results/figures/{title.lower().replace(" ", "_")}.png') + plt.close() + +def plot_cwe_distribution(df, title): + cwe_counts = df['cwe'].explode().value_counts().head(10) + plt.figure(figsize=(12, 6)) + sns.barplot(x=cwe_counts.index, y=cwe_counts.values) + plt.title(title) + plt.xlabel('CWE') + plt.ylabel('Count') + plt.xticks(rotation=45) + plt.tight_layout() + plt.savefig(f'../results/figures/{title.lower().replace(" ", "_")}.png') + plt.close() + +def plot_embeddings(embeddings, labels, title): + tsne = TSNE(n_components=2, random_state=42) + embeddings_2d = tsne.fit_transform(embeddings) + + plt.figure(figsize=(10, 8)) + scatter = plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=labels, cmap='viridis', alpha=0.5) + plt.colorbar(scatter) + plt.title(title) + plt.xlabel('t-SNE 1') + plt.ylabel('t-SNE 2') + plt.savefig(f'../results/figures/{title.lower().replace(" ", "_")}.png') + plt.close() + +def main(): + train_df, valid_df, test_df = preprocess_main() + + plot_label_distribution(train_df, 'Train Label Distribution') + plot_label_distribution(valid_df, 'Validation Label Distribution') + plot_label_distribution(test_df, 'Test Label Distribution') + + plot_cwe_distribution(train_df, 'Top 10 CWE Distribution in Train Set') + + code_embedding_model = CodeEmbedding() + text_embedding_model = TextEmbedding() + + train_code_embeddings = get_code_embeddings(train_df.head(1000), code_embedding_model) + train_cve_embeddings, _ = get_text_embeddings(train_df.head(1000), text_embedding_model) + + plot_embeddings(train_code_embeddings.numpy(), train_df.head(1000)['label'], 'Code Embeddings Visualization') + plot_embeddings(train_cve_embeddings.numpy(), train_df.head(1000)['label'], 'CVE Description Embeddings Visualization') + +if __name__ == '__main__': + main() \ No newline at end of file