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
|
||
)
|