242 lines
8.2 KiB
Python
242 lines
8.2 KiB
Python
|
# ============================================================
|
|||
|
|
# 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
|
|||
|
|
)
|