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,19 @@
+# CVSS Hyperspace Evaluator
+
+## 环境配置
+
+1. 安装 Docker 和 NVIDIA Container Toolkit。
+2. 构建 Docker 镜像：
+   ```
+   docker build -t cvss-evaluator .
+   ```
+
+## 运行代码
+
+1. 准备配置文件 `config.json`。
+2. 运行 Docker 容器：
+   ```
+   docker run --gpus all -v $(pwd)/data:/app/data -v $(pwd)/results:/app/results cvss-evaluator
+   ```
+
+结果将保存在 `./results` 目录中。

config.json

@@ -0,0 +1,34 @@
+{
+    "model_name": "meta-llama/Llama-2-7b-chat-hf",
+    "data_path": "./data/SIR_train_set.json",
+    "prompt_template_path": "./templates/cvss_prompt_template.json",
+    "output_path": "./results/evaluation_summary.json",
+    "detailed_results_path": "./results/detailed_results.csv",
+    "plot_path": "./results/prediction_vs_actual.png",
+    "batch_size": 32,
+    "tensor_parallel_size": 2,
+    "gpu_memory_utilization": 0.9,
+    "max_length": 2048,
+    "temperature": 0.7,
+    "top_p": 0.95,
+    "train_test_split": 0.2,
+    "validation_split": 0.1,
+    "random_seed": 42,
+    "logging": {
+      "level": "INFO",
+      "file": "./logs/cvss_evaluation.log"
+    },
+    "model_save_path": "./models/cvss_predictor.joblib",
+    "vectorizer_save_path": "./models/tfidf_vectorizer.joblib",
+    "feature_importance_path": "./results/feature_importance.json",
+    "confidence_interval": 0.95,
+    "memory_optimization": {
+      "target_memory_usage": 0.8,
+      "clear_cache_frequency": 100
+    },
+    "evaluation_metrics": ["mse", "rmse", "mae", "r2"],
+    "vulnerability_analysis": {
+      "severity_threshold": 7.0,
+      "top_n_features": 10
+    }
+  }

cvss_analysis/__init__.py

@@ -0,0 +1,4 @@
+from .vulnerability_analyzer import VulnerabilityAnalyzer
+from .score_predictor import ScorePredictor
+
+__all__ = ['VulnerabilityAnalyzer', 'ScorePredictor']

cvss_analysis/score_predictor.py

@@ -0,0 +1,78 @@
+import numpy as np
+from typing import Dict, Any, List
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import mean_squared_error
+import joblib
+
+class ScorePredictor:
+    def __init__(self):
+        self.vectorizer = TfidfVectorizer(max_features=1000)
+        self.model = RandomForestRegressor(n_estimators=100, random_state=42)
+        self.trained = False
+
+    def preprocess_data(self, data: List[Dict[str, Any]]) -> Dict[str, Any]:
+        descriptions = [item['description'] for item in data]
+        scores = [item['baseScore'] for item in data]
+        X = self.vectorizer.fit_transform(descriptions)
+        return {"X": X, "y": scores}
+
+    def train(self, data: List[Dict[str, Any]]):
+        processed_data = self.preprocess_data(data)
+        X_train, X_test, y_train, y_test = train_test_split(
+            processed_data["X"], processed_data["y"], test_size=0.2, random_state=42
+        )
+        self.model.fit(X_train, y_train)
+        self.trained = True
+        train_score = self.model.score(X_train, y_train)
+        test_score = self.model.score(X_test, y_test)
+        return {"train_score": train_score, "test_score": test_score}
+
+    def predict(self, description: str) -> float:
+        if not self.trained:
+            raise ValueError("Model not trained. Please train the model first.")
+        X = self.vectorizer.transform([description])
+        return self.model.predict(X)[0]
+
+    def batch_predict(self, descriptions: List[str]) -> List[float]:
+        if not self.trained:
+            raise ValueError("Model not trained. Please train the model first.")
+        X = self.vectorizer.transform(descriptions)
+        return self.model.predict(X).tolist()
+
+    def evaluate(self, data: List[Dict[str, Any]]) -> Dict[str, float]:
+        processed_data = self.preprocess_data(data)
+        predictions = self.model.predict(processed_data["X"])
+        mse = mean_squared_error(processed_data["y"], predictions)
+        rmse = np.sqrt(mse)
+        return {"mse": mse, "rmse": rmse}
+
+    def save_model(self, model_path: str, vectorizer_path: str):
+        joblib.dump(self.model, model_path)
+        joblib.dump(self.vectorizer, vectorizer_path)
+
+    def load_model(self, model_path: str, vectorizer_path: str):
+        self.model = joblib.load(model_path)
+        self.vectorizer = joblib.load(vectorizer_path)
+        self.trained = True
+
+    def feature_importance(self) -> Dict[str, float]:
+        if not self.trained:
+            raise ValueError("Model not trained. Please train the model first.")
+        feature_names = self.vectorizer.get_feature_names_out()
+        importances = self.model.feature_importances_
+        return dict(zip(feature_names, importances))
+
+    def predict_with_confidence(self, description: str) -> Dict[str, float]:
+        if not self.trained:
+            raise ValueError("Model not trained. Please train the model first.")
+        X = self.vectorizer.transform([description])
+        predictions = [tree.predict(X) for tree in self.model.estimators_]
+        mean_prediction = np.mean(predictions)
+        std_prediction = np.std(predictions)
+        return {
+            "prediction": mean_prediction,
+            "confidence_interval": (mean_prediction - 1.96 * std_prediction,
+                                    mean_prediction + 1.96 * std_prediction)
+        }

