Dockerfile

@@ -0,0 +1,21 @@
+FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04
+
+WORKDIR /app
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && apt-get install -y python3 python3-pip git wget
+
+COPY requirements.txt /app/
+RUN pip3 install --no-cache-dir -r requirements.txt
+RUN pip3 install vllm
+
+COPY . /app/
+
+ENV PYTHONUNBUFFERED=1
+ENV CUDA_VISIBLE_DEVICES=all
+
+RUN mkdir -p /root/.cache/huggingface/hub
+RUN git lfs install
+RUN python3 -c "from transformers import AutoTokenizer, AutoModelForCausalLM; model_name='deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct'; AutoTokenizer.from_pretrained(model_name); AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)"
+
+CMD ["python3", "scripts/run_inference.py"]

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,104 @@
+# CWE识别项目
+
+本项目使用大型语言模型来识别软件漏洞的 Common Weakness Enumeration (CWE) 类型。
+
+## 文件结构
+
+```
+.
+├── config
+│   └── model_config.yaml
+├── data
+│   └── test.jsonl
+├── results
+│   ├── confusion_matrix.png
+│   ├── metrics_comparison.png
+│   └── output.json
+├── scripts
+│   └── run_inference.py
+├── src
+│   ├── analysis
+│   │   ├── __init__.py
+│   │   └── metrics.py
+│   ├── data
+│   │   ├── __init__.py
+│   │   └── data_loader.py
+│   ├── model
+│   │   ├── __init__.py
+│   │   └── llm_inference.py
+│   └── utils
+│       ├── __init__.py
+│       └── preprocessing.py
+├── Dockerfile
+├── README.md
+└── requirements.txt
+```
+
+## 环境配置
+
+1. **系统要求**：
+
+   确保您的系统已安装 CUDA 12.1 和 cuDNN 8。
+
+2. **创建并激活虚拟环境**：
+
+   ```
+   python3 -m venv venv
+   source venv/bin/activate
+   ```
+
+3. **安装依赖**：
+
+   ```
+   pip install -r requirements.txt
+   pip install vllm
+   ```
+
+4. **安装 Git LFS**：
+
+   ```
+   git lfs install
+   ```
+
+5. **下载预训练模型**：
+
+   ```
+   python3 -c "from transformers import AutoTokenizer, AutoModelForCausalLM; model_name='deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct'; AutoTokenizer.from_pretrained(model_name); AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)"
+   ```
+
+## 运行代码
+
+1. 确保测试数据位于 `data/test.jsonl`。
+
+2. **运行推理脚本**：
+
+   ```
+   python scripts/run_inference.py
+   ```
+
+   脚本执行完成后，您将在 `results` 目录下看到以下输出：
+
+   - **confusion_matrix.png**：混淆矩阵可视化
+   - **metrics_comparison.png**：模型性能指标对比图
+   - **output.json**：包含详细预测结果的 JSON 文件
+
+   控制台将打印评估指标的摘要。
+
+## 使用 Docker
+
+如果您更喜欢使用 Docker，可以按照以下步骤操作：
+
+1. **构建 Docker 镜像**：
+
+   ```
+   docker build -t cwe-identification .
+   ```
+
+2. **运行 Docker 容器**：
+
+   ```
+   docker run --gpus all -v $(pwd)/data:/app/data -v $(pwd)/results:/app/results cwe-identification
+   ```
+
+   这将挂载本地的 `data` 和 `results` 目录到容器中，以便您可以轻松访问输入数据和输出结果。
+

config/model_config.yaml

@@ -0,0 +1,6 @@
+model_name: "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"
+max_model_len: 8192
+tp_size: 1
+temperature: 0.3
+max_tokens: 256
+batch_size: 32

requirements.txt

@@ -0,0 +1,9 @@
+transformers
+vllm
+pyyaml
+numpy
+pandas
+matplotlib
+seaborn
+scikit-learn
+torch

scripts/run_inference.py

