My existing implementation for getting the edge features:
another State reperesentation:
In my graph structure, where i have parallel edges between two same nodes, for the Mosfet transistor (considered as one component node) having four terminals D, G, S, B, among the four terminal, the terminals D and S are connected to the same nets, which makes the two edges parallel between the same component node and nets (these parallel edge may wrongly considered as self loop by the model)
Edge feature tensor(edge_embeddings):
edge features are having high dimentional values so i used embedding technique to reduce each feature dimensions,
# Mapping dictionaries for categorical variables
device_type_mapping = {'NMOS': 0, 'PMOS': 1, 'R': 2, 'L': 3, 'C': 4, 'I': 5, 'V': 6}
device_mapping = {'M0': 0, 'M1': 1, 'M2': 2, 'M3': 3, 'M4': 4, 'M5': 5, 'M6': 6, 'M7': 7,
                  'C0': 8, 'I0': 9, 'V1': 10}
terminal_mapping = {'D0': 0, 'G0': 1, 'S0': 2, 'B0': 3, 'D1': 4, 'G1': 5, 'S1': 6, 'B1': 7,
                    'D2': 8, 'G2': 9, 'S2': 10, 'B2': 11, 'D3': 12, 'G3': 13, 'S3': 14, 'B3': 15,
                    'D4': 16, 'G4': 17, 'S4': 18, 'B4': 19, 'D5': 20, 'G5': 21, 'S5': 22, 'B5': 23,
                    'D6': 24, 'G6': 25, 'S6': 26, 'B6': 27, 'D7': 28, 'G7': 29, 'S7': 30, 'B7': 31,
                    'C0': 32, 'I0': 33, 'V1': 34}
edge_colors_mapping = {'blue': 0, 'red': 1, 'green': 2, 'grey': 3, 'yellow': 4, 'black': 5}
parallel_edges_mapping = {'T': 0, 'F': 1}
net_mapping = {'net1': 0, 'net2': 1, 'net3': 2, 'IN1': 3, 'IN2': 4, 'VOUT': 5, 'Vbias': 6, '0': 7, 'vdd!': 8}

# Define the embedding dimensions for each edge feature
device_type_embedding_dim = 2
device_embedding_dim = 2
net_embedding_dim = 2
terminal_name_embedding_dim = 3
edge_pairs_embedding_dim = 3
edge_colors_embedding_dim = 2
parallel_edges_embedding_dim = 2

# Create embedding layers for each edge feature
device_type_embedding = nn.Embedding(num_embeddings=7, embedding_dim=device_type_embedding_dim)
device_embedding = nn.Embedding(num_embeddings=15, embedding_dim=device_embedding_dim)
net_embedding = nn.Embedding(num_embeddings=9, embedding_dim=net_embedding_dim)
terminal_name_embedding = nn.Embedding(num_embeddings=39, embedding_dim=terminal_name_embedding_dim)
edge_pairs_embedding = nn.Embedding(num_embeddings=40, embedding_dim=edge_pairs_embedding_dim)
edge_colors_embedding = nn.Embedding(num_embeddings=6, embedding_dim=edge_colors_embedding_dim)
parallel_edges_embedding = nn.Embedding(num_embeddings=2, embedding_dim=parallel_edges_embedding_dim)


# Function to map categorical values to numerical indices
def map_categorical_to_index(feature_value, mapping_dict):
    if feature_value in mapping_dict:
        return mapping_dict[feature_value]
    else:
        return -1  # Handle unknown values

# Function to create embeddings for edge features
def get_edge_embeddings(edge_features):
    embeddings = []
    for edge in edge_features:
        device_type_index = map_categorical_to_index(edge['device_type'], device_type_mapping)

        device_index = map_categorical_to_index(edge['device'], device_mapping)

        terminal_index = map_categorical_to_index(edge['terminal_name'], terminal_mapping)

        net_index = map_categorical_to_index(edge['nets'], net_mapping)

        net_embed = net_embedding(torch.tensor(net_index))        

        edge_colors_index = map_categorical_to_index(edge['edge_colors'], edge_colors_mapping)

        parallel_edges_index = map_categorical_to_index(edge['Parallel edges present'], parallel_edges_mapping)
        
        device_type_embed = device_type_embedding(torch.tensor(device_type_index))
        device_embed = device_embedding(torch.tensor(device_index))
        terminal_name_embed = terminal_name_embedding(torch.tensor(terminal_index))
        edge_colors_embed = edge_colors_embedding(torch.tensor(edge_colors_index))
        parallel_edges_embed = parallel_edges_embedding(torch.tensor(parallel_edges_index))
        edge_pair_embed = torch.cat([device_embed, net_embed], dim=0)
        edge_embed = torch.cat([device_type_embed, device_embed, terminal_name_embed, edge_colors_embed, parallel_edges_embed, edge_pair_embed], dim=0)
        embeddings.append(edge_embed)
    
    return embeddings

sample edge_embeddings of one of my total 38 edges
edge_embeddings [tensor([ 0.7124, -0.1579,  2.0348,  0.4799,  2.0526, -0.9357,  0.2265,  2.0741,
         0.5741,  1.6283, -1.4729,  2.0348,  0.4799, -0.9982, -0.4223]

Here the 'device_embed' in 'edge_embed'(edge features) are belongs to the corresponding 11 components 'M0', 'M1', 'M2', 'M3', 'M4', 'M5', 'M6', 'M7', 'C0', 'I0', 'V1', and the 'edge_pair_embed' are the edge pairs corresponding to the respective edge among total 40 edges connected between the 11 components and 9 nets.

I need you to incoperate the edge features from the above existing code 'get_edge_embeddings' in to the GNN model calss 'GATModelWithConstraints' with proper adaption of my requirements.
My code need to modify:
class GATModelWithConstraints(nn.Module):
    def init(self, node_input_dim, edge_input_dim, hidden_dim, output_dim, num_layers, num_heads, constraint_dims):
        super(GATModelWithConstraints, self).init()
        
        # Define GNN layers for node features
        self.node_convs = nn.ModuleList([GATConv(node_input_dim if i == 0 else hidden_dim * num_heads,
                                                 hidden_dim,
                                                 heads=num_heads)
                                         for i in range(num_layers)])
        
        # Define GNN layers for edge embeddings
        # Note: Assumes edge_embeddings are preprocessed; if not, adapt as needed
        self.edge_convs = nn.Linear(edge_input_dim, hidden_dim * num_layers)
        
        # An extra layer that combines node and edge features
        self.combine_features = nn.Linear(hidden_dim * num_heads + hidden_dim * num_layers, output_dim)
        
        # A layer to handle (or learn) constraints within the model
        self.constraint_fc = nn.Linear(output_dim, constraint_dims)

    def forward(self, node_features, edge_embeddings, edge_index):
        # Process node features
        x = node_features
        for conv in self.node_convs:
            x = conv(x, edge_index)
            x = F.elu(x)

        # Process edge embeddings
        e = self.edge_convs(edge_embeddings)
        e = F.elu(e)

        # Concatenate node and edge features
        combined_features = torch.cat([x, e], dim=1)
        output = self.combine_features(combined_features)

        # Apply additional constraint handling
        constraints = self.constraint_fc(output)

        return output, constraints