cvss_analysis/vulnerability_analyzer.py

@@ -0,0 +1,86 @@
+import re
+from typing import Dict, Any, List
+from collections import Counter
+
+class VulnerabilityAnalyzer:
+    def __init__(self):
+        self.vulnerability_types = [
+            "Buffer Overflow", "SQL Injection", "Cross-Site Scripting (XSS)",
+            "Remote Code Execution", "Denial of Service", "Information Disclosure",
+            "Privilege Escalation", "Memory Corruption", "Use After Free"
+        ]
+
+    def analyze_vulnerability(self, description: str) -> Dict[str, Any]:
+        analysis = {
+            "detected_types": self.detect_vulnerability_types(description),
+            "severity_indicators": self.extract_severity_indicators(description),
+            "affected_components": self.extract_affected_components(description),
+            "potential_impact": self.assess_potential_impact(description)
+        }
+        return analysis
+
+    def detect_vulnerability_types(self, description: str) -> List[str]:
+        detected_types = []
+        for vuln_type in self.vulnerability_types:
+            if re.search(vuln_type.lower(), description.lower()):
+                detected_types.append(vuln_type)
+        return detected_types
+
+    def extract_severity_indicators(self, description: str) -> Dict[str, int]:
+        severity_words = ["critical", "high", "medium", "low", "severe", "important"]
+        return {word: description.lower().count(word) for word in severity_words}
+
+    def extract_affected_components(self, description: str) -> List[str]:
+        component_pattern = r"in\s+(\w+(?:\s+\w+){0,2})"
+        matches = re.findall(component_pattern, description)
+        return list(set(matches))
+
+    def assess_potential_impact(self, description: str) -> List[str]:
+        impact_keywords = {
+            "data loss": ["data", "loss", "leak"],
+            "system compromise": ["compromise", "take over", "control"],
+            "information disclosure": ["disclose", "reveal", "expose"],
+            "service disruption": ["disrupt", "interrupt", "unavailable"]
+        }
+        potential_impacts = []
+        for impact, keywords in impact_keywords.items():
+            if any(keyword in description.lower() for keyword in keywords):
+                potential_impacts.append(impact)
+        return potential_impacts
+
+    def calculate_complexity_score(self, description: str) -> float:
+        complexity_indicators = {
+            "complex": 3,
+            "difficult": 2,
+            "easy": -1,
+            "simple": -2,
+            "straightforward": -1
+        }
+        score = 0
+        for indicator, value in complexity_indicators.items():
+            score += description.lower().count(indicator) * value
+        return max(0, min(10, score + 5))  # Normalize to 0-10 scale
+
+    def extract_cwe_ids(self, description: str) -> List[str]:
+        cwe_pattern = r"CWE-(\d+)"
+        return re.findall(cwe_pattern, description)
+
+    def analyze_attack_vector(self, description: str) -> str:
+        vectors = {
+            "network": ["remote", "network", "internet"],
+            "adjacent": ["adjacent", "nearby", "local network"],
+            "local": ["local", "physical access"],
+            "physical": ["physical", "direct access"]
+        }
+        for vector, keywords in vectors.items():
+            if any(keyword in description.lower() for keyword in keywords):
+                return vector
+        return "unknown"
+
+    def generate_summary(self, analysis: Dict[str, Any]) -> str:
+        summary = f"Vulnerability Analysis Summary:\n"
+        summary += f"- Detected Types: {', '.join(analysis['detected_types'])}\n"
+        summary += f"- Severity Indicators: {dict(analysis['severity_indicators'])}\n"
+        summary += f"- Affected Components: {', '.join(analysis['affected_components'])}\n"
+        summary += f"- Potential Impact: {', '.join(analysis['potential_impact'])}\n"
+        return summary

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
+    }
+}

data_handlers/__init__.py

@@ -0,0 +1,4 @@
+from .json_parser import JSONParser
+from .dataset_preprocessor import DatasetPreprocessor
+
+__all__ = ['JSONParser', 'DatasetPreprocessor']

data_handlers/dataset_preprocessor.py

