Fix GAN mode collapse: discriminator dominance, temperature, entropy tuning + code cleanup

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.py[cod]
+*.pth
+.venv/
+*.egg-info/
+dist/
+build/

src/config.py

@@ -53,7 +53,7 @@
 # ==================== GENERATION ====================
 NUM_SAMPLES = 10
 GENERATION_LABEL_TYPE = 0  # 0=error_based, 1=time_based, 2=union_based
-GENERATION_TEMPERATURE = 1.2  # Temperatura aumentata per favorire maggiore diversità nei payload generati
+GENERATION_TEMPERATURE = 0.8  # Lower temperature for more structured, less random generation
 
 # ==================== BPE TOKENIZER ====================
 BPE_VOCAB_SIZE = 5000

src/generate_payloads.py

@@ -3,67 +3,43 @@
 from generator import ConditionalGenerator
 from config import (
     VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM, NUM_CLASSES,
-    GENERATOR_MODEL, TOKENIZER_CONFIG, 
-    START_TOKEN, USE_CUDA, MAX_SEQ_LEN,
+    GENERATOR_MODEL, TOKENIZER_CONFIG,
+    START_TOKEN, MAX_SEQ_LEN,
     NUM_SAMPLES, GENERATION_LABEL_TYPE
 )
 
+
 def generate_payloads(num_samples=None, label_type=None, max_seq_len=None):
-    # Configurazione dei parametri del modello
     num_samples = NUM_SAMPLES if num_samples is None else num_samples
     label_type = GENERATION_LABEL_TYPE if label_type is None else label_type
     max_seq_len = MAX_SEQ_LEN if max_seq_len is None else max_seq_len
 
-    # Percorsi dei file
-    model_path = GENERATOR_MODEL
-    tokenizer_json = TOKENIZER_CONFIG
-
-    # Usa la GPU se disponibile
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Caricamento del tokenizer
-    print("Caricamento del tokenizer...")
-    tokenizer = Tokenizer.from_file(tokenizer_json)
+    print("Loading tokenizer...")
+    tokenizer = Tokenizer.from_file(TOKENIZER_CONFIG)
 
-    # Caricamento del modello
-    print("Caricamento del Generatore...")
+    print("Loading generator...")
     generator = ConditionalGenerator(VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM, NUM_CLASSES)
-
-    # Carichiamo i "pesi" salvati durante l'addestramento
-    # map_location=device serve per evitare errori se hai addestrato su GPU ma esegui su CPU
     generator.load_state_dict(torch.load(GENERATOR_MODEL, map_location=device))
     generator.to(device)
-
-    # Impostiamo il modello in modalità "Valutazione/Inferenza" (disabilita il calcolo dei gradienti)
     generator.eval()
 
-    # Generezione dei payloads
-    print(f"\nGenerazione di {num_samples} payload per la classe {label_type}...\n")
-
-    # Creiamo le etichette per i campioni che vogliamo generare (es. 0 per error_based)
+    print(f"\nGenerating {num_samples} payloads for class {label_type}...\n")
     labels = torch.tensor([label_type] * num_samples).to(device)
 
-    # Usiamo torch.no_grad() perché non dobbiamo aggiornare i pesi della rete ora
     with torch.no_grad():
-        # L'ID 2 corrisponde a [CLS] nel vocabolario BPE
-       generated_tokens = generator.sample(
-            batch_size=num_samples, 
-            start_token_id=START_TOKEN, 
-            labels=labels, 
+        generated_tokens = generator.sample(
+            batch_size=num_samples,
+            start_token_id=START_TOKEN,
+            labels=labels,
             max_seq_len=max_seq_len
-        ) 
-
-    # Decodifica e stampa
-    generated_tokens_list = generated_tokens.cpu().tolist()
-
-    for i, token_sequence in enumerate(generated_tokens_list): 
-        # Il Tokenizer ha una funzione comoda 'decode' che ritraduce gli ID in stringhe.
-        # skip_special_tokens=True rimuove in automatico [PAD], [CLS], ecc.
-        testo_decodificato = tokenizer.decode(token_sequence, skip_special_tokens=True)
+        )
 
+    for i, token_sequence in enumerate(generated_tokens.cpu().tolist()):
+        decoded = tokenizer.decode(token_sequence, skip_special_tokens=True)
         print(f"Payload {i+1} ---")
-        # Rimuoviamo gli spazi extra che il tokenizer potrebbe aver aggiunto durante la decodifica
-        print(testo_decodificato.replace(" ", ""))
+        print(decoded.replace(" ", ""))
         print("-" * 30)
 
 

src/generator.py

@@ -3,113 +3,51 @@
 import torch.nn.functional as F
 from config import GENERATION_TEMPERATURE
 
+
 class ConditionalGenerator(nn.Module):
     def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
-        """
-        Inizializza il generatore condizionato.
-
-        Parametri:
-        - vocab_size: dimensione del vocabolario (dal JSON per adesso è circa 379)
-        - embed_dim: dimensione del vettore in trasformeremo i token (circa 32 o 64)
-        - hidden_dim: La dimensione della memoria nascosta dell'LSTM (es. 64 o 128).
-        - num_classes: Il numero di etichette condizionali (es. 3, per 'ok', 'error', 'fail').
-        """
-        # Parametri di configurazione del modello
         super(ConditionalGenerator, self).__init__()
         self.hidden_dim = hidden_dim
 
-        # Livello di Embedding per i Token SQL
-        # Trasforma gli ID interi del dizionario BPE in vettori densi di numeri.
         self.token_embedding = nn.Embedding(vocab_size, embed_dim)
-
-        # Livello di Embedding per le Etichette (Condizione)
-        # Trasforma le etichette condizionali (es. 'ok', 'error', 'fail') in un vettore della stessa dimensione.
         self.label_embedding = nn.Embedding(num_classes, embed_dim)
