Create Large_Action_Model_Transformer.py

Large_Action_Model_Transformer.py

@@ -0,0 +1,329 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+
+class PositionalEncoding(nn.Module):
+    """
+    Positional Encoding for Transformer models
+    """
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(max_len, 1, d_model)
+        pe[:, 0, 0::2] = torch.sin(position * div_term)
+        pe[:, 0, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        """
+        Args:
+            x: Tensor, shape [seq_len, batch_size, embedding_dim]
+        """
+        x = x + self.pe[:x.size(0)]
+        return self.dropout(x)
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-head attention mechanism
+    """
+    def __init__(self, d_model, num_heads, dropout=0.1):
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+        
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_k = d_model // num_heads
+        
+        self.w_q = nn.Linear(d_model, d_model)
+        self.w_k = nn.Linear(d_model, d_model)
+        self.w_v = nn.Linear(d_model, d_model)
+        self.w_o = nn.Linear(d_model, d_model)
+        
+        self.dropout = nn.Dropout(dropout)
+        self.scale = torch.sqrt(torch.FloatTensor([self.d_k])).to(device)
+        
+    def forward(self, q, k, v, mask=None):
+        batch_size = q.size(0)
+        
+        # Linear projections
+        q = self.w_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        k = self.w_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        v = self.w_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        
+        # Scaled dot-product attention
+        attn = torch.matmul(q, k.transpose(-2, -1)) / self.scale
+        
+        if mask is not None:
+            attn = attn.masked_fill(mask == 0, -1e10)
+        
+        attn = F.softmax(attn, dim=-1)
+        attn = self.dropout(attn)
+        
+        output = torch.matmul(attn, v)
+        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
+        output = self.w_o(output)
+        
+        return output
+
+class PositionwiseFeedforward(nn.Module):
+    """
+    Position-wise feedforward network
+    """
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super().__init__()
+        self.fc1 = nn.Linear(d_model, d_ff)
+        self.fc2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return x
+
+class EncoderLayer(nn.Module):
+    """
+    Single encoder layer
+    """
+    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
+        super().__init__()
+        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
+        self.ffn = PositionwiseFeedforward(d_model, d_ff, dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        
+    def forward(self, x, mask=None):
+        # Self attention
+        attn_output = self.self_attn(x, x, x, mask)
+        x = x + self.dropout1(attn_output)
+        x = self.norm1(x)
+        
+        # Feedforward
+        ff_output = self.ffn(x)
+        x = x + self.dropout2(ff_output)
+        x = self.norm2(x)
+        
+        return x
+
+class DecoderLayer(nn.Module):
+    """
+    Single decoder layer
+    """
+    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
+        super().__init__()
+        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
+        self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout)
+        self.ffn = PositionwiseFeedforward(d_model, d_ff, dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+        
+    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
+        # Self attention
+        attn_output = self.self_attn(x, x, x, tgt_mask)
+        x = x + self.dropout1(attn_output)
+        x = self.norm1(x)
+        
+        # Cross attention
+        attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
+        x = x + self.dropout2(attn_output)
+        x = self.norm2(x)
+        
+        # Feedforward
+        ff_output = self.ffn(x)
+        x = x + self.dropout3(ff_output)
+        x = self.norm3(x)
+        
+        return x
+
+class Transformer(nn.Module):
+    """
+    Complete Transformer model
+    """
+    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, num_heads=8, 
+                 num_layers=6, d_ff=2048, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.encoder_embedding = nn.Embedding(src_vocab_size, d_model)
+        self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model)
+        self.pos_encoding = PositionalEncoding(d_model, dropout, max_len)
+        
+        self.encoder_layers = nn.ModuleList([
+            EncoderLayer(d_model, num_heads, d_ff, dropout) 
+            for _ in range(num_layers)
+        ])
+        
+        self.decoder_layers = nn.ModuleList([
+            DecoderLayer(d_model, num_heads, d_ff, dropout) 
+            for _ in range(num_layers)
+        ])
+        
+        self.fc_out = nn.Linear(d_model, tgt_vocab_size)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
+        # Encoder
+        src_embedded = self.dropout(self.pos_encoding(self.encoder_embedding(src)))
+        enc_output = src_embedded
+        for layer in self.encoder_layers:
+            enc_output = layer(enc_output, src_mask)
+        
+        # Decoder
+        tgt_embedded = self.dropout(self.pos_encoding(self.decoder_embedding(tgt)))
+        dec_output = tgt_embedded
+        for layer in self.decoder_layers:
+            dec_output = layer(dec_output, enc_output, src_mask, tgt_mask)
+        
+        output = self.fc_out(dec_output)
+        return output
+
+class CodeDataset(Dataset):
+    """
+    Dataset for code sequences
+    """
+    def __init__(self, sequences, max_len):
+        self.sequences = sequences
+        self.max_len = max_len
+        
+    def __len__(self):
+        return len(self.sequences)
+    
+    def __getitem__(self, idx):
+        seq = self.sequences[idx]
+        # Pad sequences to max_len
+        padded = np.zeros(self.max_len, dtype=np.int64)
+        length = min(len(seq), self.max_len)
+        padded[:length] = seq[:length]
+        return torch.tensor(padded, dtype=torch.long)
+
+def create_masks(src, tgt, pad_idx):
+    """
+    Create masks for source and target sequences
+    """
+    src_mask = (src != pad_idx).unsqueeze(1).unsqueeze(2)
+    
+    tgt_mask = (tgt != pad_idx).unsqueeze(1).unsqueeze(2)
+    seq_len = tgt.size(1)
+    nopeak_mask = (1 - torch.triu(torch.ones(1, seq_len, seq_len), diagonal=1)).bool()
+    tgt_mask = tgt_mask & nopeak_mask.to(device)
+    
+    return src_mask, tgt_mask
+
+def train_model(model, dataloader, optimizer, criterion, epochs, pad_idx):
+    """
+    Training loop for the transformer model
+    """
+    model.train()
+    
+    for epoch in range(epochs):
+        total_loss = 0
+        for src, tgt in dataloader:
+            src, tgt = src.to(device), tgt.to(device)
+            
+            # Create masks
+            src_mask, tgt_mask = create_masks(src, tgt, pad_idx)
+            
+            # Forward pass
+            optimizer.zero_grad()
+            output = model(src, tgt[:, :-1], src_mask, tgt_mask[:, :-1, :-1])
+            
+            # Calculate loss
+            output_dim = output.shape[-1]
+            output = output.contiguous().view(-1, output_dim)
+            tgt = tgt[:, 1:].contiguous().view(-1)
+            
+            loss = criterion(output, tgt)
+            
+            # Backward pass
+            loss.backward()
+            optimizer.step()
+            
+            total_loss += loss.item()
+        
+        print(f'Epoch: {epoch+1}, Loss: {total_loss / len(dataloader)}')
+
+def generate_code(model, src, max_len, start_symbol, end_symbol, pad_idx):
+    """
+    Generate code sequence using the trained model
+    """
+    model.eval()
+    src = src.to(device)
+    src_mask = (src != pad_idx).unsqueeze(1).unsqueeze(2).to(device)
+    
+    memory = model.encode(src, src_mask)
+    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(device)
+    
+    for i in range(max_len-1):
+        tgt_mask = (torch.triu(torch.ones(1, ys.size(1), ys.size(1))) == 0).transpose(0, 1)
+        tgt_mask = tgt_mask.float().masked_fill(tgt_mask == 0, float('-inf')).masked_fill(tgt_mask == 1, float(0.0))
+        
+        out = model.decode(ys, memory, src_mask, tgt_mask)
+        prob = model.fc_out(out[:, -1])
+        _, next_word = torch.max(prob, dim=1)
+        next_word = next_word.item()
+        
+        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)
+        
+        if next_word == end_symbol:
+            break
+    
+    return ys
+
+# Configuration
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Using device: {device}")
+
+# Hyperparameters
+VOCAB_SIZE = 10000  # Should be adjusted based on your actual vocabulary
+D_MODEL = 512
+NUM_HEADS = 8
+NUM_LAYERS = 6
+D_FF = 2048
+DROPOUT = 0.1
+BATCH_SIZE = 32
+EPOCHS = 10
+MAX_LEN = 100
+LEARNING_RATE = 0.0001
+PAD_IDX = 0  # Assuming 0 is the padding index
+
+# Sample data - in practice you would load your code dataset here
+# For demonstration, we'll create some dummy data
+sample_data = [np.random.randint(1, VOCAB_SIZE, size=np.random.randint(10, MAX_LEN)) for _ in range(1000)]
+dataset = CodeDataset(sample_data, MAX_LEN)
+dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+# Initialize model
+model = Transformer(
+    src_vocab_size=VOCAB_SIZE,
+    tgt_vocab_size=VOCAB_SIZE,
+    d_model=D_MODEL,
+    num_heads=NUM_HEADS,
+    num_layers=NUM_LAYERS,
+    d_ff=D_FF,
+    dropout=DROPOUT,
+    max_len=MAX_LEN
+).to(device)
+
+# Loss and optimizer
+criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
+optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+# Train the model
+train_model(model, dataloader, optimizer, criterion, EPOCHS, PAD_IDX)
+
+# Example of generating code
+start_symbol = 1  # Assuming 1 is the start token
+end_symbol = 2    # Assuming 2 is the end token
+sample_input = torch.tensor([sample_data[0][:10]], dtype=torch.long)  # First 10 tokens of first sample
+generated_code = generate_code(model, sample_input, MAX_LEN, start_symbol, end_symbol, PAD_IDX)
+print("Generated code sequence:", generated_code)