@@ -0,0 +1,72 @@
+from typing import List, Dict, Any
+import re
+import numpy as np
+from sklearn.model_selection import train_test_split
+
+class DatasetPreprocessor:
+    def __init__(self, data: List[Dict[str, Any]]):
+        self.data = data
+
+    def clean_text(self, text: str) -> str:
+        text = re.sub(r'<[^>]+>', '', text)
+        text = re.sub(r'http\S+', '', text)
+        text = re.sub(r'\s+', ' ', text).strip()
+        return text
+
+    def preprocess_data(self) -> List[Dict[str, Any]]:
+        for item in self.data:
+            if 'description' in item:
+                item['description'] = self.clean_text(item['description'])
+        return self.data
+
+    def extract_features(self) -> List[Dict[str, Any]]:
+        for item in self.data:
+            item['description_length'] = len(item.get('description', ''))
+            item['has_cve'] = 1 if 'CVE' in item.get('CVE_ID', '') else 0
+        return self.data
+
+    def normalize_scores(self, field: str = 'baseScore') -> List[Dict[str, Any]]:
+        scores = [item.get(field, 0) for item in self.data]
+        min_score, max_score = min(scores), max(scores)
+        for item in self.data:
+            if field in item:
+                item[f'normalized_{field}'] = (item[field] - min_score) / (max_score - min_score)
+        return self.data
+
+    def split_dataset(self, test_size: float = 0.2, validation_size: float = 0.1, random_state: int = 42) -> Dict[str, List[Dict[str, Any]]]:
+        train_val, test = train_test_split(self.data, test_size=test_size, random_state=random_state)
+        train, val = train_test_split(train_val, test_size=validation_size/(1-test_size), random_state=random_state)
+        return {
+            'train': train,
+            'validation': val,
+            'test': test
+        }
+
+    def balance_dataset(self, field: str, method: str = 'undersample') -> List[Dict[str, Any]]:
+        value_counts = {}
+        for item in self.data:
+            value = item.get(field)
+            value_counts[value] = value_counts.get(value, 0) + 1
+        
+        min_count = min(value_counts.values())
+        
+        balanced_data = []
+        for value, count in value_counts.items():
+            items = [item for item in self.data if item.get(field) == value]
+            if method == 'undersample':
+                balanced_data.extend(np.random.choice(items, min_count, replace=False))
+            elif method == 'oversample':
+                balanced_data.extend(np.random.choice(items, max(count, min_count), replace=True))
+        
+        return balanced_data
+
+    def create_data_batches(self, batch_size: int) -> List[List[Dict[str, Any]]]:
+        return [self.data[i:i + batch_size] for i in range(0, len(self.data), batch_size)]
+
+    def get_data_summary(self) -> Dict[str, Any]:
+        return {
+            "total_samples": len(self.data),
+            "unique_cve_ids": len(set(item.get('CVE_ID') for item in self.data)),
+            "avg_description_length": np.mean([len(item.get('description', '')) for item in self.data]),
+            "score_distribution": np.histogram([item.get('baseScore', 0) for item in self.data], bins=10)
+        }

data_handlers/json_parser.py

@@ -0,0 +1,53 @@
+import json
+from typing import List, Dict, Any
+import os
+
+class JSONParser:
+    def __init__(self):
+        self.data = []
+
+    def load_file(self, file_path: str) -> None:
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"File not found: {file_path}")
+        
+        with open(file_path, 'r', encoding='utf-8') as file:
+            self.data = json.load(file)
+
+    def load_multiple_files(self, file_paths: List[str]) -> None:
+        for file_path in file_paths:
+            self.load_file(file_path)
+
+    def get_data(self) -> List[Dict[str, Any]]:
+        return self.data
+
+    def filter_data(self, condition: callable) -> List[Dict[str, Any]]:
+        return list(filter(condition, self.data))
+
+    def get_field_values(self, field: str) -> List[Any]:
+        return [item.get(field) for item in self.data if field in item]
+
+    def save_to_file(self, file_path: str) -> None:
+        with open(file_path, 'w', encoding='utf-8') as file:
+            json.dump(self.data, file, indent=2, ensure_ascii=False)
+
+    def merge_json_files(self, file_paths: List[str], output_path: str) -> None:
+        merged_data = []
+        for file_path in file_paths:
+            with open(file_path, 'r', encoding='utf-8') as file:
+                merged_data.extend(json.load(file))
+        
+        with open(output_path, 'w', encoding='utf-8') as file:
+            json.dump(merged_data, file, indent=2, ensure_ascii=False)
+
+    def validate_json_structure(self, required_fields: List[str]) -> bool:
+        for item in self.data:
+            if not all(field in item for field in required_fields):
+                return False
+        return True
+
+    def get_statistics(self) -> Dict[str, Any]:
+        return {
+            "total_items": len(self.data),
+            "unique_keys": list(set(key for item in self.data for key in item.keys())),
+            "sample_item": self.data[0] if self.data else None
+        }

dockerfile