@@ -0,0 +1,44 @@
+import sys
+import os
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from src.model.llm_inference import initialize_model, generate_predictions
+from src.data.data_loader import load_dataset, save_jsonl
+from src.utils.preprocessing import prepare_messages, extract_cwe_id
+from src.analysis.metrics import calculate_metrics, plot_confusion_matrix, plot_metrics
+
+def main():
+    tokenizer, llm, sampling_params = initialize_model()
+
+    test_data = load_dataset('data/test.jsonl')
+
+    messages_list = [prepare_messages(item) for item in test_data]
+
+    predictions = generate_predictions(tokenizer, llm, sampling_params, messages_list)
+
+    predicted_cwes = [extract_cwe_id(pred) for pred in predictions]
+    true_cwes = [item['cwe_id'] for item in test_data]
+
+    metrics = calculate_metrics(true_cwes, predicted_cwes)
+    print("评估指标:", metrics)
+
+    unique_cwes = list(set(true_cwes + predicted_cwes))
+    plot_confusion_matrix(true_cwes, predicted_cwes, unique_cwes)
+
+    plot_metrics(metrics)
+
+    results = [
+        {
+            'function': item['function'],
+            'cve_description': item['cve_description'],
+            'true_cwe': item['cwe_id'],
+            'predicted_cwe': pred_cwe,
+            'model_output': pred
+        }
+        for item, pred_cwe, pred in zip(test_data, predicted_cwes, predictions)
+    ]
+    save_jsonl(results, 'results/output.json')
+
+
+if __name__ == "__main__":
+    main()

src/analysis/metrics.py

@@ -0,0 +1,36 @@
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, matthews_corrcoef
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+def calculate_metrics(y_true, y_pred):
+    return {
+        '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)
+    }
+
+def plot_confusion_matrix(y_true, y_pred, labels):
+    from sklearn.metrics import confusion_matrix
+    cm = confusion_matrix(y_true, y_pred, labels=labels)
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=labels, yticklabels=labels)
+    plt.title('Confusion Matrix')
+    plt.xlabel('Predicted')
+    plt.ylabel('True')
+    plt.tight_layout()
+    plt.savefig('results/confusion_matrix.png')
+    plt.close()
+
+def plot_metrics(metrics):
+    plt.figure(figsize=(10, 6))
+    plt.bar(metrics.keys(), metrics.values())
+    plt.title('Model Performance Metrics')
+    plt.ylabel('Score')
+    plt.ylim(0, 1)
+    for i, v in enumerate(metrics.values()):
+        plt.text(i, v + 0.01, f'{v:.4f}', ha='center')
+    plt.tight_layout()
+    plt.savefig('results/metrics_comparison.png')
+    plt.close()

src/data/data_loader.py

@@ -0,0 +1,22 @@
+import json
+
+def load_jsonl(file_path):
+    with open(file_path, 'r') as file:
+        return [json.loads(line) for line in file]
+
+def save_jsonl(data, file_path):
+    with open(file_path, 'w') as file:
+        for item in data:
+            json.dump(item, file)
+            file.write('\n')
+
+def load_dataset(file_path):
+    data = load_jsonl(file_path)
+    return [
+        {
+            'function': item['function'],
+            'cve_description': item['cve_description'],
+            'cwe_id': item['cwe_id'][0] if item['cwe_id'] else None
+        }
+        for item in data
+    ]

src/model/llm_inference.py

@@ -0,0 +1,22 @@
+from transformers import AutoTokenizer
+from vllm import LLM, SamplingParams
+import yaml
+
+def load_model_config():
+    with open('config/model_config.yaml', 'r') as file:
+        return yaml.safe_load(file)
+
+def initialize_model():
+    config = load_model_config()
+    model_name = config['model_name']
+    max_model_len = config['max_model_len']
+    tp_size = config['tp_size']
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    llm = LLM(model=model_name, tensor_parallel_size=tp_size, max_model_len=max_model_len, trust_remote_code=True, enforce_eager=True)
+    sampling_params = SamplingParams(temperature=config['temperature'], max_tokens=config['max_tokens'], stop_token_ids=[tokenizer.eos_token_id])
+    return tokenizer, llm, sampling_params
+
+def generate_predictions(tokenizer, llm, sampling_params, messages_list):
+    prompt_token_ids = [tokenizer.apply_chat_template(messages, add_generation_prompt=True) for messages in messages_list]
+    outputs = llm.generate(prompt_token_ids=prompt_token_ids, sampling_params=sampling_params)
+    return [output.outputs[0].text for output in outputs]

src/utils/preprocessing.py

@@ -0,0 +1,10 @@
+def prepare_messages(item):
+    return [
+        {"role": "system", "content": "You are an AI assistant specialized in identifying software vulnerabilities and determining their CWE (Common Weakness Enumeration) types. Analyze the given function and CVE description, then provide the most likely CWE ID."},
+        {"role": "user", "content": f"Function:\n{item['function']}\n\nCVE Description:\n{item['cve_description']}\n\nBased on this information, what is the most likely CWE ID for this vulnerability?"}
+    ]
+
+def extract_cwe_id(response):
+    import re
+    match = re.search(r'CWE-(\d+)', response)
+    return match.group(1) if match else None