29 changed files with 1495 additions and 14 deletions
--- 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镜像。具体要求参见赛事网站的“参赛指南”。
-
-注：推荐使用开源大模型。
-
-
-
--- a/README_V2.md
+++ b/README_V2.md
@ -0,0 +1,124 @@
+# Vulnerability Detection and Classification System
+
+## 项目概述
+
+本项目旨在开发一个用于检测和分类代码漏洞的系统。它集成了多种机器学习和深度学习技术，如 XGBoost 和基于 CodeBERT 的模型，以实现对代码漏洞的高效检测与分类。
+
+## 文件结构
+
+```
+vulnerability_detection/
+├── data/                # 数据目录
+│   ├── raw/             # 原始数据
+│   └── processed/       # 处理后的数据
+├── outputs/             # 输出目录
+│   ├── models/          # 训练好的模型
+│   ├── evaluation/      # 评估结果
+│   └── figures/         # 可视化图表
+├── scripts/             # 脚本目录
+│   ├── run_preprocessing.py
+│   ├── run_training.py
+│   └── run_evaluation.py
+├── src/                 # 源代码目录
+│   ├── data/            # 数据处理模块
+│   ├── evaluation/      # 评估模块
+│   ├── models/          # 模型定义
+│   ├── training/        # 训练模块
+│   ├── utils/           # 工具模块
+│   └── visualization/   # 可视化模块
+├── docker-compose.yml    # Docker Compose 配置文件
+├── Dockerfile            # Dockerfile
+├── Dockerfile.dashboard  # Dashboard Dockerfile
+├── Dockerfile.celery     # Celery Dockerfile
+├── nginx.conf            # Nginx 配置文件
+├── setup.py              # 安装脚本
+└── README.md             # 项目说明文件
+```
+
+## 环境配置
+
+1. **安装 Docker 和 Docker Compose**:
+
+   请确保您的系统已安装 Docker 和 Docker Compose。可以使用以下命令检查安装情况：
+
+   ```
+   docker --version
+   docker-compose --version
+   ```
+
+2. **构建 Docker 镜像**:
+
+   在项目根目录下运行以下命令构建 Docker 镜像：
+
+   ```
+   docker-compose build
+   ```
+
+3. **安装 Python 依赖**:
+
+   如果您希望在本地环境中运行代码，请确保安装 Python 3.7 或更高版本，并使用以下命令安装依赖：
+
+   ```
+   pip install -r requirements.txt
+   ```
+
+## 运行代码
+
+### 数据预处理
+
+1. **在 Docker 容器中运行数据预处理脚本**：
+
+   ```
+   docker-compose run data_preprocessor
+   ```
+
+2. **或者在本地环境中运行**：
+
+   ```
+   python scripts/run_preprocessing.py --data_dir=data/raw --output_dir=data/processed
+   ```
+
+### 模型训练
+
+1. **在 Docker 容器中运行训练脚本**：
+
+   ```
+   docker-compose run model_trainer
+   ```
+
+2. **或者在本地环境中运行**：
+
+   ```
+   python scripts/run_training.py --model=xgboost --data_dir=data/processed --output_dir=outputs/models
+   ```
+
+### 模型评估
+
+1. **在 Docker 容器中运行评估脚本**：
+
+   ```
+   docker-compose run evaluator
+   ```
+
+2. **或者在本地环境中运行**：
+
+   ```
+   python scripts/run_evaluation.py --model=xgboost --data_dir=data/processed --output_dir=outputs/evaluation
+   ```
+
+### 启动 Dashboard
+
+1. **在 Docker 容器中启动 Dashboard**：
+
+   ```
+   docker-compose up dashboard
+   ```
+
+2. **访问 Dashboard**：
+
+   打开浏览器并访问 `http://localhost:8050`，查看可视化结果。
+
+## 结果报告
+
+评估结果将保存在 `outputs/evaluation` 目录下，您可以查看生成的 JSON 文件以获取详细的评估指标。此外，所有生成的图表将保存在 `outputs/figures` 目录中。
+
--- a/Report.pdf
+++ b/Report.pdf
--- a/configs/linevul_config.yaml
+++ b/configs/linevul_config.yaml
@ -0,0 +1,43 @@
+model:
+  name: 'linevul'
+  hidden_size: 768
+  num_layers: 2
+  dropout: 0.1
+
+training:
+  batch_size: 32
+  num_epochs: 10
+  learning_rate: 2e-5
+  weight_decay: 0.01
+  warmup_steps: 500
+  max_grad_norm: 1.0
+
+data:
+  max_seq_length: 512
+  train_file: 'data/processed/train/processed_data.json'
+  val_file: 'data/processed/valid/processed_data.json'
+  test_file: 'data/processed/test/processed_data.json'
+
+tokenizer:
+  name: 'microsoft/codebert-base'
+  do_lower_case: False
+
+optimization:
+  gradient_accumulation_steps: 1
+  fp16: False
+  fp16_opt_level: 'O1'
+
+logging:
+  log_every_n_steps: 100
+  save_every_n_epochs: 1
+
+evaluation:
+  eval_batch_size: 64
+  eval_during_training: True
+  eval_steps: 1000
+
+output:
+  output_dir: 'outputs/linevul'
+  overwrite_output_dir: True
+
+seed: 42
--- a/configs/xgboost_config.yaml
+++ b/configs/xgboost_config.yaml
@ -0,0 +1,70 @@
+model:
+  name: 'xgboost'
+  classifier:
+    n_estimators: 100
+    max_depth: 6
+    learning_rate: 0.1
+    subsample: 0.8
+    colsample_bytree: 0.8
+    objective: 'multi:softprob'
+    num_class: 10  # Update this based on your actual number of classes
+  regressor:
+    n_estimators: 100
+    max_depth: 6
+    learning_rate: 0.1
+    subsample: 0.8
+    colsample_bytree: 0.8
+    objective: 'reg:squarederror'
+
+training:
+  early_stopping_rounds: 10
+  eval_metric:
+    classifier: ['mlogloss', 'merror']
+    regressor: ['rmse', 'mae']
+
+data:
+  train_file: 'data/processed/train/features.npy'
+  val_file: 'data/processed/valid/features.npy'
+  test_file: 'data/processed/test/features.npy'
+  label_file:
+    train: 'data/processed/train/labels.npy'
+    val: 'data/processed/valid/labels.npy'
+    test: 'data/processed/test/labels.npy'
+  cvss_file:
+    train: 'data/processed/train/cvss_scores.npy'
+    val: 'data/processed/valid/cvss_scores.npy'
+    test: 'data/processed/test/cvss_scores.npy'
+
+feature_engineering:
+  use_feature_selection: True
+  feature_selection_method: 'mutual_info'
+  n_features_to_select: 100
+
+hyperparameter_tuning:
+  perform_tuning: True
+  cv: 5
+  n_iter: 100
+  param_distributions:
+    classifier:
+      n_estimators: [50, 100, 200, 300]
+      max_depth: [3, 4, 5, 6, 7, 8]
+      learning_rate: [0.01, 0.05, 0.1, 0.2]
+      subsample: [0.6, 0.7, 0.8, 0.9]
+      colsample_bytree: [0.6, 0.7, 0.8, 0.9]
+    regressor:
+      n_estimators: [50, 100, 200, 300]
+      max_depth: [3, 4, 5, 6, 7, 8]
+      learning_rate: [0.01, 0.05, 0.1, 0.2]
+      subsample: [0.6, 0.7, 0.8, 0.9]
+      colsample_bytree: [0.6, 0.7, 0.8, 0.9]
+
+output:
+  output_dir: 'outputs/xgboost'
+  save_model: True
+  save_feature_importance: True
+
+logging:
+  log_level: 'INFO'
+  log_file: 'logs/xgboost_training.log'
+
+seed: 42
--- a/data/raw/test.jsonl
+++ b/data/raw/test.jsonl
--- a/data/raw/train.jsonl
+++ b/data/raw/train.jsonl
--- a/data/raw/valid.jsonl
+++ b/data/raw/valid.jsonl
--- a/docker-compose.yml
+++ b/docker-compose.yml
@ -0,0 +1,141 @@
+version: '3.8'
+
+services:
+  vulnerability_detector:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    image: vulnerability_detector:latest
+    container_name: vul_detector
+    volumes:
+      - ./data:/app/data
+      - ./outputs:/app/outputs
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=all
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    depends_on:
+      - redis
+      - postgres
+    networks:
+      - vul_net
+
+  data_preprocessor:
+    image: vulnerability_detector:latest
+    command: python scripts/run_preprocessing.py
+    volumes:
+      - ./data:/app/data
+    depends_on:
+      - vulnerability_detector
+    networks:
+      - vul_net
+
+  model_trainer:
+    image: vulnerability_detector:latest
+    command: python scripts/run_training.py
+    volumes:
+      - ./data:/app/data
+      - ./outputs:/app/outputs
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=all
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    depends_on:
+      - data_preprocessor
+    networks:
+      - vul_net
+
+  evaluator:
+    image: vulnerability_detector:latest
+    command: python scripts/run_evaluation.py
+    volumes:
+      - ./data:/app/data
+      - ./outputs:/app/outputs
+    depends_on:
+      - model_trainer
+    networks:
+      - vul_net
+
+  dashboard:
+    build:
+      context: .
+      dockerfile: Dockerfile.dashboard
+    ports:
+      - "8050:8050"
+    volumes:
+      - ./outputs:/app/outputs
+    depends_on:
+      - evaluator
+    networks:
+      - vul_net
+
+  redis:
+    image: redis:alpine
+    ports:
+      - "6379:6379"
+    volumes:
+      - redis_data:/data
+    networks:
+      - vul_net
+
+  postgres:
+    image: postgres:13
+    environment:
+      POSTGRES_DB: vuldb
+      POSTGRES_USER: vuluser
+      POSTGRES_PASSWORD: vulpass
+    volumes:
+      - postgres_data:/var/lib/postgresql/data
+    networks:
+      - vul_net
+
+  rabbitmq:
+    image: rabbitmq:3-management
+    ports:
+      - "5672:5672"
+      - "15672:15672"
+    networks:
+      - vul_net
+
+  celery_worker:
+    build:
+      context: .
+      dockerfile: Dockerfile.celery
+    command: celery -A tasks worker --loglevel=info
+    volumes:
+      - ./data:/app/data
+      - ./outputs:/app/outputs
+    depends_on:
+      - rabbitmq
+      - redis
+    networks:
+      - vul_net
+
+  nginx:
+    image: nginx:alpine
+    ports:
+      - "80:80"
+    volumes:
+      - ./nginx.conf:/etc/nginx/nginx.conf:ro
+    depends_on:
+      - dashboard
+    networks:
+      - vul_net
+
+volumes:
+  redis_data:
+  postgres_data:
+
+networks:
+  vul_net:
+    driver: bridge
--- a/docs/model_architecture.md
+++ b/docs/model_architecture.md
--- a/docs/project_structure.md
+++ b/docs/project_structure.md
--- a/requirements.txt
+++ b/requirements.txt
@ -0,0 +1,11 @@
+numpy
+pandas
+scikit-learn
+torch
+transformers
+xgboost
+matplotlib
+seaborn
+dash
+plotly
+pyyaml
--- a/scripts/run_evaluation.py
+++ b/scripts/run_evaluation.py
@ -0,0 +1,96 @@
+import sys
+import os
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from src.models.linevul import LineVulModel
+from src.models.xgboost_classifier import XGBoostVulClassifier, XGBoostVulRegressor
+from src.data.data_loader import DataLoader
+from src.evaluation.evaluator import Evaluator
+from src.visualization.plot_utils import VisualizationUtils
+import argparse
+import torch
+import numpy as np
+import json
+
+def load_linevul_model(model_path, num_labels):
+    model = LineVulModel(num_labels=num_labels)
+    model.load_state_dict(torch.load(model_path))
+    return model
+
+def load_xgboost_models(classifier_path, regressor_path):
+    classifier = XGBoostVulClassifier()
+    classifier.model.load_model(classifier_path)
+    regressor = XGBoostVulRegressor()
+    regressor.model.load_model(regressor_path)
+    return classifier, regressor
+
+def main(args):
+    data_loader = DataLoader(args.data_dir)
+    evaluator = Evaluator()
+    viz_utils = VisualizationUtils()
+
+    test_data = {
+        'features': np.load(os.path.join(args.data_dir, 'test', 'features.npy')),
+        'labels': np.load(os.path.join(args.data_dir, 'test', 'labels.npy')),
+        'cvss_scores': np.load(os.path.join(args.data_dir, 'test', 'cvss_scores.npy'))
+    }
+
+    if args.model == 'linevul':
+        model = load_linevul_model(args.model_path, num_labels=len(set(test_data['labels'])))
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        model.to(device)
+        model.eval()
+
+        test_dataset = torch.utils.data.TensorDataset(
+            torch.tensor(test_data['features']).float(),
+            torch.tensor(test_data['labels']).long()
+        )
+        test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size)
+
+        all_preds = []
+        all_labels = []
+        with torch.no_grad():
+            for batch in test_loader:
+                inputs, labels = batch
+                inputs = inputs.to(device)
+                outputs = model(inputs)
+                _, preds = torch.max(outputs, 1)
+                all_preds.extend(preds.cpu().numpy())
+                all_labels.extend(labels.numpy())
+
+        results = evaluator.evaluate_classification(all_labels, all_preds)
+    elif args.model == 'xgboost':
+        classifier, regressor = load_xgboost_models(args.classifier_path, args.regressor_path)
+        class_preds = classifier.predict(test_data['features'])
+        reg_preds = regressor.predict(test_data['features'])
+
+        results = evaluator.evaluate_classification(test_data['labels'], class_preds)
+        results.update(evaluator.evaluate_regression(test_data['cvss_scores'], reg_preds))
+    else:
+        raise ValueError(f"Unsupported model: {args.model}")
+
+    evaluator.evaluate_and_print(results)
+
+    viz_utils.plot_confusion_matrix(test_data['labels'], all_preds if args.model == 'linevul' else class_preds,
+                                    labels=list(set(test_data['labels'])), title=f'Confusion Matrix - {args.model}')
+
+    if args.model == 'xgboost':
+        feature_importance = classifier.feature_importance()
+        viz_utils.plot_feature_importance(feature_importance, [f'Feature_{i}' for i in range(len(feature_importance))],
+                                          title=f'Feature Importance - {args.model}')
+
+    with open(os.path.join(args.output_dir, f'{args.model}_evaluation_results.json'), 'w') as f:
+        json.dump(results, f, indent=2)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run evaluation for vulnerability detection models")
+    parser.add_argument('--model', type=str, choices=['linevul', 'xgboost'], required=True, help='Model to evaluate')
+    parser.add_argument('--data_dir', type=str, default='data/processed', help='Directory containing processed data')
+    parser.add_argument('--model_path', type=str, help='Path to the trained LineVul model')
+    parser.add_argument('--classifier_path', type=str, help='Path to the trained XGBoost classifier')
+    parser.add_argument('--regressor_path', type=str, help='Path to the trained XGBoost regressor')
+    parser.add_argument('--output_dir', type=str, default='outputs/evaluation', help='Directory to save evaluation results')
+    parser.add_argument('--batch_size', type=int, default=32, help='Batch size for evaluation')
+    args = parser.parse_args()
+
+    main(args)
--- a/scripts/run_preprocessing.py
+++ b/scripts/run_preprocessing.py
@ -0,0 +1,40 @@
+import sys
+import os
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from src.data.data_loader import DataLoader
+from src.data.preprocessor import Preprocessor
+from src.data.feature_extractor import FeatureExtractor
+import argparse
+import json
+
+def main(args):
+    data_loader = DataLoader(args.data_dir)
+    preprocessor = Preprocessor()
+    feature_extractor = FeatureExtractor()
+
+    for dataset in ['train', 'valid', 'test']:
+        print(f"Processing {dataset} dataset...")
+        data = data_loader.load_jsonl(f'{dataset}.jsonl')
+        processed_data = preprocessor.process_data(data)
+        features = feature_extractor.extract_features(processed_data)
+
+        output_dir = os.path.join(args.output_dir, dataset)
+        os.makedirs(output_dir, exist_ok=True)
+
+        with open(os.path.join(output_dir, 'processed_data.json'), 'w') as f:
+            json.dump(processed_data, f)
+
+        np.save(os.path.join(output_dir, 'features.npy'), features['features'])
+        np.save(os.path.join(output_dir, 'labels.npy'), features['labels'])
+        np.save(os.path.join(output_dir, 'cvss_scores.npy'), features['cvss_scores'])
+
+    print("Preprocessing completed successfully.")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run preprocessing on vulnerability detection datasets")
+    parser.add_argument('--data_dir', type=str, default='data/raw', help='Directory containing raw data files')
+    parser.add_argument('--output_dir', type=str, default='data/processed', help='Directory to save processed data')
+    args = parser.parse_args()
+
+    main(args)
--- a/scripts/run_training.py
+++ b/scripts/run_training.py
@ -0,0 +1,102 @@
+import sys
+import os
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from src.models.linevul import LineVulModel, LineVulTrainer
+from src.models.xgboost_classifier import XGBoostVulClassifier, XGBoostVulRegressor
+from src.data.data_loader import DataLoader
+from src.evaluation.evaluator import Evaluator
+import argparse
+import torch
+import numpy as np
+from sklearn.model_selection import train_test_split
+
+def train_linevul(args, train_data, val_data):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = LineVulModel(num_labels=len(set(train_data['labels'])))
+    trainer = LineVulTrainer(model, device, learning_rate=args.learning_rate)
+    
+    train_dataset = torch.utils.data.TensorDataset(
+        torch.tensor(train_data['features']).float(),
+        torch.tensor(train_data['labels']).long()
+    )
+    val_dataset = torch.utils.data.TensorDataset(
+        torch.tensor(val_data['features']).float(),
+        torch.tensor(val_data['labels']).long()
+    )
+    
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
+    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size)
+    
+    trainer.train(train_loader, args.num_epochs)
+    
+    _, preds, labels = trainer.evaluate(val_loader)
+    return model, preds, labels
+
+def train_xgboost(args, train_data, val_data):
+    classifier = XGBoostVulClassifier(
+        n_estimators=args.n_estimators,
+        max_depth=args.max_depth,
+        learning_rate=args.learning_rate
+    )
+    classifier.train(train_data['features'], train_data['labels'])
+    
+    regressor = XGBoostVulRegressor(
+        n_estimators=args.n_estimators,
+        max_depth=args.max_depth,
+        learning_rate=args.learning_rate
+    )
+    regressor.train(train_data['features'], train_data['cvss_scores'])
+    
+    class_preds = classifier.predict(val_data['features'])
+    reg_preds = regressor.predict(val_data['features'])
+    
+    return classifier, regressor, class_preds, reg_preds
+
+def main(args):
+    data_loader = DataLoader(args.data_dir)
+    
+    train_data = {
+        'features': np.load(os.path.join(args.data_dir, 'train', 'features.npy')),
+        'labels': np.load(os.path.join(args.data_dir, 'train', 'labels.npy')),
+        'cvss_scores': np.load(os.path.join(args.data_dir, 'train', 'cvss_scores.npy'))
+    }
+    
+    val_data = {
+        'features': np.load(os.path.join(args.data_dir, 'valid', 'features.npy')),
+        'labels': np.load(os.path.join(args.data_dir, 'valid', 'labels.npy')),
+        'cvss_scores': np.load(os.path.join(args.data_dir, 'valid', 'cvss_scores.npy'))
+    }
+    
+    if args.model == 'linevul':
+        model, preds, labels = train_linevul(args, train_data, val_data)
+        torch.save(model.state_dict(), os.path.join(args.output_dir, 'linevul_model.pth'))
+    elif args.model == 'xgboost':
+        classifier, regressor, class_preds, reg_preds = train_xgboost(args, train_data, val_data)
+        classifier.model.save_model(os.path.join(args.output_dir, 'xgboost_classifier.json'))
+        regressor.model.save_model(os.path.join(args.output_dir, 'xgboost_regressor.json'))
+    else:
+        raise ValueError(f"Unsupported model: {args.model}")
+    
+    evaluator = Evaluator()
+    if args.model == 'linevul':
+        results = evaluator.evaluate_classification(labels, preds)
+    else:
+        results = evaluator.evaluate_classification(val_data['labels'], class_preds)
+        results.update(evaluator.evaluate_regression(val_data['cvss_scores'], reg_preds))
+    
+    evaluator.evaluate_and_print(results)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run training for vulnerability detection models")
+    parser.add_argument('--model', type=str, choices=['linevul', 'xgboost'], required=True, help='Model to train')
+    parser.add_argument('--data_dir', type=str, default='data/processed', help='Directory containing processed data')
+    parser.add_argument('--output_dir', type=str, default='outputs/models', help='Directory to save trained models')
+    parser.add_argument('--batch_size', type=int, default=32, help='Batch size for training')
+    parser.add_argument('--num_epochs', type=int, default=10, help='Number of epochs for training')
+    parser.add_argument('--learning_rate', type=float, default=2e-5, help='Learning rate')
+    parser.add_argument('--n_estimators', type=int, default=100, help='Number of estimators for XGBoost')
+    parser.add_argument('--max_depth', type=int, default=3, help='Max depth for XGBoost')
+    args = parser.parse_args()
+
+    main(args)
--- a/setup.py
+++ b/setup.py
@ -0,0 +1,41 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='vulnerability_detection',
+    version='0.1.0',
+    packages=find_packages(exclude=['tests*']),
+    description='A vulnerability detection and classification system',
+    long_description=open('README.md').read(),
+    long_description_content_type='text/markdown',
+    url='https://github.com/yourusername/vulnerability_detection',
+    install_requires=[
+        'numpy',
+        'pandas',
+        'scikit-learn',
+        'torch',
+        'transformers',
+        'xgboost',
+        'matplotlib',
+        'seaborn',
+        'dash',
+        'plotly',
+        'pyyaml',
+    ],
+    classifiers=[
+        'Development Status :: 3 - Alpha',
+        'Intended Audience :: Developers',
+        'License :: OSI Approved :: MIT License',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
+    ],
+    python_requires='>=3.7',
+    entry_points={
+        'console_scripts': [
+            'run_preprocessing=scripts.run_preprocessing:main',
+            'run_training=scripts.run_training:main',
+            'run_evaluation=scripts.run_evaluation:main',
+        ],
+    },
+)
--- a/src/data/data_loader.py
+++ b/src/data/data_loader.py
@ -0,0 +1,31 @@
+import json
+import os
+from typing import List, Dict, Any
+
+class DataLoader:
+    def __init__(self, data_dir: str):
+        self.data_dir = data_dir
+
+    def load_jsonl(self, file_name: str) -> List[Dict[str, Any]]:
+        file_path = os.path.join(self.data_dir, file_name)
+        data = []
+        with open(file_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                data.append(json.loads(line.strip()))
+        return data
+
+    def load_all_data(self) -> Dict[str, List[Dict[str, Any]]]:
+        return {
+            'train': self.load_jsonl('train.jsonl'),
+            'valid': self.load_jsonl('valid.jsonl'),
+            'test': self.load_jsonl('test.jsonl')
+        }
+
+    def save_processed_data(self, data: List[Dict[str, Any]], file_name: str):
+        output_dir = os.path.join(self.data_dir, 'processed')
+        os.makedirs(output_dir, exist_ok=True)
+        output_path = os.path.join(output_dir, file_name)
+        with open(output_path, 'w', encoding='utf-8') as f:
+            for item in data:
+                json.dump(item, f)
+                f.write('\n')
--- a/src/data/feature_extractor.py
+++ b/src/data/feature_extractor.py
@ -0,0 +1,34 @@
+from typing import List, Dict, Any
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+from transformers import RobertaTokenizer, RobertaModel
+import torch
+
+class FeatureExtractor:
+    def __init__(self):
+        self.tfidf_vectorizer = TfidfVectorizer(max_features=1000)
+        self.tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
+        self.model = RobertaModel.from_pretrained("microsoft/codebert-base")
+
+    def extract_tfidf_features(self, data: List[Dict[str, Any]]) -> np.ndarray:
+        code_samples = [' '.join(sample['tokens']) for sample in data]
+        return self.tfidf_vectorizer.fit_transform(code_samples).toarray()
+
+    def extract_codebert_features(self, data: List[Dict[str, Any]]) -> np.ndarray:
+        code_samples = [' '.join(sample['tokens']) for sample in data]
+        inputs = self.tokenizer(code_samples, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        return outputs.last_hidden_state[:, 0, :].numpy()
+
+    def extract_features(self, data: List[Dict[str, Any]]) -> Dict[str, np.ndarray]:
+        tfidf_features = self.extract_tfidf_features(data)
+        codebert_features = self.extract_codebert_features(data)
+        
+        combined_features = np.hstack([tfidf_features, codebert_features])
+        
+        return {
+            'features': combined_features,
+            'labels': np.array([sample['cwe_id'] for sample in data]),
+            'cvss_scores': np.array([sample['cvss'] for sample in data])
+        }
--- a/src/data/preprocessor.py
+++ b/src/data/preprocessor.py
@ -0,0 +1,29 @@
+import re
+from typing import List, Dict, Any
+
+class Preprocessor:
+    def __init__(self):
+        self.max_sequence_length = 512
+
+    def clean_code(self, code: str) -> str:
+        code = re.sub(r'\/\/.*?$|\/\*.*?\*\/|\s+', ' ', code, flags=re.MULTILINE)
+        return code.strip()
+
+    def tokenize_code(self, code: str) -> List[str]:
+        tokens = re.findall(r'\w+|[^\w\s]', code)
+        return tokens[:self.max_sequence_length]
+
+    def process_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+        cleaned_code = self.clean_code(sample['function'])
+        tokens = self.tokenize_code(cleaned_code)
+        
+        processed_sample = {
+            'function_id': sample['function_id'],
+            'tokens': tokens,
+            'cwe_id': sample['cwe_id'][0] if sample['cwe_id'] else 'CWE-0',
+            'cvss': float(sample['cvss']) if sample['cvss'] else 0.0
+        }
+        return processed_sample
+
+    def process_data(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        return [self.process_sample(sample) for sample in data]
--- a/src/evaluation/evaluator.py
+++ b/src/evaluation/evaluator.py
@ -0,0 +1,62 @@
+import numpy as np
+from typing import Dict, Any
+from .metrics import MetricsCalculator
+
+class Evaluator:
+    def __init__(self):
+        self.metrics_calculator = MetricsCalculator()
+
+    def evaluate_classification(self, y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, Any]:
+        overall_metrics = self.metrics_calculator.calculate_metrics(y_true, y_pred)
+        class_metrics = self.metrics_calculator.calculate_class_metrics(y_true, y_pred)
+        
+        evaluation_results = {
+            'overall_metrics': overall_metrics,
+            'class_metrics': class_metrics
+        }
+        
+        return evaluation_results
+
+    def evaluate_regression(self, y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
+        return self.metrics_calculator.calculate_cvss_metrics(y_true, y_pred)
+
+    def evaluate_model(self, model: Any, X_test: np.ndarray, y_test: np.ndarray, task: str) -> Dict[str, Any]:
+        if task == 'classification':
+            y_pred = model.predict(X_test)
+            return self.evaluate_classification(y_test, y_pred)
+        elif task == 'regression':
+            y_pred = model.predict(X_test)
+            return self.evaluate_regression(y_test, y_pred)
+        else:
+            raise ValueError(f"Unsupported task: {task}")
+
+    def cross_validate(self, model: Any, X: np.ndarray, y: np.ndarray, cv: int, task: str) -> Dict[str, Any]:
+        from sklearn.model_selection import cross_val_predict
+        
+        y_pred = cross_val_predict(model, X, y, cv=cv)
+        
+        if task == 'classification':
+            return self.evaluate_classification(y, y_pred)
+        elif task == 'regression':
+            return self.evaluate_regression(y, y_pred)
+        else:
+            raise ValueError(f"Unsupported task: {task}")
+
+    def evaluate_and_print(self, evaluation_results: Dict[str, Any]):
+        print("Evaluation Results:")
+        if 'overall_metrics' in evaluation_results:
+            print("\nOverall Metrics:")
+            for metric, value in evaluation_results['overall_metrics'].items():
+                print(f"{metric.capitalize()}: {value:.4f}")
+        
+        if 'class_metrics' in evaluation_results:
+            print("\nClass-wise Metrics:")
+            for cls, metrics in evaluation_results['class_metrics'].items():
+                print(f"\nClass {cls}:")
+                for metric, value in metrics.items():
+                    print(f"  {metric.capitalize()}: {value:.4f}")
+        
+        if 'mse' in evaluation_results:
+            print("\nRegression Metrics:")
+            for metric, value in evaluation_results.items():
+                print(f"{metric.upper()}: {value:.4f}")
--- a/src/evaluation/metrics.py
+++ b/src/evaluation/metrics.py
@ -0,0 +1,49 @@
+import numpy as np
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, matthews_corrcoef
+from typing import Dict, Any
+
+class MetricsCalculator:
+    @staticmethod
+    def calculate_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
+        accuracy = accuracy_score(y_true, y_pred)
+        precision = precision_score(y_true, y_pred, average='weighted')
+        recall = recall_score(y_true, y_pred, average='weighted')
+        f1 = f1_score(y_true, y_pred, average='weighted')
+        mcc = matthews_corrcoef(y_true, y_pred)
+        
+        return {
+            'accuracy': accuracy,
+            'precision': precision,
+            'recall': recall,
+            'f1': f1,
+            'mcc': mcc
+        }
+
+    @staticmethod
+    def calculate_class_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, Dict[str, float]]:
+        classes = np.unique(np.concatenate((y_true, y_pred)))
+        class_metrics = {}
+        
+        for cls in classes:
+            cls_y_true = (y_true == cls)
+            cls_y_pred = (y_pred == cls)
+            
+            class_metrics[str(cls)] = {
+                'precision': precision_score(cls_y_true, cls_y_pred, average='binary'),
+                'recall': recall_score(cls_y_true, cls_y_pred, average='binary'),
+                'f1': f1_score(cls_y_true, cls_y_pred, average='binary')
+            }
+        
+        return class_metrics
+
+    @staticmethod
+    def calculate_cvss_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
+        mse = np.mean((y_true - y_pred) ** 2)
+        mae = np.mean(np.abs(y_true - y_pred))
+        rmse = np.sqrt(mse)
+        
+        return {
+            'mse': mse,
+            'mae': mae,
+            'rmse': rmse
+        }
--- a/src/models/linevul.py
+++ b/src/models/linevul.py
@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+from transformers import RobertaModel, RobertaConfig
+
+class LineVulModel(nn.Module):
+    def __init__(self, num_labels, hidden_size=768, num_layers=2):
+        super(LineVulModel, self).__init__()
+        self.config = RobertaConfig.from_pretrained('microsoft/codebert-base')
+        self.config.num_labels = num_labels
+        self.roberta = RobertaModel.from_pretrained('microsoft/codebert-base', config=self.config)
+        self.dropout = nn.Dropout(0.1)
+        self.lstm = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True, bidirectional=True)
+        self.classifier = nn.Linear(hidden_size * 2, num_labels)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        lstm_output, _ = self.lstm(sequence_output)
+        logits = self.classifier(lstm_output[:, -1, :])
+        return logits
+
+class LineVulTrainer:
+    def __init__(self, model, device, learning_rate=2e-5):
+        self.model = model
+        self.device = device
+        self.model.to(self.device)
+        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=learning_rate)
+        self.criterion = nn.CrossEntropyLoss()
+
+    def train(self, train_dataloader, num_epochs):
+        self.model.train()
+        for epoch in range(num_epochs):
+            total_loss = 0
+            for batch in train_dataloader:
+                input_ids = batch['input_ids'].to(self.device)
+                attention_mask = batch['attention_mask'].to(self.device)
+                labels = batch['labels'].to(self.device)
+
+                self.optimizer.zero_grad()
+                outputs = self.model(input_ids, attention_mask)
+                loss = self.criterion(outputs, labels)
+                loss.backward()
+                self.optimizer.step()
+
+                total_loss += loss.item()
+            
+            print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(train_dataloader):.4f}")
+
+    def evaluate(self, eval_dataloader):
+        self.model.eval()
+        total_eval_loss = 0
+        all_preds = []
+        all_labels = []
+
+        with torch.no_grad():
+            for batch in eval_dataloader:
+                input_ids = batch['input_ids'].to(self.device)
+                attention_mask = batch['attention_mask'].to(self.device)
+                labels = batch['labels'].to(self.device)
+
+                outputs = self.model(input_ids, attention_mask)
+                loss = self.criterion(outputs, labels)
+                total_eval_loss += loss.item()
+
+                preds = torch.argmax(outputs, dim=1)
+                all_preds.extend(preds.cpu().numpy())
+                all_labels.extend(labels.cpu().numpy())
+
+        avg_eval_loss = total_eval_loss / len(eval_dataloader)
+        return avg_eval_loss, all_preds, all_labels
+
+    def predict(self, test_dataloader):
+        self.model.eval()
+        all_preds = []
+
+        with torch.no_grad():
+            for batch in test_dataloader:
+                input_ids = batch['input_ids'].to(self.device)
+                attention_mask = batch['attention_mask'].to(self.device)
+
+                outputs = self.model(input_ids, attention_mask)
+                preds = torch.argmax(outputs, dim=1)
+                all_preds.extend(preds.cpu().numpy())
+
+        return all_preds
--- a/src/models/xgboost_classifier.py
+++ b/src/models/xgboost_classifier.py
@ -0,0 +1,62 @@
+import numpy as np
+from xgboost import XGBClassifier
+from sklearn.model_selection import GridSearchCV
+
+class XGBoostVulClassifier:
+    def __init__(self, n_estimators=100, max_depth=3, learning_rate=0.1):
+        self.model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth, learning_rate=learning_rate)
+
+    def train(self, X_train, y_train):
+        self.model.fit(X_train, y_train)
+
+    def predict(self, X):
+        return self.model.predict(X)
+
+    def predict_proba(self, X):
+        return self.model.predict_proba(X)
+
+    def feature_importance(self):
+        return self.model.feature_importances_
+
+    def hyperparameter_tuning(self, X_train, y_train, param_grid, cv=5):
+        grid_search = GridSearchCV(estimator=self.model, param_grid=param_grid, cv=cv, scoring='f1_weighted', n_jobs=-1)
+        grid_search.fit(X_train, y_train)
+        self.model = grid_search.best_estimator_
+        return grid_search.best_params_, grid_search.best_score_
+
+class XGBoostVulRegressor:
+    def __init__(self, n_estimators=100, max_depth=3, learning_rate=0.1):
+        self.model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth, learning_rate=learning_rate, objective='reg:squarederror')
+
+    def train(self, X_train, y_train):
+        self.model.fit(X_train, y_train)
+
+    def predict(self, X):
+        return self.model.predict(X)
+
+    def feature_importance(self):
+        return self.model.feature_importances_
+
+    def hyperparameter_tuning(self, X_train, y_train, param_grid, cv=5):
+        grid_search = GridSearchCV(estimator=self.model, param_grid=param_grid, cv=cv, scoring='neg_mean_squared_error', n_jobs=-1)
+        grid_search.fit(X_train, y_train)
+        self.model = grid_search.best_estimator_
+        return grid_search.best_params_, -grid_search.best_score_
+
+class EnsembleVulModel:
+    def __init__(self, linevul_model, xgboost_model, alpha=0.5):
+        self.linevul_model = linevul_model
+        self.xgboost_model = xgboost_model
+        self.alpha = alpha
+
+    def predict(self, linevul_input, xgboost_input):
+        linevul_preds = self.linevul_model.predict(linevul_input)
+        xgboost_preds = self.xgboost_model.predict(xgboost_input)
+        ensemble_preds = self.alpha * linevul_preds + (1 - self.alpha) * xgboost_preds
+        return np.round(ensemble_preds).astype(int)
+
+    def predict_proba(self, linevul_input, xgboost_input):
+        linevul_proba = self.linevul_model.predict_proba(linevul_input)
+        xgboost_proba = self.xgboost_model.predict_proba(xgboost_input)
+        ensemble_proba = self.alpha * linevul_proba + (1 - self.alpha) * xgboost_proba
+        return ensemble_proba
--- a/src/training/train_linevul.py
+++ b/src/training/train_linevul.py
@ -0,0 +1,72 @@
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+from transformers import RobertaTokenizer
+from src.models.linevul import LineVulModel, LineVulTrainer
+from src.data.data_loader import DataLoader as CustomDataLoader
+from src.data.preprocessor import Preprocessor
+import numpy as np
+
+def prepare_linevul_data(data, tokenizer, max_length=512):
+    input_ids = []
+    attention_masks = []
+    labels = []
+
+    for item in data:
+        encoded = tokenizer.encode_plus(
+            ' '.join(item['tokens']),
+            add_special_tokens=True,
+            max_length=max_length,
+            padding='max_length',
+            truncation=True,
+            return_attention_mask=True,
+            return_tensors='pt'
+        )
+        input_ids.append(encoded['input_ids'])
+        attention_masks.append(encoded['attention_mask'])
+        labels.append(item['cwe_id'])
+
+    input_ids = torch.cat(input_ids, dim=0)
+    attention_masks = torch.cat(attention_masks, dim=0)
+    labels = torch.tensor(labels)
+
+    return TensorDataset(input_ids, attention_masks, labels)
+
+def train_linevul(config):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    data_loader = CustomDataLoader(config['data_dir'])
+    preprocessor = Preprocessor()
+
+    train_data = data_loader.load_jsonl('train.jsonl')
+    valid_data = data_loader.load_jsonl('valid.jsonl')
+
+    train_data = preprocessor.process_data(train_data)
+    valid_data = preprocessor.process_data(valid_data)
+
+    tokenizer = RobertaTokenizer.from_pretrained('microsoft/codebert-base')
+
+    train_dataset = prepare_linevul_data(train_data, tokenizer)
+    valid_dataset = prepare_linevul_data(valid_data, tokenizer)
+
+    train_dataloader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=config['batch_size'])
+
+    num_labels = len(set([item['cwe_id'] for item in train_data]))
+    model = LineVulModel(num_labels)
+
+    trainer = LineVulTrainer(model, device, learning_rate=config['learning_rate'])
+    trainer.train(train_dataloader, config['num_epochs'])
+
+    eval_loss, eval_preds, eval_labels = trainer.evaluate(valid_dataloader)
+    print(f"Validation Loss: {eval_loss:.4f}")
+
+    return model, trainer, (eval_preds, eval_labels)
+
+if __name__ == "__main__":
+    config = {
+        'data_dir': '../../data/raw',
+        'batch_size': 32,
+        'learning_rate': 2e-5,
+        'num_epochs': 5
+    }
+    train_linevul(config)
--- a/src/training/train_xgboost.py
+++ b/src/training/train_xgboost.py
@ -0,0 +1,69 @@
+# src/training/train_xgboost.py
+
+from src.models.xgboost_classifier import XGBoostVulClassifier, XGBoostVulRegressor
+from src.data.data_loader import DataLoader
+from src.data.preprocessor import Preprocessor
+from src.data.feature_extractor import FeatureExtractor
+import numpy as np
+from sklearn.model_selection import train_test_split
+
+def train_xgboost(config):
+    data_loader = DataLoader(config['data_dir'])
+    preprocessor = Preprocessor()
+    feature_extractor = FeatureExtractor()
+
+    train_data = data_loader.load_jsonl('train.jsonl')
+    train_data = preprocessor.process_data(train_data)
+
+    features = feature_extractor.extract_features(train_data)
+
+    X = features['features']
+    y_class = features['labels']
+    y_reg = features['cvss_scores']
+
+    X_train, X_val, y_class_train, y_class_val, y_reg_train, y_reg_val = train_test_split(
+        X, y_class, y_reg, test_size=0.2, random_state=42
+    )
+
+    classifier = XGBoostVulClassifier(
+        n_estimators=config['n_estimators'],
+        max_depth=config['max_depth'],
+        learning_rate=config['learning_rate']
+    )
+    classifier.train(X_train, y_class_train)
+
+    regressor = XGBoostVulRegressor(
+        n_estimators=config['n_estimators'],
+        max_depth=config['max_depth'],
+        learning_rate=config['learning_rate']
+    )
+    regressor.train(X_train, y_reg_train)
+
+    if config['perform_hyperparameter_tuning']:
+        param_grid = {
+            'n_estimators': [50, 100, 200],
+            'max_depth': [3, 5, 7],
+            'learning_rate': [0.01, 0.1, 0.3]
+        }
+        best_params_class, best_score_class = classifier.hyperparameter_tuning(X_train, y_class_train, param_grid)
+        best_params_reg, best_score_reg = regressor.hyperparameter_tuning(X_train, y_reg_train, param_grid)
+
+        print("Best Classification Parameters:", best_params_class)
+        print("Best Classification Score:", best_score_class)
+        print("Best Regression Parameters:", best_params_reg)
+        print("Best Regression Score:", best_score_reg)
+
+    class_preds = classifier.predict(X_val)
+    reg_preds = regressor.predict(X_val)
+
+    return classifier, regressor, (class_preds, y_class_val), (reg_preds, y_reg_val)
+
+if __name__ == "__main__":
+    config = {
+        'data_dir': '../../data/raw',
+        'n_estimators': 100,
+        'max_depth': 3,
+        'learning_rate': 0.1,
+        'perform_hyperparameter_tuning': True
+    }
+    train_xgboost(config)
--- a/src/utils/code_tokenizer.py
+++ b/src/utils/code_tokenizer.py
@ -0,0 +1,58 @@
+import re
+from typing import List
+
+class CodeTokenizer:
+    def __init__(self):
+        self.keywords = set(['if', 'else', 'for', 'while', 'return', 'break', 'continue', 'int', 'float', 'char', 'double', 'void'])
+        self.operators = set(['+', '-', '*', '/', '%', '=', '==', '!=', '<', '>', '<=', '>=', '&&', '||', '!', '&', '|', '^', '~', '<<', '>>', '++', '--'])
+        self.separators = set(['(', ')', '{', '}', '[', ']', ';', ',', '.'])
+
+    def tokenize(self, code: str) -> List[str]:
+        tokens = []
+        current_token = ''
+        
+        for char in code:
+            if char.isalnum() or char == '_':
+                current_token += char
+            else:
+                if current_token:
+                    tokens.append(self.classify_token(current_token))
+                    current_token = ''
+                if not char.isspace():
+                    tokens.append(self.classify_token(char))
+        
+        if current_token:
+            tokens.append(self.classify_token(current_token))
+        
+        return tokens
+
+    def classify_token(self, token: str) -> str:
+        if token in self.keywords:
+            return f'KEYWORD_{token.upper()}'
+        elif token in self.operators:
+            return f'OPERATOR_{token}'
+        elif token in self.separators:
+            return f'SEPARATOR_{token}'
+        elif token.isdigit():
+            return 'NUMBER'
+        elif re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', token):
+            return 'IDENTIFIER'
+        else:
+            return 'UNKNOWN'
+
+    def detokenize(self, tokens: List[str]) -> str:
+        code = ''
+        for token in tokens:
+            if token.startswith('KEYWORD_'):
+                code += token.split('_')[1].lower() + ' '
+            elif token.startswith('OPERATOR_'):
+                code += token.split('_')[1] + ' '
+            elif token.startswith('SEPARATOR_'):
+                code += token.split('_')[1]
+            elif token == 'NUMBER':
+                code += '0 '
+            elif token == 'IDENTIFIER':
+                code += 'id '
+            else:
+                code += token + ' '
+        return code.strip()
--- a/src/utils/data_augmentation.py
+++ b/src/utils/data_augmentation.py
@ -0,0 +1,69 @@
+import random
+from typing import List, Dict, Any
+
+class DataAugmentor:
+    def __init__(self, tokenizer: Any):
+        self.tokenizer = tokenizer
+
+    def augment_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+        augmented_sample = sample.copy()
+        tokens = self.tokenizer.tokenize(sample['function'])
+        
+        augmentation_methods = [
+            self.insert_random_token,
+            self.delete_random_token,
+            self.swap_random_tokens,
+            self.replace_with_synonym
+        ]
+        
+        method = random.choice(augmentation_methods)
+        augmented_tokens = method(tokens)
+        
+        augmented_sample['function'] = self.tokenizer.detokenize(augmented_tokens)
+        return augmented_sample
+
+    def insert_random_token(self, tokens: List[str]) -> List[str]:
+        if len(tokens) < 2:
+            return tokens
+        insert_position = random.randint(0, len(tokens) - 1)
+        insert_token = random.choice(tokens)
+        return tokens[:insert_position] + [insert_token] + tokens[insert_position:]
+
+    def delete_random_token(self, tokens: List[str]) -> List[str]:
+        if len(tokens) < 2:
+            return tokens
+        delete_position = random.randint(0, len(tokens) - 1)
+        return tokens[:delete_position] + tokens[delete_position + 1:]
+
+    def swap_random_tokens(self, tokens: List[str]) -> List[str]:
+        if len(tokens) < 2:
+            return tokens
+        pos1, pos2 = random.sample(range(len(tokens)), 2)
+        tokens[pos1], tokens[pos2] = tokens[pos2], tokens[pos1]
+        return tokens
+
+    def replace_with_synonym(self, tokens: List[str]) -> List[str]:
+        synonyms = {
+            'KEYWORD_IF': 'KEYWORD_WHEN',
+            'KEYWORD_ELSE': 'KEYWORD_OTHERWISE',
+            'KEYWORD_FOR': 'KEYWORD_REPEAT',
+            'KEYWORD_WHILE': 'KEYWORD_UNTIL',
+            'KEYWORD_RETURN': 'KEYWORD_YIELD',
+            'OPERATOR_+': 'OPERATOR_PLUS',
+            'OPERATOR_-': 'OPERATOR_MINUS',
+            'OPERATOR_*': 'OPERATOR_MULTIPLY',
+            'OPERATOR_/': 'OPERATOR_DIVIDE'
+        }
+        
+        for i, token in enumerate(tokens):
+            if token in synonyms:
+                tokens[i] = synonyms[token]
+        
+        return tokens
+
+    def augment_dataset(self, data: List[Dict[str, Any]], augmentation_factor: int = 2) -> List[Dict[str, Any]]:
+        augmented_data = data.copy()
+        for _ in range(augmentation_factor - 1):
+            for sample in data:
+                augmented_data.append(self.augment_sample(sample))
+        return augmented_data
--- a/src/visualization/dashboard.py
+++ b/src/visualization/dashboard.py
@ -0,0 +1,112 @@
+import dash
+import dash_core_components as dcc
+import dash_html_components as html
+from dash.dependencies import Input, Output
+import plotly.graph_objs as go
+import pandas as pd
+import numpy as np
+
+class VulnerabilityDashboard:
+    def __init__(self, data: pd.DataFrame):
+        self.data = data
+        self.app = dash.Dash(__name__)
+        self.setup_layout()
+        self.setup_callbacks()
+
+    def setup_layout(self):
+        self.app.layout = html.Div([
+            html.H1('Vulnerability Detection Dashboard'),
+            
+            html.Div([
+                html.Div([
+                    dcc.Dropdown(
+                        id='metric-dropdown',
+                        options=[
+                            {'label': 'Accuracy', 'value': 'accuracy'},
+                            {'label': 'Precision', 'value': 'precision'},
+                            {'label': 'Recall', 'value': 'recall'},
+                            {'label': 'F1 Score', 'value': 'f1_score'}
+                        ],
+                        value='accuracy'
+                    ),
+                    dcc.Graph(id='performance-graph')
+                ], className='six columns'),
+                
+                html.Div([
+                    dcc.Dropdown(
+                        id='vulnerability-type-dropdown',
+                        options=[{'label': vul_type, 'value': vul_type} for vul_type in self.data['vulnerability_type'].unique()],
+                        value=self.data['vulnerability_type'].unique()[0]
+                    ),
+                    dcc.Graph(id='vulnerability-distribution')
+                ], className='six columns')
+            ], className='row'),
+            
+            html.Div([
+                html.Div([
+                    dcc.Graph(id='confusion-matrix')
+                ], className='six columns'),
+                
+                html.Div([
+                    dcc.Graph(id='feature-importance')
+                ], className='six columns')
+            ], className='row')
+        ])
+
+    def setup_callbacks(self):
+        @self.app.callback(
+            Output('performance-graph', 'figure'),
+            [Input('metric-dropdown', 'value')]
+        )
+        def update_performance_graph(selected_metric):
+            return go.Figure(data=[
+                go.Scatter(x=self.data['date'], y=self.data[selected_metric], mode='lines+markers')
+            ], layout=go.Layout(title=f'{selected_metric.capitalize()} Over Time'))
+
+        @self.app.callback(
+            Output('vulnerability-distribution', 'figure'),
+            [Input('vulnerability-type-dropdown', 'value')]
+        )
+        def update_vulnerability_distribution(selected_type):
+            type_data = self.data[self.data['vulnerability_type'] == selected_type]
+            return go.Figure(data=[
+                go.Histogram(x=type_data['severity'])
+            ], layout=go.Layout(title=f'Severity Distribution for {selected_type}'))
+
+        @self.app.callback(
+            Output('confusion-matrix', 'figure'),
+            [Input('metric-dropdown', 'value')]
+        )
+        def update_confusion_matrix(selected_metric):
+            cm = np.random.randint(0, 100, size=(4, 4))  # Placeholder for actual confusion matrix
+            return go.Figure(data=[
+                go.Heatmap(z=cm, x=['Pred 0', 'Pred 1', 'Pred 2', 'Pred 3'],
+                           y=['True 0', 'True 1', 'True 2', 'True 3'])
+            ], layout=go.Layout(title='Confusion Matrix'))
+
+        @self.app.callback(
+            Output('feature-importance', 'figure'),
+            [Input('metric-dropdown', 'value')]
+        )
+        def update_feature_importance(selected_metric):
+            features = ['Feature 1', 'Feature 2', 'Feature 3', 'Feature 4', 'Feature 5']
+            importance = np.random.rand(5)  # Placeholder for actual feature importance
+            return go.Figure(data=[
+                go.Bar(x=features, y=importance)
+            ], layout=go.Layout(title='Feature Importance'))
+
+    def run_server(self):
+        self.app.run_server(debug=True)
+
+if __name__ == '__main__':
+    data = pd.DataFrame({
+        'date': pd.date_range(start='2023-01-01', periods=100),
+        'accuracy': np.random.rand(100),
+        'precision': np.random.rand(100),
+        'recall': np.random.rand(100),
+        'f1_score': np.random.rand(100),
+        'vulnerability_type': np.random.choice(['Type A', 'Type B', 'Type C'], 100),
+        'severity': np.random.randint(1, 6, 100)
+    })
+    dashboard = VulnerabilityDashboard(data)
+    dashboard.run_server()
--- a/src/visualization/plot_utils.py
+++ b/src/visualization/plot_utils.py
@ -0,0 +1,94 @@
+import matplotlib.pyplot as plt
+import seaborn as sns
+import numpy as np
+from sklearn.metrics import confusion_matrix
+from typing import List, Dict, Any
+
+class VisualizationUtils:
+    def __init__(self):
+        plt.style.use('seaborn')
+        sns.set_palette("deep")
+        plt.rcParams['figure.figsize'] = (12, 8)
+        plt.rcParams['font.size'] = 12
+
+    def plot_confusion_matrix(self, y_true: np.ndarray, y_pred: np.ndarray, labels: List[str], title: str):
+        cm = confusion_matrix(y_true, y_pred)
+        plt.figure(figsize=(12, 10))
+        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=labels, yticklabels=labels)
+        plt.title(title)
+        plt.xlabel('Predicted')
+        plt.ylabel('True')
+        plt.tight_layout()
+        plt.savefig(f'outputs/figures/{title.lower().replace(" ", "_")}.png')
+        plt.close()
+
+    def plot_feature_importance(self, feature_importance: np.ndarray, feature_names: List[str], title: str):
+        sorted_idx = np.argsort(feature_importance)
+        pos = np.arange(sorted_idx.shape[0]) + .5
+        
+        plt.figure(figsize=(12, len(feature_names) * 0.3))
+        plt.barh(pos, feature_importance[sorted_idx], align='center')
+        plt.yticks(pos, np.array(feature_names)[sorted_idx])
+        plt.title(title)
+        plt.xlabel('Importance')
+        plt.tight_layout()
+        plt.savefig(f'outputs/figures/{title.lower().replace(" ", "_")}.png')
+        plt.close()
+
+    def plot_learning_curve(self, train_scores: List[float], val_scores: List[float], metric: str):
+        plt.figure()
+        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}')
+        plt.xlabel('Epoch')
+        plt.ylabel(metric)
+        plt.legend()
+        plt.tight_layout()
+        plt.savefig(f'outputs/figures/learning_curve_{metric.lower()}.png')
+        plt.close()
+
+    def plot_roc_curve(self, fpr: np.ndarray, tpr: np.ndarray, roc_auc: float):
+        plt.figure()
+        plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.2f})')
+        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
+        plt.xlim([0.0, 1.0])
+        plt.ylim([0.0, 1.05])
+        plt.xlabel('False Positive Rate')
+        plt.ylabel('True Positive Rate')
+        plt.title('Receiver Operating Characteristic (ROC) Curve')
+        plt.legend(loc="lower right")
+        plt.tight_layout()
+        plt.savefig('outputs/figures/roc_curve.png')
+        plt.close()
+
+    def plot_precision_recall_curve(self, precision: np.ndarray, recall: np.ndarray, average_precision: float):
+        plt.figure()
+        plt.step(recall, precision, color='b', alpha=0.2, where='post')
+        plt.fill_between(recall, precision, step='post', alpha=0.2, color='b')
+        plt.xlabel('Recall')
+        plt.ylabel('Precision')
+        plt.ylim([0.0, 1.05])
+        plt.xlim([0.0, 1.0])
+        plt.title(f'Precision-Recall curve: AP={average_precision:.2f}')
+        plt.tight_layout()
+        plt.savefig('outputs/figures/precision_recall_curve.png')
+        plt.close()
+
+    def plot_distribution(self, data: np.ndarray, title: str):
+        plt.figure()
+        sns.histplot(data, kde=True)
+        plt.title(title)
+        plt.xlabel('Value')
+        plt.ylabel('Frequency')
+        plt.tight_layout()
+        plt.savefig(f'outputs/figures/{title.lower().replace(" ", "_")}_distribution.png')
+        plt.close()
+
+    def plot_correlation_matrix(self, data: np.ndarray, feature_names: List[str], title: str):
+        corr = np.corrcoef(data.T)
+        plt.figure(figsize=(12, 10))
+        sns.heatmap(corr, annot=True, fmt='.2f', cmap='coolwarm', xticklabels=feature_names, yticklabels=feature_names)
+        plt.title(title)
+        plt.tight_layout()
+        plt.savefig(f'outputs/figures/{title.lower().replace(" ", "_")}_correlation.png')
+        plt.close()