@@ -0,0 +1,69 @@
+FROM quay.io/centos/centos:stream8
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PATH="/root/miniconda3/bin:${PATH}"
+ARG PATH="/root/miniconda3/bin:${PATH}"
+
+RUN dnf -y update && dnf -y install \
+    wget \
+    git \
+    gcc \
+    gcc-c++ \
+    make \
+    cmake \
+    boost-devel \
+    gperftools-devel \
+    protobuf-devel \
+    openblas-devel \
+    atlas-devel \
+    lapack-devel \
+    blas-devel \
+    gfortran \
+    pkgconfig \
+    hdf5-devel \
+    lmdb-devel \
+    leveldb-devel \
+    snappy-devel \
+    libjpeg-turbo-devel \
+    libpng-devel \
+    libtiff-devel \
+    libwebp-devel \
+    opencv-devel \
+    zeromq-devel \
+    zlib-devel \
+    jemalloc-devel \
+    && dnf clean all
+
+# Install CUDA
+RUN dnf -y install https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-repo-rhel8-11.7.1-1.x86_64.rpm
+RUN dnf -y module install nvidia-driver:latest-dkms
+RUN dnf -y install cuda-11-7
+ENV PATH="/usr/local/cuda-11.7/bin:${PATH}"
+ENV LD_LIBRARY_PATH="/usr/local/cuda-11.7/lib64:${LD_LIBRARY_PATH}"
+
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
+    && mkdir /root/.conda \
+    && bash Miniconda3-latest-Linux-x86_64.sh -b \
+    && rm -f Miniconda3-latest-Linux-x86_64.sh
+
+RUN conda create -n cvss_env python=3.9 -y
+ENV PATH /root/miniconda3/envs/cvss_env/bin:$PATH
+
+COPY requirements.txt /tmp/
+RUN pip install --no-cache-dir -r /tmp/requirements.txt
+
+RUN conda install -c pytorch magma-cuda117 -y
+RUN conda install -c conda-forge gperftools -y
+
+WORKDIR /app
+COPY . /app
+
+RUN git clone https://github.com/NVIDIA/apex && \
+    cd apex && \
+    pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
+
+RUN pip install --no-cache-dir git+https://github.com/huggingface/transformers.git@main
+
+ENV LD_PRELOAD=/usr/lib64/libtcmalloc.so.4
+
+CMD ["python", "main.py", "--config", "config.json"]

evaluation/__init__.py

@@ -0,0 +1,4 @@
+from .mse_calculator import MSECalculator
+from .rmse_calculator import RMSECalculator
+
+__all__ = ['MSECalculator', 'RMSECalculator']

evaluation/mse_calculator.py

@@ -0,0 +1,50 @@
+import numpy as np
+from typing import List, Dict, Any
+
+class MSECalculator:
+    def __init__(self):
+        self.predictions = []
+        self.actual_values = []
+
+    def add_sample(self, prediction: float, actual: float):
+        self.predictions.append(prediction)
+        self.actual_values.append(actual)
+
+    def add_batch(self, predictions: List[float], actual_values: List[float]):
+        if len(predictions) != len(actual_values):
+            raise ValueError("Predictions and actual values must have the same length")
+        self.predictions.extend(predictions)
+        self.actual_values.extend(actual_values)
+
+    def calculate_mse(self) -> float:
+        if len(self.predictions) == 0:
+            raise ValueError("No samples added yet")
+        return np.mean(np.square(np.array(self.predictions) - np.array(self.actual_values)))
+
+    def calculate_mse_per_sample(self) -> List[float]:
+        return [(pred - actual) ** 2 for pred, actual in zip(self.predictions, self.actual_values)]
+
+    def get_summary(self) -> Dict[str, Any]:
+        mse = self.calculate_mse()
+        return {
+            "mse": mse,
+            "samples": len(self.predictions),
+            "max_error": max(self.calculate_mse_per_sample()),
+            "min_error": min(self.calculate_mse_per_sample())
+        }
+
+    def reset(self):
+        self.predictions.clear()
+        self.actual_values.clear()
+
+    def get_data(self) -> Dict[str, List[float]]:
+        return {
+            "predictions": self.predictions,
+            "actual_values": self.actual_values
+        }
+
+    def calculate_weighted_mse(self, weights: List[float]) -> float:
+        if len(weights) != len(self.predictions):
+            raise ValueError("Number of weights must match number of samples")
+        squared_errors = np.square(np.array(self.predictions) - np.array(self.actual_values))
+        return np.average(squared_errors, weights=weights)

evaluation/rmse_calculator.py

