Buckets:

tahamajs
/

Sysmem2_in_AI

Files

xet

tahamajs/Sysmem2_in_AI / ComputerAssignments /CA6_systematic_generalization /src /models /neural_architectures.py

tahamajs

5 months ago

download

raw

15.8 kB

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Dict, Tuple, Any, Optional, Union
import math
from einops import rearrange, repeat
import numpy as np

class PositionalEncoding(nn.Module):

    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)

        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pe[: x.size(0), :]

class MultiHeadCompositionalAttention(nn.Module):

    def __init__(self, d_model: int, num_heads: int = 8, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0

        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 = math.sqrt(self.d_k)


        self.compositional_bias = nn.Parameter(torch.randn(num_heads, 1, 1))

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        batch_size = query.size(0)


        Q = (
            self.w_q(query)
            .view(batch_size, -1, self.num_heads, self.d_k)
            .transpose(1, 2)
        )
        K = self.w_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = (
            self.w_v(value)
            .view(batch_size, -1, self.num_heads, self.d_k)
            .transpose(1, 2)
        )


        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale


        scores = scores + self.compositional_bias

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)


        context = torch.matmul(attention_weights, V)
        context = (
            context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        )

        return self.w_o(context), attention_weights

class ModularNetwork(nn.Module):

    def __init__(
        self,
        primitive_vocab_size: int,
        operator_vocab_size: int,
        embed_dim: int = 64,
        hidden_dim: int = 128,
        num_layers: int = 2,
    ):
        super().__init__()

        self.embed_dim = embed_dim
        self.hidden_dim = hidden_dim


        self.primitive_embeddings = nn.Embedding(primitive_vocab_size, embed_dim)
        self.operator_embeddings = nn.Embedding(operator_vocab_size, embed_dim)


        self.primitive_processor = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, embed_dim),
            nn.LayerNorm(embed_dim),
        )

        self.operator_processor = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, embed_dim),
            nn.LayerNorm(embed_dim),
        )


        composition_layers = []
        for i in range(num_layers):
            composition_layers.extend(
                [
                    nn.Linear(embed_dim * 3 if i == 0 else embed_dim, hidden_dim),
                    nn.ReLU(),
                    nn.Dropout(0.1),
                    nn.Linear(hidden_dim, embed_dim),
                    nn.LayerNorm(embed_dim),
                ]
            )

        self.composer = nn.Sequential(*composition_layers)


        self.output_layer = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, 1),
        )

    def forward(
        self,
        primitive1_id: torch.Tensor,
        operator_id: torch.Tensor,
        primitive2_id: torch.Tensor,
    ) -> torch.Tensor:

        prim1_embed = self.primitive_embeddings(primitive1_id)
        op_embed = self.operator_embeddings(operator_id)
        prim2_embed = self.primitive_embeddings(primitive2_id)


        prim1_processed = self.primitive_processor(prim1_embed)
        op_processed = self.operator_processor(op_embed)
        prim2_processed = self.primitive_processor(prim2_embed)


        combined = torch.cat([prim1_processed, op_processed, prim2_processed], dim=-1)
        composed = self.composer(combined)

        return self.output_layer(composed)

