初始化项目，由ModelHub XC社区提供模型
Model: hypnonyx/traffico Source: Original Platform
This commit is contained in:
ModelHub XC
241
TrafficoDataset/train_gemma3_traffico.py
Normal file
241
TrafficoDataset/train_gemma3_traffico.py
Normal file
@@ -0,0 +1,241 @@
|
||||
# ============================================================
|
||||
# Gemma-3-270M – Analisi Traffico di Rete TCP/IP
|
||||
# Unsloth + LoRA + Dataset JSONL (CIC-IDS2017 + UNSW-NB15)
|
||||
# ============================================================
|
||||
# Struttura basata sul tuo script, adattata per il dominio
|
||||
# di analisi del traffico di rete con mappatura MITRE ATT&CK.
|
||||
#
|
||||
# PREREQUISITI:
|
||||
# Google Colab con runtime GPU (T4 basta)
|
||||
# !pip install --no-deps unsloth
|
||||
# !pip install transformers datasets trl peft accelerate sentencepiece
|
||||
#
|
||||
# FILE NECESSARI (nella stessa cartella dello script):
|
||||
# dataset.jsonl ← generato dalla script apposita
|
||||
# ============================================================
|
||||
|
||||
# ---------- INSTALL (Colab) ----------
|
||||
# !pip install --no-deps unsloth
|
||||
# !pip install transformers datasets trl peft accelerate sentencepiece
|
||||
|
||||
# ---------- IMPORT ----------
|
||||
from unsloth import FastModel
|
||||
from unsloth.chat_templates import get_chat_template, train_on_responses_only
|
||||
import torch
|
||||
from datasets import load_dataset
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
|
||||
# ---------- CONFIG ----------
|
||||
MODEL_NAME = "unsloth/gemma-3-270m-it"
|
||||
DATASET_PATH = "dataset_traffico.jsonl" # <== il JSONL che abbiamo generato
|
||||
OUTPUT_DIR = "outputs"
|
||||
MAX_SEQ_LENGTH = 512 # 512 basta per questi prompt, risparmia memoria
|
||||
|
||||
# ---------- LOAD MODEL ----------
|
||||
model, tokenizer = FastModel.from_pretrained(
|
||||
model_name = MODEL_NAME,
|
||||
max_seq_length = MAX_SEQ_LENGTH,
|
||||
load_in_4bit = False,
|
||||
load_in_8bit = False,
|
||||
full_finetuning = False,
|
||||
)
|
||||
|
||||
# ---------- LoRA ----------
|
||||
# Configurazione più aggressiva sul rank (r=64) per un dominio specifico come questo.
|
||||
# Target modules: tutti i proiettori del transformer.
|
||||
model = FastModel.get_peft_model(
|
||||
model,
|
||||
r = 64,
|
||||
target_modules = [
|
||||
"q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",
|
||||
],
|
||||
lora_alpha = 64,
|
||||
lora_dropout = 0,
|
||||
bias = "none",
|
||||
use_gradient_checkpointing = "unsloth",
|
||||
random_state = 3407,
|
||||
)
|
||||
|
||||
# ---------- CHAT TEMPLATE (Gemma-3) ----------
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template = "gemma3",
|
||||
)
|
||||
|
||||
# ---------- LOAD DATASET ----------
|
||||
dataset = load_dataset(
|
||||
"json",
|
||||
data_files = DATASET_PATH,
|
||||
split = "train",
|
||||
)
|
||||
|
||||
print(f"Dataset caricato: {len(dataset)} righe")
|
||||
print(f"Campi presenti: {dataset.column_names}")
|
||||
print(f"\nEsempio riga 0:")
|
||||
print(dataset[0])
|
||||
|
||||
# ---------- CONVERT TO CHATML ----------
|
||||
# Il JSONL ha campi: instruction, input, output
|
||||
# Li convertiamo nel formato conversations [system, user, assistant]
|
||||
# che Gemma-3 si aspetta.
|
||||
def convert_to_chatml(example):
|
||||
system_prompt = example["instruction"]
|
||||
|
||||
# Se c'è un campo 'context' lo aggiungiamo al system prompt
|
||||
if "context" in example and example["context"]:
|
||||
system_prompt += f"\nContesto: {example['context']}."
|
||||
|
||||
return {
|
||||
"conversations": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": example["input"]},
|
||||
{"role": "assistant", "content": example["output"]},
|
||||
]
|
||||
}
|
||||
|
||||
dataset = dataset.map(convert_to_chatml)
|
||||
|
||||
# ---------- APPLY GEMMA-3 TEMPLATE ----------
|
||||
# Applica il template di chat di Gemma-3 a ogni esempio.
|
||||
# Questo produce la stringa finale che il modello vedrà durante il training.
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["conversations"]
|
||||
texts = [
|
||||
tokenizer.apply_chat_template(
|
||||
convo,
|
||||
tokenize = False,
|
||||
add_generation_prompt = False,
|
||||
).removeprefix("<bos>")
|
||||
for convo in convos
|
||||
]
|
||||
return {"text": texts}
|
||||
|
||||
dataset = dataset.map(formatting_prompts_func, batched=True)
|
||||
|
||||
# Verifica come appare un prompt formattato
|
||||
print("\n" + "=" * 60)
|
||||
print(" PROMPT FORMATTATO (esempio)")
|
||||
print("=" * 60)
|
||||
print(dataset[0]["text"])
|
||||
print("=" * 60)
|
||||
|
||||
# ---------- TRAINER ----------
|
||||
trainer = SFTTrainer(
|
||||
model = model,
|
||||
tokenizer = tokenizer,
|
||||
train_dataset = dataset,
|
||||
eval_dataset = None,
|
||||
args = SFTConfig(
|
||||
dataset_text_field = "text",
|
||||
per_device_train_batch_size = 4,
|
||||
gradient_accumulation_steps = 4, # batch effettivo = 4 * 4 = 16
|
||||
warmup_steps = 10,
|
||||
max_steps = 500, # ~500 step su 10k righe con batch 16
|
||||
learning_rate = 2e-5,
|
||||
logging_steps = 25,
|
||||
optim = "adamw_8bit",
|
||||
weight_decay = 0.001,
|
||||
lr_scheduler_type = "linear",
|
||||
seed = 3407,
|
||||
output_dir = OUTPUT_DIR,
|
||||
report_to = "none",
|
||||
),
|
||||
)
|
||||
|
||||
# ---------- TRAIN ONLY ON ASSISTANT ----------
|
||||
# Fondamentale: il modello calcola il loss SOLO sulla risposta dell'assistant,
|
||||
# non sul prompt. Così non "impara" a ripetere la domanda.
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part = "<start_of_turn>user\n",
|
||||
response_part = "<start_of_turn>model\n",
|
||||
)
|
||||
|
||||
# ---------- TRAIN ----------
|
||||
trainer.train()
|
||||
|
||||
# ---------- SAVE LoRA ----------
|
||||
model.save_pretrained("gemma3-traffico-rete-lora")
|
||||
tokenizer.save_pretrained("gemma3-traffico-rete-lora")
|
||||
print("\n✓ Modello LoRA salvato in: gemma3-traffico-rete-lora/")
|
||||
model.save_pretrained_merged(
|
||||
"gemma3-traffico-rete-lora", # cartella output
|
||||
tokenizer,
|
||||
save_method="merged_16bit" # Float16 per GGUF
|
||||
)
|
||||
model.save_pretrained_gguf(
|
||||
"gemma3-traffico-rete-lora",
|
||||
tokenizer,
|
||||
quantization_method = "BF16", # For now only Q8_0, BF16, F16 supported
|
||||
)
|
||||
# ---------- INFERENCE: TEST ----------
|
||||
# Dopo il training, prova il modello con alcuni flussi di esempio.
|
||||
from transformers import TextStreamer
|
||||
|
||||
test_cases = [
|
||||
# Caso 1: profilo tipico DoS (masse enormi di byte src, pochissimi dst, durata minima)
|
||||
"Protocollo: tcp | Porta dst: 80 | Byte src: 480000 | Byte dst: 40 | Pacchetti: 5200 | Durata: 0.015s",
|
||||
# Caso 2: traffico normale HTTPS
|
||||
"Protocollo: tcp | Porta dst: 443 | Byte src: 1500 | Byte dst: 6200 | Pacchetti: 9 | Durata: 3.200s",
|
||||
# Caso 3: profilo PortScan (tanti dst diversi, pochi byte, durata quasi zero)
|
||||
"Protocollo: tcp | Porta dst: 22 | Byte src: 60 | Byte dst: 0 | Pacchetti: 1 | Durata: 0.002s",
|
||||
# Caso 4: profilo Brute Force su SSH
|
||||
"Protocollo: tcp | Porta dst: 22 | Byte src: 3200 | Byte dst: 8500 | Pacchetti: 45 | Durata: 1.800s",
|
||||
# Caso 5: profilo Infiltration / esfiltrazioni dati
|
||||
"Protocollo: tcp | Porta dst: 443 | Byte src: 8000 | Byte dst: 120000 | Pacchetti: 200 | Durata: 25.500s",
|
||||
]
|
||||
|
||||
streamer = TextStreamer(tokenizer, skip_prompt=True)
|
||||
|
||||
for i, test_input in enumerate(test_cases, 1):
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"Analizza il seguente flusso di traffico di rete TCP/IP. "
|
||||
"Classifica se è traffico normale o un attacco. "
|
||||
"Se è un attacco, indica la categoria e la tecnica MITRE ATT&CK corrispondente."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": test_input},
|
||||
]
|
||||
|
||||
text = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
).removeprefix("<bos>")
|
||||
|
||||
print(f"\n{'─' * 60}")
|
||||
print(f" TEST {i}: {test_input[:80]}...")
|
||||
print(f"{'─' * 60}")
|
||||
print(" Risposta: ", end="")
|
||||
|
||||
_ = model.generate(
|
||||
**tokenizer(text, return_tensors="pt").to("cuda"),
|
||||
max_new_tokens = 128,
|
||||
temperature = 0.3, # bassa temperatura = risposte più deterministe
|
||||
top_p = 0.9,
|
||||
top_k = 40,
|
||||
streamer = streamer,
|
||||
)
|
||||
|
||||
# ---------- SAVE MERGED (opzionale) ----------
|
||||
# Unisce i pesi LoRA al modello base e salva come modello completo.
|
||||
# Utile per deployare senza dipendenza da PEFT.
|
||||
#
|
||||
model.save_pretrained_merged(
|
||||
"gemma3-traffico-rete-merged",
|
||||
tokenizer,
|
||||
save_method = "merged_16bit",
|
||||
)
|
||||
#
|
||||
# ---------- SAVE GGUF (opzionale) ----------
|
||||
# Formato GGUF per inferenza locale con llama.cpp / Ollama.
|
||||
#
|
||||
model.save_pretrained_gguf(
|
||||
"gemma3-traffico-rete-gguf",
|
||||
tokenizer,
|
||||
quantization_method = "Q8_0", # Q8_0 = buon equilibrio qualità/dimensione
|
||||
)
|
||||
Reference in New Issue
Block a user