@@ -0,0 +1,63 @@
+import numpy as np
+from typing import List, Dict, Any
+
+class RMSECalculator:
+    def __init__(self):
+        self.predictions = []
+        self.actual_values = []
+
+    def add_sample(self, prediction: float, actual: float):
+        self.predictions.append(prediction)
+        self.actual_values.append(actual)
+
+    def add_batch(self, predictions: List[float], actual_values: List[float]):
+        if len(predictions) != len(actual_values):
+            raise ValueError("Predictions and actual values must have the same length")
+        self.predictions.extend(predictions)
+        self.actual_values.extend(actual_values)
+
+    def calculate_rmse(self) -> float:
+        if len(self.predictions) == 0:
+            raise ValueError("No samples added yet")
+        return np.sqrt(np.mean(np.square(np.array(self.predictions) - np.array(self.actual_values))))
+
+    def calculate_rmse_per_sample(self) -> List[float]:
+        return [np.sqrt((pred - actual) ** 2) for pred, actual in zip(self.predictions, self.actual_values)]
+
+    def get_summary(self) -> Dict[str, Any]:
+        rmse = self.calculate_rmse()
+        return {
+            "rmse": rmse,
+            "samples": len(self.predictions),
+            "max_error": max(self.calculate_rmse_per_sample()),
+            "min_error": min(self.calculate_rmse_per_sample())
+        }
+
+    def reset(self):
+        self.predictions.clear()
+        self.actual_values.clear()
+
+    def get_data(self) -> Dict[str, List[float]]:
+        return {
+            "predictions": self.predictions,
+            "actual_values": self.actual_values
+        }
+
+    def calculate_weighted_rmse(self, weights: List[float]) -> float:
+        if len(weights) != len(self.predictions):
+            raise ValueError("Number of weights must match number of samples")
+        squared_errors = np.square(np.array(self.predictions) - np.array(self.actual_values))
+        return np.sqrt(np.average(squared_errors, weights=weights))
+
+    def calculate_rmse_with_confidence_interval(self, confidence_level: float = 0.95) -> Dict[str, float]:
+        rmse = self.calculate_rmse()
+        n = len(self.predictions)
+        se = rmse / np.sqrt(2 * n)
+        t_value = np.abs(np.random.standard_t(df=n-1, size=10000)).mean()  # Approximating t-value
+        margin_of_error = t_value * se
+        return {
+            "rmse": rmse,
+            "lower_bound": max(0, rmse - margin_of_error),
+            "upper_bound": rmse + margin_of_error,
+            "confidence_level": confidence_level
+        }

llama_core/__init__.py

@@ -0,0 +1,4 @@
+from .model_loader import ModelLoader
+from .inference_engine import InferenceEngine
+
+__all__ = ['ModelLoader', 'InferenceEngine']

llama_core/inference_engine.py

@@ -0,0 +1,82 @@
+import torch
+from typing import List, Dict, Any
+from .model_loader import ModelLoader
+
+class InferenceEngine:
+    def __init__(self, model_loader: ModelLoader, max_length: int = 2048, temperature: float = 0.7, top_p: float = 0.95):
+        self.model_loader = model_loader
+        self.max_length = max_length
+        self.temperature = temperature
+        self.top_p = top_p
+
+    def generate(self, prompt: str) -> str:
+        model, tokenizer = self.model_loader.load()
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=self.max_length,
+                temperature=self.temperature,
+                top_p=self.top_p,
+                do_sample=True
+            )
+        
+        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    def batch_generate(self, prompts: List[str]) -> List[str]:
+        model, tokenizer = self.model_loader.load()
+        batch_inputs = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True).to(model.device)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **batch_inputs,
+                max_length=self.max_length,
+                temperature=self.temperature,
+                top_p=self.top_p,
+                do_sample=True,
+                num_return_sequences=1
+            )
+        
+        return [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
+
+    def extract_cvss_score(self, generated_text: str) -> float:
+        try:
+            score_text = generated_text.split("CVSS Score:")[-1].strip().split()[0]
+            return float(score_text)
+        except:
+            return 0.0
+
+    def infer_cvss_scores(self, vulnerability_descriptions: List[str]) -> List[Dict[str, Any]]:
+        prompts = [f"Analyze the following vulnerability and provide a CVSS score:\n\n{desc}\n\nCVSS Score:" for desc in vulnerability_descriptions]
+        generated_texts = self.batch_generate(prompts)
+        
+        results = []
+        for desc, text in zip(vulnerability_descriptions, generated_texts):
+            score = self.extract_cvss_score(text)
+            results.append({
+                "description": desc,
+                "generated_text": text,
+                "cvss_score": score
+            })
+        
+        return results
+
+    def update_generation_params(self, max_length: int = None, temperature: float = None, top_p: float = None):
+        if max_length is not None:
+            self.max_length = max_length
+        if temperature is not None:
+            self.temperature = temperature
+        if top_p is not None:
+            self.top_p = top_p
+
+    def get_generation_params(self) -> Dict[str, Any]:
+        return {
+            "max_length": self.max_length,
+            "temperature": self.temperature,
+            "top_p": self.top_p
+        }
+
+    def warm_up(self):
+        dummy_prompt = "This is a warm-up prompt to initialize the model."
+        self.generate(dummy_prompt)

llama_core/model_loader.py

@@ -0,0 +1,54 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from typing import Tuple, Optional
+
+class ModelLoader:
+    def __init__(self, model_name: str, device: str = "cuda", low_cpu_mem_usage: bool = True):
+        self.model_name = model_name
+        self.device = device
+        self.low_cpu_mem_usage = low_cpu_mem_usage
+        self.model = None
+        self.tokenizer = None
+
+    def load(self) -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
+        if self.model is None or self.tokenizer is None:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float16,
+                low_cpu_mem_usage=self.low_cpu_mem_usage,
+                device_map="auto"
+            )
+        return self.model, self.tokenizer
+
+    def unload(self):
+        if self.model is not None:
+            del self.model
+        if self.tokenizer is not None:
+            del self.tokenizer
+        torch.cuda.empty_cache()
+        self.model = None
+        self.tokenizer = None
+
+    def reload(self) -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
+        self.unload()
+        return self.load()
+
+    def get_model_info(self) -> dict:
+        if self.model is None:
+            return {"error": "Model not loaded"}
+        return {
+            "model_name": self.model_name,
+            "device": self.device,
+            "parameters": sum(p.numel() for p in self.model.parameters()),
+            "loaded": self.model is not None
+        }
+
+    def is_loaded(self) -> bool:
+        return self.model is not None and self.tokenizer is not None
+
+    def get_tokenizer(self) -> Optional[AutoTokenizer]:
+        return self.tokenizer
+
+    def get_model(self) -> Optional[AutoModelForCausalLM]:
+        return self.model