class AttentionComposer(nn.Module):

    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 64,
        hidden_dim: int = 128,
        num_heads: int = 8,
        num_layers: int = 3,
        dropout: float = 0.1,
    ):
        super().__init__()

        self.embed_dim = embed_dim
        self.num_heads = num_heads


        self.embeddings = nn.Embedding(vocab_size, embed_dim)
        self.pos_encoding = PositionalEncoding(embed_dim)


        self.attention_layers = nn.ModuleList(
            [
                MultiHeadCompositionalAttention(embed_dim, num_heads, dropout)
                for _ in range(num_layers)
            ]
        )


        self.ffn_layers = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(embed_dim, hidden_dim),
                    nn.ReLU(),
                    nn.Dropout(dropout),
                    nn.Linear(hidden_dim, embed_dim),
                    nn.Dropout(dropout),
                )
                for _ in range(num_layers)
            ]
        )


        self.layer_norms1 = nn.ModuleList(
            [nn.LayerNorm(embed_dim) for _ in range(num_layers)]
        )
        self.layer_norms2 = nn.ModuleList(
            [nn.LayerNorm(embed_dim) for _ in range(num_layers)]
        )


        self.output_layer = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 1),
        )

    def forward(
        self, token_ids: torch.Tensor, mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:

        embeddings = self.embeddings(token_ids)
        embeddings = self.pos_encoding(embeddings.transpose(0, 1)).transpose(0, 1)


        x = embeddings
        for i, (attention, ffn, norm1, norm2) in enumerate(
            zip(
                self.attention_layers,
                self.ffn_layers,
                self.layer_norms1,
                self.layer_norms2,
            )
        ):

            attn_out, _ = attention(x, x, x, mask)
            x = norm1(x + attn_out)


            ffn_out = ffn(x)
            x = norm2(x + ffn_out)


        pooled = x.mean(dim=1)
        return self.output_layer(pooled)

class GraphComposer(nn.Module):

    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 64,
        num_layers: int = 3,
        dropout: float = 0.1,
    ):
        super().__init__()

        self.embeddings = nn.Embedding(vocab_size, embed_dim)
        self.num_layers = num_layers


        self.message_layers = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(embed_dim, embed_dim), nn.ReLU(), nn.Dropout(dropout)
                )
                for _ in range(num_layers)
            ]
        )

        self.update_layers = nn.ModuleList(
            [nn.GRU(embed_dim, embed_dim, batch_first=True) for _ in range(num_layers)]
        )


        self.attention_layers = nn.ModuleList(
            [
                nn.MultiheadAttention(
                    embed_dim, num_heads=4, batch_first=True, dropout=dropout
                )
                for _ in range(num_layers)
            ]
        )


        self.output_layer = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embed_dim * 2, 1),
        )

    def forward(
        self, node_ids: torch.Tensor, adjacency_matrix: torch.Tensor
    ) -> torch.Tensor:

        node_features = self.embeddings(node_ids)


        for layer_idx in range(self.num_layers):

            messages = self.message_layers[layer_idx](node_features)


            attended_messages, _ = self.attention_layers[layer_idx](
                messages, messages, messages
            )


            node_features, _ = self.update_layers[layer_idx](attended_messages)


        graph_representation = node_features.mean(dim=1)
        return self.output_layer(graph_representation)

class TreeLSTM(nn.Module):

    def __init__(self, input_size: int, hidden_size: int, dropout: float = 0.1):
        super().__init__()

        self.hidden_size = hidden_size


        self.i_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)
        self.fl_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)
        self.fr_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)
        self.o_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)
        self.u_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)

        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        input_vec: torch.Tensor,
        left_state: Tuple[torch.Tensor, torch.Tensor],
        right_state: Tuple[torch.Tensor, torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        hl, cl = left_state
        hr, cr = right_state


        combined_input = torch.cat([input_vec, hl, hr], dim=-1)
        combined_input = self.dropout(combined_input)


        i = torch.sigmoid(self.i_gate(combined_input))
        fl = torch.sigmoid(self.fl_gate(combined_input))
        fr = torch.sigmoid(self.fr_gate(combined_input))
        o = torch.sigmoid(self.o_gate(combined_input))
        u = torch.tanh(self.u_gate(combined_input))


        c = i * u + fl * cl + fr * cr


        h = o * torch.tanh(c)

        return h, c

class MetaLearningComposer(nn.Module):

    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 64,
        meta_lr: float = 0.01,
        inner_lr: float = 0.001,
    ):
        super().__init__()

        self.embed_dim = embed_dim
        self.meta_lr = meta_lr
        self.inner_lr = inner_lr


        self.base_embedding = nn.Embedding(vocab_size, embed_dim)
        self.base_composer = nn.Sequential(
            nn.Linear(embed_dim * 2, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, 1),
        )


        self.meta_embedding = nn.Embedding(vocab_size, embed_dim)
        self.adaptation_network = nn.Sequential(
            nn.Linear(embed_dim * 3, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, embed_dim),
        )


        self.fast_weights = nn.Parameter(torch.randn(embed_dim))

    def forward(
        self,
        support_examples: List[Tuple[torch.Tensor, torch.Tensor]],
        query_examples: torch.Tensor,
    ) -> torch.Tensor:

        adapted_params = self.adapt_to_support(support_examples)


        query_embeddings = self.base_embedding(query_examples)
        adapted_embeddings = self.apply_adaptation(query_embeddings, adapted_params)


        if query_embeddings.shape[1] >= 2:
            combined = torch.cat(
                [adapted_embeddings[:, 0], adapted_embeddings[:, 1]], dim=-1
            )
        else:
            combined = adapted_embeddings.mean(dim=1)

        return self.base_composer(combined)

    def adapt_to_support(
        self, support_examples: List[Tuple[torch.Tensor, torch.Tensor]]
    ) -> torch.Tensor:
        if not support_examples:
            return self.fast_weights


        support_embeddings = []
        support_targets = []

        for inputs, targets in support_examples:
            embeddings = self.base_embedding(inputs)
            support_embeddings.append(embeddings.mean(dim=1))
            support_targets.append(targets)

        if support_embeddings:
            support_tensor = torch.stack(support_embeddings)
            target_tensor = torch.stack(support_targets)


            adaptation_signal = torch.cat(
                [support_tensor.mean(dim=0), target_tensor.mean(dim=0).unsqueeze(0)],
                dim=0,
            )
            adapted_params = self.adaptation_network(adaptation_signal.unsqueeze(0))
            return adapted_params.squeeze(0)

        return self.fast_weights

    def apply_adaptation(
        self, embeddings: torch.Tensor, adapted_params: torch.Tensor
    ) -> torch.Tensor:
        return embeddings + adapted_params.unsqueeze(0).unsqueeze(0)

