x.py

import os
import glob
import stat
import xml.etree.ElementTree as ET
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import defaultdict
from typing import List, Dict, Any, Optional
from colorama import Fore, Style, init
from accelerate import Accelerator
from torch.utils.data import DataLoader, TensorDataset
from torch.cuda.amp import GradScaler, autocast

# Initialize colorama
init(autoreset=True)

# Set file path and output path
file_path = 'data/'
output_path = 'output/'

# Create output path if it doesn't exist
if not os.path.exists(output_path):
    os.makedirs(output_path)
    os.chmod(output_path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)  # Set full r/w permissions

# Ensure necessary files are created with full r/w permissions
def ensure_file(file_path):
    if not os.path.exists(file_path):
        with open(file_path, 'w') as f:
            pass
        os.chmod(file_path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)  # Set full r/w permissions

# Define a simple memory augmentation module
class MemoryAugmentationLayer(nn.Module):
    def __init__(self, size: int):
        super(MemoryAugmentationLayer, self).__init__()
        self.memory = nn.Parameter(torch.randn(size))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.memory

class HybridAttentionLayer(nn.Module):
    def __init__(self, size: int):
        super(HybridAttentionLayer, self).__init__()
        self.attention = nn.MultiheadAttention(size, num_heads=8)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.unsqueeze(1)  # Add sequence dimension
        attn_output, _ = self.attention(x, x, x)
        return attn_output.squeeze(1)

class DynamicFlashAttentionLayer(nn.Module):
    def __init__(self, size: int):
        super(DynamicFlashAttentionLayer, self).__init__()
        self.attention = nn.MultiheadAttention(size, num_heads=8)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.unsqueeze(1)  # Add sequence dimension
        attn_output, _ = self.attention(x, x, x)
        return attn_output.squeeze(1)

class DynamicModel(nn.Module):
    def __init__(self, sections: Dict[str, List[Dict[str, Any]]]):
        super(DynamicModel, self).__init__()
        self.sections = nn.ModuleDict()
        if not sections:
            sections = {
                'default': [{
                    'input_size': 128,
                    'output_size': 256,
                    'activation': 'relu',
                    'batch_norm': True,
                    'dropout': 0.1
                }]
            }
        for section_name, layers in sections.items():
            self.sections[section_name] = nn.ModuleList()
            for layer_params in layers:
                print(f"Creating layer in section '{section_name}' with params: {layer_params}")
                self.sections[section_name].append(self.create_layer(layer_params))

    def create_layer(self, layer_params: Dict[str, Any]) -> nn.Module:
        layers = []
        layers.append(nn.Linear(layer_params['input_size'], layer_params['output_size']))
        if layer_params.get('batch_norm', False):
            layers.append(nn.BatchNorm1d(layer_params['output_size']))
        activation = layer_params.get('activation', 'relu')
        if activation == 'relu':
            layers.append(nn.ReLU(inplace=True))
        elif activation == 'tanh':
            layers.append(nn.Tanh())
        elif activation == 'sigmoid':
            layers.append(nn.Sigmoid())
        elif activation == 'leaky_relu':
            layers.append(nn.LeakyReLU(negative_slope=0.01, inplace=True))
        elif activation == 'elu':
            layers.append(nn.ELU(alpha=1.0, inplace=True))
        elif activation is not None:
            raise ValueError(f"Unsupported activation function: {activation}")
        if dropout_rate := layer_params.get('dropout', 0.0):
            layers.append(nn.Dropout(p=dropout_rate))
        if hidden_layers := layer_params.get('hidden_layers', []):
            for hidden_layer_params in hidden_layers:
                layers.append(self.create_layer(hidden_layer_params))
        if layer_params.get('memory_augmentation', True):
            layers.append(MemoryAugmentationLayer(layer_params['output_size']))
        if layer_params.get('hybrid_attention', True):
            layers.append(HybridAttentionLayer(layer_params['output_size']))
        if layer_params.get('dynamic_flash_attention', True):
            layers.append(DynamicFlashAttentionLayer(layer_params['output_size']))
        return nn.Sequential(*layers)

    def forward(self, x: torch.Tensor, section_name: Optional[str] = None) -> torch.Tensor:
        if section_name is not None:
            if section_name not in self.sections:
                raise KeyError(f"Section '{section_name}' not found in model")
            for layer in self.sections[section_name]:
                x = layer(x)
        else:
            for section_name, layers in self.sections.items():
                for layer in layers:
                    x = layer(x)
        return x

def parse_xml_file(file_path: str) -> List[Dict[str, Any]]:
    tree = ET.parse(file_path)
    root = tree.getroot()
    layers = []
    for layer in root.findall('.//layer'):
        layer_params = {}
        layer_params['input_size'] = int(layer.get('input_size', 128))
        layer_params['output_size'] = int(layer.get('output_size', 256))
        layer_params['activation'] = layer.get('activation', 'relu').lower()
        if layer_params['activation'] not in ['relu', 'tanh', 'sigmoid', 'none']:
            raise ValueError(f"Unsupported activation function: {layer_params['activation']}")
        if layer_params['input_size'] <= 0 or layer_params['output_size'] <= 0:
            raise ValueError("Layer dimensions must be positive integers")
        layers.append(layer_params)
    if not layers:
        layers.append({
            'input_size': 128,
            'output_size': 256,
            'activation': 'relu'
        })
    return layers

def create_model_from_folder(folder_path: str) -> DynamicModel:
    sections = defaultdict(list)
    if not os.path.exists(folder_path):
        print(f"Warning: Folder {folder_path} does not exist. Creating model with default configuration.")
        return DynamicModel({})
    xml_files_found = False
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.endswith('.xml'):
                xml_files_found = True
                file_path = os.path.join(root, file)
                try:
                    layers = parse_xml_file(file_path)
                    section_name = os.path.basename(root).replace('.', '_')
                    sections[section_name].extend(layers)
                except Exception as e:
                    print(f"Error processing {file_path}: {str(e)}")
    if not xml_files_found:
        print("Warning: No XML files found. Creating model with default configuration.")
        return DynamicModel({})
    return DynamicModel(dict(sections))

def main():
    print(Fore.CYAN + "Starting conversion...")

    # Create the dynamic model from the folder
    model = create_model_from_folder(file_path)
    print(f"Created dynamic PyTorch model with sections: {list(model.sections.keys())}")

    # Print the model architecture
    print(model)

    # Ensure the input tensor size matches the expected input size
    first_section = next(iter(model.sections.keys()))
    first_layer = model.sections[first_section][0]
    input_features = first_layer[0].in_features
    sample_input = torch.randn(1, input_features)
    output = model(sample_input)
    print(f"Sample output shape: {output.shape}")

    # Training setup
    accelerator = Accelerator()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    num_epochs = 10
    dataset = TensorDataset(
        torch.randn(100, input_features),
        torch.randint(0, 2, (100,))
    )
    train_dataloader = DataLoader(
        dataset,
        batch_size=8,  # Reduced batch size
        shuffle=True
    )
    model, optimizer, train_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader
    )
    scaler = GradScaler()  # Mixed precision training

    # Training loop
    for epoch in range(num_epochs):
        model.train()
        total_loss = 0
        for batch_idx, (inputs, labels) in enumerate(train_dataloader):
            optimizer.zero_grad()
            with autocast():  # Mixed precision training
                outputs = model(inputs)
                loss = criterion(outputs, labels)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            total_loss += loss.item()
        avg_loss = total_loss / len(train_dataloader)
        print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss:.4f}")

if __name__ == "__main__":
    main()