main.py

@@ -0,0 +1,95 @@
+import argparse
+import json
+from data_handlers import JSONParser, DatasetPreprocessor
+from llama_core import ModelLoader, InferenceEngine
+from cvss_analysis import VulnerabilityAnalyzer, ScorePredictor
+from evaluation import MSECalculator, RMSECalculator
+from utils import PromptGenerator, ResultFormatter
+from vllm_accelerator import ParallelProcessor, MemoryOptimizer
+
+def load_config(config_path):
+    with open(config_path, 'r') as f:
+        return json.load(f)
+
+def main(config_path):
+    config = load_config(config_path)
+
+    # Initialize components
+    json_parser = JSONParser()
+    dataset_preprocessor = DatasetPreprocessor([])
+    model_loader = ModelLoader(config['model_name'])
+    inference_engine = InferenceEngine(model_loader)
+    vulnerability_analyzer = VulnerabilityAnalyzer()
+    score_predictor = ScorePredictor()
+    mse_calculator = MSECalculator()
+    rmse_calculator = RMSECalculator()
+    prompt_generator = PromptGenerator(config['prompt_template_path'])
+    result_formatter = ResultFormatter()
+    parallel_processor = ParallelProcessor(config['model_name'])
+    memory_optimizer = MemoryOptimizer()
+
+    # Load and preprocess data
+    json_parser.load_file(config['data_path'])
+    raw_data = json_parser.get_data()
+    dataset_preprocessor.data = raw_data
+    preprocessed_data = dataset_preprocessor.preprocess_data()
+    
+    # Split dataset
+    dataset_splits = dataset_preprocessor.split_dataset()
+    train_data, val_data, test_data = dataset_splits['train'], dataset_splits['validation'], dataset_splits['test']
+
+    # Train score predictor
+    score_predictor.train(train_data)
+
+    # Generate prompts for test data
+    test_prompts = [prompt_generator.generate_prompt(item) for item in test_data]
+
+    # Use parallel processor for inference
+    batch_size = config['batch_size']
+    all_predictions = []
+    for i in range(0, len(test_prompts), batch_size):
+        batch_prompts = test_prompts[i:i+batch_size]
+        batch_results = parallel_processor.process_batch(batch_prompts)
+        all_predictions.extend([inference_engine.extract_cvss_score(result['generated_text']) for result in batch_results])
+
+    # Evaluate predictions
+    actual_scores = [item['baseScore'] for item in test_data]
+    mse_calculator.add_batch(all_predictions, actual_scores)
+    rmse_calculator.add_batch(all_predictions, actual_scores)
+
+    mse = mse_calculator.calculate_mse()
+    rmse = rmse_calculator.calculate_rmse()
+
+    # Analyze vulnerabilities
+    for item, prediction in zip(test_data, all_predictions):
+        analysis = vulnerability_analyzer.analyze_vulnerability(item['description'])
+        result_formatter.add_result(prediction, item['baseScore'], {
+            'cve_id': item['CVE_ID'],
+            'analysis': analysis
+        })
+
+    # Generate and save results
+    results_summary = {
+        'mse': mse,
+        'rmse': rmse,
+        'model_info': model_loader.get_model_info(),
+        'data_summary': dataset_preprocessor.get_data_summary(),
+        'memory_usage': memory_optimizer.get_memory_usage()
+    }
+
+    with open(config['output_path'], 'w') as f:
+        json.dump(results_summary, f, indent=2)
+
+    result_formatter.to_csv(config['detailed_results_path'])
+    result_formatter.plot_prediction_vs_actual(config['plot_path'])
+
+    print(f"Evaluation complete. MSE: {mse}, RMSE: {rmse}")
+    print(f"Results saved to {config['output_path']}")
+    print(f"Detailed results saved to {config['detailed_results_path']}")
+    print(f"Plot saved to {config['plot_path']}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="CVSS Score Prediction System")
+    parser.add_argument('--config', type=str, required=True, help='Path to configuration file')
+    args = parser.parse_args()
+    main(args.config)

requirements.txt

@@ -0,0 +1,49 @@
+torch==2.0.1+cu117
+transformers==4.30.2
+vllm==0.1.4
+numpy==1.23.5
+pandas==2.0.3
+scikit-learn==1.3.0
+matplotlib==3.7.2
+seaborn==0.12.2
+joblib==1.3.1
+tqdm==4.65.0
+pyarrow==12.0.1
+fastapi==0.100.0
+uvicorn[standard]==0.22.0
+gunicorn==20.1.0
+pydantic==1.10.11
+redis==4.6.0
+pymongo==4.4.1
+psycopg2-binary==2.9.6
+sqlalchemy==2.0.19
+alembic==1.11.1
+pytest==7.4.0
+hypothesis==6.82.0
+black==23.7.0
+isort==5.12.0
+mypy==1.4.1
+flake8==6.0.0
+pre-commit==3.3.3
+wandb==0.15.5
+mlflow==2.5.0
+optuna==3.2.0
+ray[tune]==2.5.1
+deepspeed==0.9.5
+fairscale==0.4.13
+bitsandbytes==0.41.1
+triton==2.0.0
+einops==0.6.1
+xformers==0.0.20
+accelerate==0.21.0
+datasets==2.13.1
+evaluate==0.4.0
+tokenizers==0.13.3
+sentencepiece==0.1.99
+sacremoses==0.0.53
+protobuf==3.20.3
+onnx==1.14.0
+onnxruntime-gpu==1.15.1
+tensorrt==8.6.1
+pycuda==2022.2.2
+cupy-cuda11x==12.2.0

utils/__init__.py

@@ -0,0 +1,4 @@
+from .prompt_generator import PromptGenerator
+from .result_formatter import ResultFormatter
+
+__all__ = ['PromptGenerator', 'ResultFormatter']

utils/prompt_generator.py

@@ -0,0 +1,31 @@
+import json
+from typing import Dict, Any
+
+class PromptGenerator:
+    def __init__(self, template_path: str):
+        with open(template_path, 'r') as f:
+            self.template = json.load(f)
+
+    def generate_prompt(self, vulnerability_data: Dict[str, Any]) -> str:
+        prompt = self.template['prefix']
+        prompt += f"CVE ID: {vulnerability_data.get('CVE_ID', 'N/A')}\n"
+        prompt += f"Description: {vulnerability_data.get('description', 'N/A')}\n"
+        prompt += f"Issue URL: {vulnerability_data.get('Issue_Url_new', 'N/A')}\n"
+        prompt += f"Repository: {vulnerability_data.get('Repo_new', 'N/A')}\n"
+        prompt += f"Created At: {vulnerability_data.get('Issue_Created_At', 'N/A')}\n"
+        prompt += self.template['task_description']
+        prompt += self.template['suffix']
+        return prompt
+
+    def generate_batch_prompts(self, vulnerability_data_list: List[Dict[str, Any]]) -> List[str]:
+        return [self.generate_prompt(data) for data in vulnerability_data_list]
+
+    def update_template(self, new_template: Dict[str, str]):
+        self.template.update(new_template)
+
+    def get_template(self) -> Dict[str, str]:
+        return self.template.copy()
+
+    def reset_template(self):
+        with open(self.template_path, 'r') as f:
+            self.template = json.load(f)

utils/result_formatter.py

@@ -0,0 +1,68 @@
+import json
+from typing import List, Dict, Any
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+class ResultFormatter:
+    def __init__(self):
+        self.results = []
+
+    def add_result(self, prediction: float, actual: float, metadata: Dict[str, Any]):
+        self.results.append({
+            'prediction': prediction,
+            'actual': actual,
+            'metadata': metadata
+        })
+
+    def to_json(self, output_path: str):
+        with open(output_path, 'w') as f:
+            json.dump(self.results, f, indent=2)
+
+    def to_csv(self, output_path: str):
+        df = pd.DataFrame(self.results)
+        df.to_csv(output_path, index=False)
+
+    def generate_summary(self) -> Dict[str, Any]:
+        df = pd.DataFrame(self.results)
+        return {
+            'total_samples': len(df),
+            'mean_prediction': df['prediction'].mean(),
+            'mean_actual': df['actual'].mean(),
+            'median_prediction': df['prediction'].median(),
+            'median_actual': df['actual'].median(),
+            'min_prediction': df['prediction'].min(),
+            'max_prediction': df['prediction'].max(),
+            'min_actual': df['actual'].min(),
+            'max_actual': df['actual'].max()
+        }
+
+    def plot_prediction_vs_actual(self, output_path: str):
+        df = pd.DataFrame(self.results)
+        plt.figure(figsize=(10, 6))
+        sns.scatterplot(x='actual', y='prediction', data=df)
+        plt.title('Predicted vs Actual CVSS Scores')
+        plt.xlabel('Actual CVSS Score')
+        plt.ylabel('Predicted CVSS Score')
+        plt.savefig(output_path)
+        plt.close()
+
+    def plot_error_distribution(self, output_path: str):
+        df = pd.DataFrame(self.results)
+        df['error'] = df['prediction'] - df['actual']
+        plt.figure(figsize=(10, 6))
+        sns.histplot(df['error'], kde=True)
+        plt.title('Distribution of Prediction Errors')
+        plt.xlabel('Prediction Error')
+        plt.ylabel('Frequency')
+        plt.savefig(output_path)
+        plt.close()
+
+    def clear_results(self):
+        self.results.clear()
+
+    def get_results(self) -> List[Dict[str, Any]]:
+        return self.results.copy()
+
+    def filter_results(self, condition: callable) -> List[Dict[str, Any]]:
+        return list(filter(condition, self.results))

vllm_accelerator/__init__.py

@@ -0,0 +1,4 @@
+from .parallel_processor import ParallelProcessor
+from .memory_optimizer import MemoryOptimizer
+
+__all__ = ['ParallelProcessor', 'MemoryOptimizer']

vllm_accelerator/memory_optimizer.py

@@ -0,0 +1,43 @@
+import torch
+import gc
+from typing import Optional
+
+class MemoryOptimizer:
+    def __init__(self, target_memory_usage: float = 0.8):
+        self.target_memory_usage = target_memory_usage
+
+    def optimize(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def get_memory_usage(self) -> Dict[str, float]:
+        current_device = torch.cuda.current_device()
+        allocated = torch.cuda.memory_allocated(current_device) / (1024 ** 3)
+        reserved = torch.cuda.memory_reserved(current_device) / (1024 ** 3)
+        return {
+            'allocated_gb': allocated,
+            'reserved_gb': reserved,
+            'utilization': allocated / reserved if reserved > 0 else 0
+        }
+
+    def check_memory_pressure(self) -> bool:
+        usage = self.get_memory_usage()
+        return usage['utilization'] > self.target_memory_usage
+
+    def try_reduce_batch_size(self, current_batch_size: int) -> Optional[int]:
+        if self.check_memory_pressure():
+            new_batch_size = max(1, current_batch_size // 2)
+            return new_batch_size if new_batch_size < current_batch_size else None
+        return None
+
+    def clear_cuda_cache(self):
+        torch.cuda.empty_cache()
+
+    def reset_peak_memory_stats(self):
+        torch.cuda.reset_peak_memory_stats()
+
+    def get_peak_memory_usage(self) -> float:
+        return torch.cuda.max_memory_allocated() / (1024 ** 3)
+
+    def __del__(self):
+        self.clear_cuda_cache()

vllm_accelerator/parallel_processor.py

@@ -0,0 +1,39 @@
+import torch
+from vllm import LLM, SamplingParams
+from typing import List, Dict, Any
+
+class ParallelProcessor:
+    def __init__(self, model_path: str, tensor_parallel_size: int = 1, gpu_memory_utilization: float = 0.9):
+        self.llm = LLM(model=model_path, tensor_parallel_size=tensor_parallel_size, gpu_memory_utilization=gpu_memory_utilization)
+        self.sampling_params = SamplingParams(temperature=0.7, top_p=0.95, max_tokens=256)
+
+    def process_batch(self, prompts: List[str]) -> List[Dict[str, Any]]:
+        outputs = self.llm.generate(prompts, self.sampling_params)
+        results = []
+        for output in outputs:
+            results.append({
+                'prompt': output.prompt,
+                'generated_text': output.outputs[0].text,
+                'token_count': len(output.outputs[0].token_ids),
+                'finish_reason': output.outputs[0].finish_reason
+            })
+        return results
+
+    def update_sampling_params(self, **kwargs):
+        self.sampling_params = SamplingParams(**kwargs)
+
+    def get_model_info(self) -> Dict[str, Any]:
+        return {
+            'model_name': self.llm.model_config.model_name,
+            'tensor_parallel_size': self.llm.worker.model_runner.tensor_parallel_size,
+            'gpu_memory_utilization': self.llm.worker.model_runner.gpu_memory_utilization
+        }
+
+    def warm_up(self):
+        dummy_prompt = "This is a warm-up prompt."
+        self.process_batch([dummy_prompt])
+
+    def __del__(self):
+        if hasattr(self, 'llm'):
+            del self.llm
+        torch.cuda.empty_cache()