class HierarchicalComposer(nn.Module):

    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 64,
        hidden_dim: int = 128,
        num_levels: int = 3,
        dropout: float = 0.1,
    ):
        super().__init__()

        self.num_levels = num_levels
        self.embeddings = nn.Embedding(vocab_size, embed_dim)


        self.composition_layers = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(embed_dim if level == 0 else embed_dim, embed_dim),
                    nn.ReLU(),
                    nn.Dropout(dropout),
                    nn.Linear(embed_dim, embed_dim),
                    nn.LayerNorm(embed_dim),
                )
                for level in range(num_levels)
            ]
        )


        self.level_attention = nn.ModuleList(
            [
                nn.MultiheadAttention(
                    embed_dim, num_heads=4, batch_first=True, dropout=dropout
                )
                for _ in range(num_levels)
            ]
        )


        self.cross_level_attention = nn.ModuleList(
            [
                nn.MultiheadAttention(
                    embed_dim, num_heads=2, batch_first=True, dropout=dropout
                )
                for _ in range(num_levels - 1)
            ]
        )


        self.output_layer = nn.Sequential(
            nn.Linear(embed_dim * num_levels, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, token_sequence: torch.Tensor) -> torch.Tensor:

        embeddings = self.embeddings(token_sequence)

        level_representations = []
        current_repr = embeddings


        for level in range(self.num_levels):

            composed = self.composition_layers[level](current_repr)


            attended, _ = self.level_attention[level](composed, composed, composed)


            level_repr = attended.mean(dim=1)
            level_representations.append(level_repr)


            if level < self.num_levels - 1:

                current_repr = attended[:, ::2]


                if level > 0:
                    cross_attended, _ = self.cross_level_attention[level - 1](
                        current_repr,
                        level_representations[-2].unsqueeze(1),
                        level_representations[-2].unsqueeze(1),
                    )
                    current_repr = cross_attended.squeeze(1)


        combined_repr = torch.cat(level_representations, dim=-1)
        return self.output_layer(combined_repr)

Xet Storage Details

Size:: 15.8 kB
Xet hash:: 2b52248a95d819926384982b9bdb3c646bc419659afe2e240cb2dcd1ad246a6e

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from typing import List, Dict, Tuple, Any, Optional, Union
	import math
	from einops import rearrange, repeat
	import numpy as np

	class PositionalEncoding(nn.Module):

	def __init__(self, d_model: int, max_len: int = 5000):
	super().__init__()

	pe = torch.zeros(max_len, d_model)
	position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
	div_term = torch.exp(
	torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
	)

	pe[:, 0::2] = torch.sin(position * div_term)
	pe[:, 1::2] = torch.cos(position * div_term)
	pe = pe.unsqueeze(0).transpose(0, 1)

	self.register_buffer("pe", pe)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return x + self.pe[: x.size(0), :]

	class MultiHeadCompositionalAttention(nn.Module):

	def __init__(self, d_model: int, num_heads: int = 8, dropout: float = 0.1):
	super().__init__()
	assert d_model % num_heads == 0

	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 = math.sqrt(self.d_k)


	self.compositional_bias = nn.Parameter(torch.randn(num_heads, 1, 1))

	def forward(
	self,
	query: torch.Tensor,
	key: torch.Tensor,
	value: torch.Tensor,
	mask: Optional[torch.Tensor] = None,
	) -> Tuple[torch.Tensor, torch.Tensor]:
	batch_size = query.size(0)


	Q = (
	self.w_q(query)
	.view(batch_size, -1, self.num_heads, self.d_k)
	.transpose(1, 2)
	)
	K = self.w_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
	V = (
	self.w_v(value)
	.view(batch_size, -1, self.num_heads, self.d_k)
	.transpose(1, 2)
	)


	scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale


	scores = scores + self.compositional_bias

	if mask is not None:
	scores = scores.masked_fill(mask == 0, -1e9)

	attention_weights = F.softmax(scores, dim=-1)
	attention_weights = self.dropout(attention_weights)


	context = torch.matmul(attention_weights, V)
	context = (
	context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
	)

	return self.w_o(context), attention_weights

	class ModularNetwork(nn.Module):

	def __init__(
	self,
	primitive_vocab_size: int,
	operator_vocab_size: int,
	embed_dim: int = 64,
	hidden_dim: int = 128,
	num_layers: int = 2,
	):
	super().__init__()

	self.embed_dim = embed_dim
	self.hidden_dim = hidden_dim


	self.primitive_embeddings = nn.Embedding(primitive_vocab_size, embed_dim)
	self.operator_embeddings = nn.Embedding(operator_vocab_size, embed_dim)


	self.primitive_processor = nn.Sequential(
	nn.Linear(embed_dim, hidden_dim),
	nn.ReLU(),
	nn.Dropout(0.1),
	nn.Linear(hidden_dim, embed_dim),
	nn.LayerNorm(embed_dim),
	)

	self.operator_processor = nn.Sequential(
	nn.Linear(embed_dim, hidden_dim),
	nn.ReLU(),
	nn.Dropout(0.1),
	nn.Linear(hidden_dim, embed_dim),
	nn.LayerNorm(embed_dim),
	)


	composition_layers = []
	for i in range(num_layers):
	composition_layers.extend(
	[
	nn.Linear(embed_dim * 3 if i == 0 else embed_dim, hidden_dim),
	nn.ReLU(),
	nn.Dropout(0.1),
	nn.Linear(hidden_dim, embed_dim),
	nn.LayerNorm(embed_dim),
	]
	)

	self.composer = nn.Sequential(*composition_layers)


	self.output_layer = nn.Sequential(
	nn.Linear(embed_dim, hidden_dim),
	nn.ReLU(),
	nn.Dropout(0.1),
	nn.Linear(hidden_dim, 1),
	)

	def forward(
	self,
	primitive1_id: torch.Tensor,
	operator_id: torch.Tensor,
	primitive2_id: torch.Tensor,
	) -> torch.Tensor:

	prim1_embed = self.primitive_embeddings(primitive1_id)
	op_embed = self.operator_embeddings(operator_id)
	prim2_embed = self.primitive_embeddings(primitive2_id)


	prim1_processed = self.primitive_processor(prim1_embed)
	op_processed = self.operator_processor(op_embed)
	prim2_processed = self.primitive_processor(prim2_embed)


	combined = torch.cat([prim1_processed, op_processed, prim2_processed], dim=-1)
	composed = self.composer(combined)

	return self.output_layer(composed)

	class AttentionComposer(nn.Module):

	def __init__(
	self,
	vocab_size: int,
	embed_dim: int = 64,
	hidden_dim: int = 128,
	num_heads: int = 8,
	num_layers: int = 3,
	dropout: float = 0.1,
	):
	super().__init__()

	self.embed_dim = embed_dim
	self.num_heads = num_heads


	self.embeddings = nn.Embedding(vocab_size, embed_dim)
	self.pos_encoding = PositionalEncoding(embed_dim)


	self.attention_layers = nn.ModuleList(
	[
	MultiHeadCompositionalAttention(embed_dim, num_heads, dropout)
	for _ in range(num_layers)
	]
	)


	self.ffn_layers = nn.ModuleList(
	[
	nn.Sequential(
	nn.Linear(embed_dim, hidden_dim),
	nn.ReLU(),
	nn.Dropout(dropout),
	nn.Linear(hidden_dim, embed_dim),
	nn.Dropout(dropout),
	)
	for _ in range(num_layers)
	]
	)


	self.layer_norms1 = nn.ModuleList(
	[nn.LayerNorm(embed_dim) for _ in range(num_layers)]
	)
	self.layer_norms2 = nn.ModuleList(
	[nn.LayerNorm(embed_dim) for _ in range(num_layers)]
	)


	self.output_layer = nn.Sequential(
	nn.Linear(embed_dim, hidden_dim),
	nn.ReLU(),
	nn.Dropout(dropout),
	nn.Linear(hidden_dim, 1),
	)

	def forward(
	self, token_ids: torch.Tensor, mask: Optional[torch.Tensor] = None
	) -> torch.Tensor:

	embeddings = self.embeddings(token_ids)
	embeddings = self.pos_encoding(embeddings.transpose(0, 1)).transpose(0, 1)


	x = embeddings
	for i, (attention, ffn, norm1, norm2) in enumerate(
	zip(
	self.attention_layers,
	self.ffn_layers,
	self.layer_norms1,
	self.layer_norms2,
	)
	):

	attn_out, _ = attention(x, x, x, mask)
	x = norm1(x + attn_out)


	ffn_out = ffn(x)
	x = norm2(x + ffn_out)


	pooled = x.mean(dim=1)
	return self.output_layer(pooled)

	class GraphComposer(nn.Module):

	def __init__(
	self,
	vocab_size: int,
	embed_dim: int = 64,
	num_layers: int = 3,
	dropout: float = 0.1,
	):
	super().__init__()

	self.embeddings = nn.Embedding(vocab_size, embed_dim)
	self.num_layers = num_layers


	self.message_layers = nn.ModuleList(
	[
	nn.Sequential(
	nn.Linear(embed_dim, embed_dim), nn.ReLU(), nn.Dropout(dropout)
	)
	for _ in range(num_layers)
	]
	)

	self.update_layers = nn.ModuleList(
	[nn.GRU(embed_dim, embed_dim, batch_first=True) for _ in range(num_layers)]
	)


	self.attention_layers = nn.ModuleList(
	[
	nn.MultiheadAttention(
	embed_dim, num_heads=4, batch_first=True, dropout=dropout
	)
	for _ in range(num_layers)
	]
	)


	self.output_layer = nn.Sequential(
	nn.Linear(embed_dim, embed_dim * 2),
	nn.ReLU(),
	nn.Dropout(dropout),
	nn.Linear(embed_dim * 2, 1),
	)

	def forward(
	self, node_ids: torch.Tensor, adjacency_matrix: torch.Tensor
	) -> torch.Tensor:

	node_features = self.embeddings(node_ids)


	for layer_idx in range(self.num_layers):

	messages = self.message_layers[layer_idx](node_features)


	attended_messages, _ = self.attention_layers[layer_idx](
	messages, messages, messages
	)


	node_features, _ = self.update_layers[layer_idx](attended_messages)


	graph_representation = node_features.mean(dim=1)
	return self.output_layer(graph_representation)

	class TreeLSTM(nn.Module):

	def __init__(self, input_size: int, hidden_size: int, dropout: float = 0.1):
	super().__init__()

	self.hidden_size = hidden_size


	self.i_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)
	self.fl_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)
	self.fr_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)
	self.o_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)
	self.u_gate = nn.Linear(input_size + 2 * hidden_size, hidden_size)

	self.dropout = nn.Dropout(dropout)

	def forward(
	self,
	input_vec: torch.Tensor,
	left_state: Tuple[torch.Tensor, torch.Tensor],
	right_state: Tuple[torch.Tensor, torch.Tensor],
	) -> Tuple[torch.Tensor, torch.Tensor]:
	hl, cl = left_state
	hr, cr = right_state


	combined_input = torch.cat([input_vec, hl, hr], dim=-1)
	combined_input = self.dropout(combined_input)


	i = torch.sigmoid(self.i_gate(combined_input))
	fl = torch.sigmoid(self.fl_gate(combined_input))
	fr = torch.sigmoid(self.fr_gate(combined_input))
	o = torch.sigmoid(self.o_gate(combined_input))
	u = torch.tanh(self.u_gate(combined_input))


	c = i * u + fl * cl + fr * cr


	h = o * torch.tanh(c)

	return h, c

	class MetaLearningComposer(nn.Module):

	def __init__(
	self,
	vocab_size: int,
	embed_dim: int = 64,
	meta_lr: float = 0.01,
	inner_lr: float = 0.001,
	):
	super().__init__()

	self.embed_dim = embed_dim
	self.meta_lr = meta_lr
	self.inner_lr = inner_lr


	self.base_embedding = nn.Embedding(vocab_size, embed_dim)
	self.base_composer = nn.Sequential(
	nn.Linear(embed_dim * 2, embed_dim),
	nn.ReLU(),
	nn.Linear(embed_dim, embed_dim),
	nn.ReLU(),
	nn.Linear(embed_dim, 1),
	)


	self.meta_embedding = nn.Embedding(vocab_size, embed_dim)
	self.adaptation_network = nn.Sequential(
	nn.Linear(embed_dim * 3, embed_dim),
	nn.ReLU(),
	nn.Linear(embed_dim, embed_dim),
	)


	self.fast_weights = nn.Parameter(torch.randn(embed_dim))

	def forward(
	self,
	support_examples: List[Tuple[torch.Tensor, torch.Tensor]],
	query_examples: torch.Tensor,
	) -> torch.Tensor:

	adapted_params = self.adapt_to_support(support_examples)


	query_embeddings = self.base_embedding(query_examples)
	adapted_embeddings = self.apply_adaptation(query_embeddings, adapted_params)


	if query_embeddings.shape[1] >= 2:
	combined = torch.cat(
	[adapted_embeddings[:, 0], adapted_embeddings[:, 1]], dim=-1
	)
	else:
	combined = adapted_embeddings.mean(dim=1)

	return self.base_composer(combined)

	def adapt_to_support(
	self, support_examples: List[Tuple[torch.Tensor, torch.Tensor]]
	) -> torch.Tensor:
	if not support_examples:
	return self.fast_weights


	support_embeddings = []
	support_targets = []

	for inputs, targets in support_examples:
	embeddings = self.base_embedding(inputs)
	support_embeddings.append(embeddings.mean(dim=1))
	support_targets.append(targets)

	if support_embeddings:
	support_tensor = torch.stack(support_embeddings)
	target_tensor = torch.stack(support_targets)


	adaptation_signal = torch.cat(
	[support_tensor.mean(dim=0), target_tensor.mean(dim=0).unsqueeze(0)],
	dim=0,
	)
	adapted_params = self.adaptation_network(adaptation_signal.unsqueeze(0))
	return adapted_params.squeeze(0)

	return self.fast_weights

	def apply_adaptation(
	self, embeddings: torch.Tensor, adapted_params: torch.Tensor
	) -> torch.Tensor:
	return embeddings + adapted_params.unsqueeze(0).unsqueeze(0)

	class HierarchicalComposer(nn.Module):

	def __init__(
	self,
	vocab_size: int,
	embed_dim: int = 64,
	hidden_dim: int = 128,
	num_levels: int = 3,
	dropout: float = 0.1,
	):
	super().__init__()

	self.num_levels = num_levels
	self.embeddings = nn.Embedding(vocab_size, embed_dim)


	self.composition_layers = nn.ModuleList(
	[
	nn.Sequential(
	nn.Linear(embed_dim if level == 0 else embed_dim, embed_dim),
	nn.ReLU(),
	nn.Dropout(dropout),
	nn.Linear(embed_dim, embed_dim),
	nn.LayerNorm(embed_dim),
	)
	for level in range(num_levels)
	]
	)


	self.level_attention = nn.ModuleList(
	[
	nn.MultiheadAttention(
	embed_dim, num_heads=4, batch_first=True, dropout=dropout
	)
	for _ in range(num_levels)
	]
	)


	self.cross_level_attention = nn.ModuleList(
	[
	nn.MultiheadAttention(
	embed_dim, num_heads=2, batch_first=True, dropout=dropout
	)
	for _ in range(num_levels - 1)
	]
	)


	self.output_layer = nn.Sequential(
	nn.Linear(embed_dim * num_levels, hidden_dim),
	nn.ReLU(),
	nn.Dropout(dropout),
	nn.Linear(hidden_dim, 1),
	)

	def forward(self, token_sequence: torch.Tensor) -> torch.Tensor:

	embeddings = self.embeddings(token_sequence)

	level_representations = []
	current_repr = embeddings


	for level in range(self.num_levels):

	composed = self.composition_layers[level](current_repr)


	attended, _ = self.level_attention[level](composed, composed, composed)


	level_repr = attended.mean(dim=1)
	level_representations.append(level_repr)


	if level < self.num_levels - 1:

	current_repr = attended[:, ::2]


	if level > 0:
	cross_attended, _ = self.cross_level_attention[level - 1](
	current_repr,
	level_representations[-2].unsqueeze(1),
	level_representations[-2].unsqueeze(1),
	)
	current_repr = cross_attended.squeeze(1)


	combined_repr = torch.cat(level_representations, dim=-1)
	return self.output_layer(combined_repr)