El repositorio de TrafficLLM , un marco de adaptación universal de LLM para aprender una representación sólida del tráfico para todos los LLM de código abierto en escenarios del mundo real y mejorar la generalización en diversas tareas de análisis de tráfico.
Nota: este código está basado en ChatGLM2 y Llama2. Muchas gracias a los autores.
[2024.10.28] Hemos actualizado el código de adaptación para usar GLM4 para construir TrafficLLM, que tiene una velocidad de ajuste e inferencia más rápida que ChatGLM2. Vaya a Adapt2GLM4 para obtener más detalles.
TrafficLLM se basa en un sofisticado marco de ajuste que utiliza lenguaje natural y datos de tráfico, que propone las siguientes técnicas para mejorar la utilidad de grandes modelos de lenguaje en el análisis del tráfico de red.
Tokenización de dominio de tráfico. Para superar la brecha de modalidad entre el lenguaje natural y los datos de tráfico heterogéneos, TrafficLLM introduce la tokenización de dominio de tráfico para procesar las diversas entradas de las tareas de generación y detección de tráfico para la adaptación de LLM. Este mecanismo extiende efectivamente el tokenizador nativo de LLM al entrenar el modelo de tokenización especializado en corpus de dominio de tráfico a gran escala.
Tubería de ajuste de doble etapa. TrafficLLM emplea un proceso de ajuste de dos etapas para lograr el sólido aprendizaje de representación de LLM en diferentes tareas de dominio de tráfico. El canal entrena a LLM para comprender instrucciones y aprender patrones de tráfico relacionados con tareas en diferentes etapas, lo que se basa en la comprensión de tareas de TrafficLLM y las capacidades de razonamiento del tráfico para diversas tareas de detección y generación de tráfico.
Adaptación extensible con ajuste fino efectivo de parámetros (EA-PEFT). Para adaptar LLM para la generalización a nuevos entornos de tráfico, TrafficLLM propone una adaptación extensible con ajuste fino efectivo de parámetros (EA-PEFT) para actualizar los parámetros del modelo con baja sobrecarga. La técnica divide las capacidades del modelo en diferentes modelos PEFT, lo que ayuda a minimizar los costos en escenarios dinámicos generados por cambios en los patrones de tráfico.
Lanzamos los conjuntos de datos de capacitación de TrafficLLM, que contienen más de 0,4 millones de datos de tráfico y 9.000 instrucciones humanas para la adaptación de LLM en diferentes tareas de análisis de tráfico.
Instruction Datasets : los conjuntos de datos de instrucción se utilizan para ayudar a LLM a aprender el conocimiento del dominio de las tareas de generación o detección de tráfico y comprender qué tarea debe realizarse en diferentes escenarios.
Traffic Datasets : los conjuntos de datos de tráfico contienen los datos de ajuste de tráfico que extrajimos de los conjuntos de datos de tráfico públicos, lo que ayuda a LLM a aprender el patrón de tráfico en diferentes tareas posteriores.
Para construir el corpus de lenguaje natural como instrucciones humanas en TrafficLLM, recopilamos 9209 instrucciones de tareas específicas supervisadas por expertos y asistentes de IA. Las estadísticas se muestran a continuación:
| Tareas principales | Tareas posteriores | Abrev. | #Muestra |
|---|---|---|---|
| Detección de tráfico | Detección de tráfico de malware | MTD | 1.0K |
| Detección de redes robot | BND | 1.1K | |
| Detección de DoH malicioso | TDM | 0,6K | |
| Detección de ataques web | TACO | 0,6K | |
| Detección de ataques APT | DAA | 0,6K | |
| Detección de VPN cifrada | EVE | 1.2K | |
| Detección de comportamiento de Tor | Por determinar | 0,6K | |
| Clasificación de aplicaciones cifradas | CAO | 0,6K | |
| Huellas dactilares del sitio web | WF | 0,6K | |
| Deriva conceptual | CD | 0,6K | |
| Generación de tráfico | Generación de tráfico de malware | MTG | 0,6K |
| Generación de tráfico de botnets | BTG | 0,1K | |
| Generación de VPN cifrada | EVG | 0.4K | |
| Generación de aplicaciones cifradas | EAG | 0,6K |
Para evaluar el rendimiento de TrafficLLM en varios escenarios de red, extrajimos más de 0,4 millones de datos de ajuste de conjuntos de datos de tráfico disponibles públicamente para medir las capacidades de TrafficLLM para detectar o generar tráfico malicioso y benigno. Las estadísticas se muestran a continuación:
| Conjuntos de datos | Tareas | Abrev. | #Muestra |
|---|---|---|---|
| USTC TFC 2016 | Detección de tráfico de malware | MTD | 50,7K |
| Red de bots ISCX 2014 | Detección de redes robot | BND | 25.0K |
| DoHBrw 2020 | Detección de DoH malicioso | TDM | 47,8K |
| CSIC 2010 | Detección de ataques web | TACO | 34,5K |
| DAPT 2020 | Detección de ataques APT | DAA | 10.0K |
| VPN ISCX 2016 | Detección de VPN cifrada | EVE | 64,8K |
| ISCX Tor 2016 | Detección de comportamiento de Tor | Por determinar | 40.0K |
| CSTNET 2023 | Clasificación de aplicaciones cifradas | CAO | 97,6K |
| CW-100 2018 | Huellas dactilares del sitio web | WF | 7.4K |
| APLICACIÓN-53 2023 | Deriva conceptual | CD | 109,8K |
1. Preparación del entorno 2. Capacitación TrafficLLM 2.1. Preparación del punto de control previamente entrenado 2.2. Conjunto de datos de preprocesamiento 2.3. Entrenamiento del tokenizador de dominio de tráfico (opcional) 2.4. Ajuste de la instrucción del lenguaje neuronal 2.5. Ajuste del tráfico para tareas específicas 2.6. Adaptación extensible con PEFT (EA-PEFT) 3. Evaluación de TrafficLLM 3.1. Preparación de puntos de control y datos 3.2. Evaluación en ejecuciónTabla de contenidos:
1. Preparación del entorno [Volver al inicio]
Clone el repositorio e instale el entorno requerido ejecutando los siguientes comandos.
conda crear -n tráfico python = 3.9 conda activar Trafficllm# Clonar nuestro clon de TrafficLLMgit https://github.com/ZGC-LLM-Safety/TrafficLLM.gitcd TrafficLLM# Instalar las bibliotecas requeridaspip install -r requisitos.txt# Si Trainingpip instala rouge_chinese nltk jieba datasets
TrafficLLM emplea tres técnicas principales: tokenización de dominio de tráfico para procesar instrucciones y datos de tráfico, proceso de ajuste de dos etapas para comprender la semántica del texto y aprender patrones de tráfico en diferentes tareas, y EA-PEFT para actualizar Parámetros del modelo para la adaptación a nuevos escenarios.
TrafficLLM se entrena en base a LLM de código abierto existentes. Siga las instrucciones para preparar los puntos de control.
ChatGLM2 : Prepare el modelo base ChatGLM, que es un LLM de código abierto con requisitos de implementación livianos. Descargue sus pesos aquí. Generalmente utilizamos el modelo v2 con parámetros 6B.
Other LLMs : para adaptar otros LLM para tareas de análisis de tráfico, puede reutilizar los datos de capacitación en el repositorio y modificar sus guiones de capacitación de acuerdo con las instrucciones oficiales. Por ejemplo, se requiere Llama2 para registrar el nuevo conjunto de datos en las configuraciones.
Para extraer datos de entrenamiento adecuados para el aprendizaje LLM a partir de los conjuntos de datos de tráfico sin procesar, diseñamos extractores especializados para preprocesar conjuntos de datos de tráfico para diferentes tareas. El código de preprocesamiento contiene los siguientes parámetros para configurar.
input : la ruta del conjunto de datos de tráfico sin procesar (la ruta del directorio principal que contiene subdirectorios etiquetados. Cada subdirectorio etiquetado contiene los archivos .pcap sin procesar que se procesarán previamente).
dataset_name : el nombre del conjunto de datos de tráfico sin procesar (ayuda a juzgar si el nombre se ha registrado en los códigos de TrafficLLM).
traffic_task : Tareas de detección o tareas de generación.
granularity : granularidad a nivel de paquete o de flujo.
output_path : ruta del conjunto de datos de entrenamiento de salida.
output_name : nombre del conjunto de datos de entrenamiento de salida.
Esta es una instancia para preprocesar conjuntos de datos de tráfico sin procesar para tareas de detección de tráfico a nivel de paquetes.
preproceso de cd python preprocess_dataset.py --input /Su/Raw/Conjunto de datos/Ruta --dataset_name /Su/Raw/Conjunto de datos/Nombre --tráfico_detección de tareas --paquete de granularidad --output_path /Su/Salida/Conjunto de datos/Ruta --output_name /Su /Salida/Conjunto de datos/Nombre
TrafficLLM presenta un tokenizador de dominio de tráfico para manejar el lenguaje neuronal y los datos de tráfico. Si desea entrenar un tokenizador personalizado con su propio conjunto de datos, modifique model_name y data_path en el código.
model_name : la ruta del modelo base que contiene el tokenizador nativo.
data_path : los conjuntos de datos de entrenamiento extraídos del proceso de preprocesamiento.
Siga el comando para usar el código.
tokenización de cd python Traffic_tokenizer.py
Preparar datos: Los datos de ajuste de instrucciones del lenguaje neuronal son nuestros conjuntos de datos de instrucciones recopilados para la comprensión de las tareas de análisis de tráfico.
Comience a ajustar: después de los pasos mencionados anteriormente, puede comenzar la primera etapa de ajuste utilizando Trafficllm_stage1.sh. Hay un ejemplo como el siguiente:
PRE_SEQ_LEN=128
LR=2e-2
NUM_GPUS=1exportar CUDA_VISIBLE_DEVICES=1
torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py
--do_train
--train_file ../datasets/instructions/instructions.json
--validation_file ../datasets/instructions/instructions.json
--preprocesamiento_num_trabajadores 10
--prompt_column instrucción
--respuesta_columna salida
--overwrite_cache
--cache_dir ../cache
--model_name_or_path ../models/chatglm2/chatglm2-6b
--output_dir ../models/chatglm2/peft/instruction
--overwrite_output_dir
--max_source_length 1024
--max_target_length 32
--per_device_train_batch_size 1
--per_device_eval_batch_size 1
--gradient_acumulación_pasos 16
--predict_with_generate
--max_pasos 20000
--logging_steps 10
--save_steps 4000
--tasa_aprendizaje $LR
--pre_seq_len $PRE_SEQ_LEN Preparar datos: Los conjuntos de datos de ajuste de tráfico específicos de la tarea son los conjuntos de datos de entrenamiento extraídos del paso de preprocesamiento para diferentes tareas posteriores.
Comience a ajustar: después de los pasos mencionados anteriormente, puede comenzar la optimización de la segunda etapa utilizando Trafficllm_stage2.sh. Hay un ejemplo como el siguiente:
PRE_SEQ_LEN=128
LR=2e-2
NUM_GPUS=1exportar CUDA_VISIBLE_DEVICES=1
torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py
--do_train
--train_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_train.json
--validation_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_train.json
--preprocesamiento_num_trabajadores 10
--prompt_column instrucción
--respuesta_columna salida
--overwrite_cache
--cache_dir ../cache
--model_name_or_path ../models/chatglm2/chatglm2-6b
--output_dir ../models/chatglm2/peft/ustc-tfc-2016-detection-packet
--overwrite_output_dir
--max_source_length 1024
--max_target_length 32
--per_device_train_batch_size 1
--per_device_eval_batch_size 1
--gradient_acumulación_pasos 16
--predict_with_generate
--max_pasos 20000
--logging_steps 10
--save_steps 4000
--tasa_aprendizaje $LR
--pre_seq_len $PRE_SEQ_LEN TrafficLLM emplea EA-PEFT para organizar los modelos de ajuste fino efectivo de parámetros (PEFT) con una adaptación extensible, que puede ayudar a TrafficLLM a adaptarse fácilmente a nuevos entornos. El adaptador TrafficLLM permite operaciones flexibles para actualizar modelos antiguos o registrar nuevas tareas.
model_name : la ruta del modelo base.
tuning_data : el nuevo conjunto de datos del entorno.
adaptation_task : Actualizar o registrar (actualizar modelos antiguos o registrar nuevas tareas).
task_name : el nombre de la tarea posterior que se actualizará o insertará.
Hay un ejemplo para actualizar TrafficLLM con tareas de detección de tráfico de malware (MTD).
cd EA-PEFT python ea-peft.py --model_name /Su/Base/Modelo/Ruta --tuning_data /Su/Nuevo/Conjunto de datos/Ruta --adaptation_task update --task_name MTD
Puntos de control: puede intentar evaluar TrafficLLM utilizando su propio modelo o nuestros puntos de control publicados.
Datos: durante el paso de preprocesamiento, dividimos conjuntos de datos de prueba y creamos archivos de etiquetas para diferentes conjuntos de datos para su evaluación. Consulte los códigos de preprocesamiento.
Para medir la efectividad de TrafficLLM para diferentes tareas posteriores, ejecute los códigos de evaluación.
model_name : la ruta del modelo base.
traffic_task : Tareas de detección o tareas de generación.
test_file : los conjuntos de datos de prueba extraídos durante los pasos de preprocesamiento.
label_file : el archivo de etiqueta extraído durante los pasos de preprocesamiento.
ptuning_path : la ruta del modelo PEFT para la evaluación de tareas específicas.
Hay un ejemplo para ejecutar la evaluación de tareas MTD.
python evaluacion.py --model_name /Su/Base/Modelo/Ruta --traffic_task deteccion --test_file datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_test.json --label_file datasets/ustc-tfc-2016/ustc- tfc-2016_label.json --ptuning_path models/chatglm2/peft/ustc-tfc-2016-detection-packet/checkpoints-20000/
Puede implementar TrafficLLM en su dispositivo local. Primero, configure la ruta del modelo en config.json para registrar los modelos PEFT recopilados en los pasos de capacitación. Hay un ejemplo para registrar 6 tareas en TrafficLLM:
{ "model_path": "models/chatglm2/chatglm2-6b/", "peft_path": "models/chatglm2/peft/", "peft_set": { "NLP": " instrucción/checkpoint-8000/", "MTD": "ustc-tfc-2016-detection-packet/checkpoint-10000/", "BND": "iscx-botnet-2014-detection-packet/checkpoint-5000/", "WAD": "csic-2010-detection-packet/checkpoint-6000/", "AAD": "dapt-2020-detection-packet/checkpoint-20000/", "EVD": "iscx-vpn-2016-detection -packet/checkpoint-4000/", "TBD": "iscx-tor-2016-detection-packet/checkpoint-10000/"
}, "tareas": { "Detección de tráfico de malware": "MTD", "Detección de botnets": "BND", "Detección de ataques web": "WAD", "Detección de ataques APT": "AAD", "Detección de VPN cifrada ": "EVD", "Detección de comportamiento Tor": "TBD"
}
} Luego deberías agregar el aviso previo en la función prepromt de inference.py y Trafficllm_server.py. El mensaje previo es el texto de prefijo utilizado en los datos de entrenamiento durante el ajuste del tráfico específico de la tarea.
Para chatear con TrafficLLM en modo terminal, puede ejecutar el siguiente comando:
python inference.py --config=config.json --prompt="Su texto de instrucciones +: + Datos de tráfico"
web Inicie la demostración del sitio web de TrafficLLM usando el siguiente comando:
streamlit run Trafficllm_server.py
Esta demostración ejecuta un servidor web de TrafficLLM. Acceda http://Your-Server-IP:Port para chatear en el chatbox.
Muchas gracias al trabajo relacionado ChatGLM2 y Llama2 que sirve como base para nuestro marco y códigos. El diseño del edificio TrafficLLM está inspirado en ET-BERT y GraphGPT. Gracias por sus maravillosos trabajos.
