tdbtxdy
/
competition-vd
forked from Eshe/competition-vd
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Compare commits

merge into: tdbtxdy:master
Branches Tags
tdbtxdy:master
Eshe:master
...
pull from: Eshe:master
Branches Tags
Eshe:master
tdbtxdy:master

No commits in common. "master" and "master" have entirely different histories.
master ... master

31 changed files with 14 additions and 184062 deletions
Show Stats Expand all files Collapse all files

14
README.md Normal file
View File

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

19
README_V2.md
View File

@ -1,19 +0,0 @@
# 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` 目录中。

BIN
Report.pdf
View File

Binary file not shown.

34
config.json
View File

@ -1,34 +0,0 @@
{
"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
}
}

4
cvss_analysis/__init__.py
View File

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

78
cvss_analysis/score_predictor.py
View File

@ -1,78 +0,0 @@
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)
}

86
cvss_analysis/vulnerability_analyzer.py
View File

@ -1,86 +0,0 @@
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

91483
data/SIR_dataset_processed.json
View File

File diff suppressed because one or more lines are too long

9232
data/SIR_test_set.json
View File

File diff suppressed because one or more lines are too long

73114
data/SIR_train_set.json
View File

File diff suppressed because one or more lines are too long

9141
data/SIR_validation_set.json
View File

File diff suppressed because one or more lines are too long

46
data/label_word_ids.json
View File

@ -1,46 +0,0 @@
{
"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
}
}

37
data/label_word_ids_CVSS2.json
View File

@ -1,37 +0,0 @@
{
"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
}
}

4
data_handlers/__init__.py
View File

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

72
data_handlers/dataset_preprocessor.py
View File

@ -1,72 +0,0 @@
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)
}

53
data_handlers/json_parser.py
View File

@ -1,53 +0,0 @@
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
}

69
dockerfile
View File

@ -1,69 +0,0 @@
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"]

4
evaluation/__init__.py
View File

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

50
evaluation/mse_calculator.py
View File

@ -1,50 +0,0 @@
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)

63
evaluation/rmse_calculator.py
View File

@ -1,63 +0,0 @@
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
}

4
llama_core/__init__.py
View File

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

82
llama_core/inference_engine.py
View File

@ -1,82 +0,0 @@
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)

54
llama_core/model_loader.py
View File

@ -1,54 +0,0 @@
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

95
main.py
View File

@ -1,95 +0,0 @@
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)

49
requirements.txt
View File

@ -1,49 +0,0 @@
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

4
utils/__init__.py
View File

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

31
utils/prompt_generator.py
View File

@ -1,31 +0,0 @@
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)

68
utils/result_formatter.py
View File

@ -1,68 +0,0 @@
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))

4
vllm_accelerator/__init__.py
View File

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

43
vllm_accelerator/memory_optimizer.py
View File

@ -1,43 +0,0 @@
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()

39
vllm_accelerator/parallel_processor.py
View File

@ -1,39 +0,0 @@
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()