-
-        # Livello LSTM 
-        # Prende in input la concatenazione del token e della condizione
-        # La dimensione di ingresso è embed_dim (token) + embed_dim (label) = embed_dim * 2
+        # LSTM input: token embedding concatenated with label embedding
         self.lstm = nn.LSTM(embed_dim * 2, hidden_dim, batch_first=True)
-
-        # Livello lineare (Output)
-        # Mappa l'output della LSTM (hidden_dim) alla dimensione del vocabolario
-        # Serve per calcolare la probabiltà del token successivo da generare.
         self.fc = nn.Linear(hidden_dim, vocab_size)
 
     def forward(self, x, labels, hidden=None):
-        """
-        Passaggio in avanti (forward pass) del modello.
-        Usato durante il pre-training Maximum Likelihood Estimation (MLE).
-
-        Parametri:
-        - x: Tensore dei token di input (batch_size, seq_len) - sequenza di token SQL codificati come ID interi
-        - labels: Tensore delle etichette condizionali (batch_size, 1) - ad esempio, 0 per 'ok', 1 per 'error', 2 per 'fail'
-        - hidden: Stato nascosto iniziale per la LSTM 
-        """
         batch_size, seq_len = x.size()
 
-        # Otteniamo la rappresentazione vettoriale dei token SQL
-        emb_x = self.token_embedding(x)  # (batch_size, seq_len, embed_dim)
+        emb_x = self.token_embedding(x)  # (B, T, embed_dim)
 
-        # Otteniamo la rappresentazione vettoriale delle etichette condizionali
-        emb_label = self.label_embedding(labels)  # (batch_size, 1, embed_dim)
-
-        # Assicuriamoci che labels abbia una dimensione aggiuntiva per l'espansione
+        emb_label = self.label_embedding(labels)  # (B, 1, embed_dim) or (B, embed_dim)
         if emb_label.dim() == 2:
-            emb_label = emb_label.unsqueeze(1)  # (batch_size, 1, embed_dim)
-
-        # Espandiamo l'etichetta per tutta la lunghezza della sequenza
-        # Vogliamo che l'LSTM ricordi la condizione ad ogni singolo token che processa
-        emb_label = emb_label.expand(batch_size, seq_len, -1)  # (batch_size, seq_len, embed_dim)
+            emb_label = emb_label.unsqueeze(1)
+        emb_label = emb_label.expand(batch_size, seq_len, -1)  # (B, T, embed_dim)
 
-        # Concatenazione del token con l'etichetta
-        lstm_input = torch.cat([emb_x,emb_label], dim=2)  # (batch_size, seq_len, embed_dim * 2)
-
-        # Passare i dati nell'LSTM
-        out, hidden = self.lstm(lstm_input, hidden)  # out: (batch_size, seq_len, hidden_dim)
-
-        # Otteniamo i punteggi finali per ogni token del vocabolario
-        logits = self.fc(out)  # (batch_size, seq_len, vocab_size)
+        lstm_input = torch.cat([emb_x, emb_label], dim=2)  # (B, T, embed_dim * 2)
+        out, hidden = self.lstm(lstm_input, hidden)          # (B, T, hidden_dim)
+        logits = self.fc(out)                                 # (B, T, vocab_size)
 
         return logits, hidden
-    
+
     def sample(self, batch_size, start_token_id, labels, max_seq_len):
-        """
-        Metodo per generare (campionare) nuovi payload da zero.
-        Il generatore parte dal token iniziale e genera un token alla volta.
-        """
-        # Creazione del primo input: una colonna piena del token iniziale
-        inputs = torch.LongTensor(batch_size, 1).fill_(start_token_id).to(labels.device)  # (batch_size, 1)
+        """Autoregressively generate a sequence of tokens starting from start_token_id."""
+        inputs = torch.LongTensor(batch_size, 1).fill_(start_token_id).to(labels.device)
 
-        # Se c'è una GPU disponibile, spostiamo il modello e i tensori su di essa
         if next(self.parameters()).is_cuda:
             inputs = inputs.cuda()
             labels = labels.cuda()
 
-        hidden = None # Stato nascosto iniziale
-        samples = [] # Lista per memorizzare i token generati
+        hidden = None
+        samples = []
 
         for _ in range(max_seq_len):
-            # Passiamo il token attuale e l'etichetta attraverso la rete
-            logits, hidden = self.forward(inputs, labels, hidden)  # logits: (batch_size, 1, vocab_size)
-
-            # Estraiamo l'output per l'ultimo token calcolato
-            logits = logits[:, -1, :]  # (batch_size, vocab_size)
-
-            # Applichiamo una temperatura per controllare la casualità della generazione
-            logits = logits / GENERATION_TEMPERATURE
-            # Convertiamo i punteggi in probabilità usando softmax
-            probs = F.softmax(logits, dim=-1)  # (batch_size, vocab_size)
-
-            # Campioniamo il token successivo in base alle probabilità
-            next_token = torch.multinomial(probs, 1)  # (batch_size, 1)   
-
-            # Aggiungiamo il token campionato alla nostra lista
+            logits, hidden = self.forward(inputs, labels, hidden)
+            logits = logits[:, -1, :] / GENERATION_TEMPERATURE  # (B, vocab_size)
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, 1)             # (B, 1)
             samples.append(next_token)
-
-            # Il token campionato diventa l'input per il prossimo passo
             inputs = next_token
 
-        # uniamo tutti i token generati in un unico tensore
-        samples = torch.cat(samples, dim=1)  # (batch_size, max_seq_len)
-        return samples    
+        return torch.cat(samples, dim=1)  # (B, max_seq_len)