Swarming algorithms like PSO, Ant Colony, Sakana, and more in PyTorch 😊 
from swarms_torch import ParticleSwarmOptimization

# test
pso = ParticleSwarmOptimization(goal="Attention is all you need", n_particles=100)
pso.optimize(iterations=1000)

example.py
kyegomez/swarms-pytorch
import torch
from swarms_torch.fish_school import Fish, FishSchool

# Create random source and target sequences
src = torch.randn(10, 32, 512)
tgt = torch.randn(10, 32, 512)

# Create random labels
labels = torch.randint(0, 512, (10, 32))

# Create a fish and train it on the random data
fish = Fish(512, 8, 6)
fish.train(src, tgt, labels)
print(fish.food)  # Print the fish's food

# Create a fish school and optimize it on the random data
school = FishSchool(10, 512, 8, 6, 100)
school.forward(src, tgt, labels)
print(school.fish[0].food)  # Print the first fish's food

fish_school_example.py
kyegomez/swarms-pytorch
from swarms_torch.ant_colony_swarm import AntColonyOptimization

# Usage:
goal_string = "Hello ACO"
aco = AntColonyOptimization(goal_string, num_iterations=1000)
best_solution = aco.optimize()
print("Best Matched String:", best_solution)

playground/ant_colony.py
kyegomez/swarms-pytorch
import torch
from swarms_torch.neuronal_transformer import NNTransformer

x = torch.randn(1, 10)

network = NNTransformer(
    # transformer cells
    neuron_count=5,
    # num states
    num_states=10,
    # input dim
    input_dim=10,
    # output dim
    output_dim=10,
    # nhead
    nhead=2,
)


output = network(x)
print(output)

playground/nnt.py
kyegomez/swarms-pytorch
from swarms_torch import SPO

# Example Usage
goal_str = "Attention is all you need"
optimizer = SPO(goal_str)
optimizer.optimize()
print(f"Optimized String: {optimizer.best_string()}")

playground/sop.py
kyegomez/swarms-pytorch
from swarms_torch.ant_colony_swarm import AntColonyOptimization
from swarms_torch.cellular_transformer import CellularSwarm
from swarms_torch.fish_school import Fish, FishSchool
from swarms_torch.neuronal_transformer import NNTransformer
from swarms_torch.particle_swarm import ParticleSwarmOptimization
from swarms_torch.queen_bee import QueenBeeGa
from swarms_torch.spiral_optimization import SPO
from swarms_torch.multi_swarm_pso import MultiSwarmPSO
from swarms_torch.transformer_pso import Particle, TransformerParticleSwarmOptimization


__all__ = [
    "ParticleSwarmOptimization",
    "AntColonyOptimization",
    "QueenBeeGa",
    "NNTransformer",
    "CellularSwarm",
    "SPO",
    "Fish",
    "FishSchool",
    "MultiSwarmPSO",
    "Particle",
    "TransformerParticleSwarmOptimization",
]

swarms_torch/__init__.py
kyegomez/swarms-pytorch
import torch


class AntColonyOptimization:
    """
    Ant Colony Optimization
    Overview: https://en.wikipedia.org/wiki/Ant_colony_optimization_algorithms

    How does it work?
    1. Initialize pheromone levels for each path
    2. For each ant, choose the next path based on the pheromone levels
    3. Update the pheromone levels
    4. Repeat step 2 to 3 until the maximum number of iterations is reached

    Parameters
    ----------
    goal: str
        The goal string to be optimized
    num_ants: int
        Number of ants
    evaporation_rate: float
        Evaporation rate

    Usage
    -----
    from swarms_torch import AntColonyOptimization

    goal_string = "Hello ACO"
    aco = AntColonyOptimization(goal_string, num_iterations=1000)
    best_solution = aco.optimize()

    print("Best Matched String:", best_solution)

    Features to implement
    --------
    1. Add a stopping criterion
    2. Add a callback function to track the progress
    3. Add a function to plot the pheromone levels
    4. Add a function to plot the ants
    5. Add a function to plot the best solution

    """

    def __init__(
        self,
        goal: str = None,
        num_ants: int = 10000,
        evaporation_rate: float = 0.1,
        alpha: int = 1,
        beta: int = 1,
        num_iterations: int = 10010,
    ):
        self.goal = torch.tensor([ord(c) for c in goal], dtype=torch.float32)
        self.num_ants = num_ants
        self.evaporation_rate = evaporation_rate
        self.alpha = alpha
        self.beta = beta
        self.num_iterations = num_iterations
        # Pheromone levels can be initialized for different paths
        # (architectures)
        self.pheromones = torch.ones(num_ants)
        self.solutions = []

    def fitness(self, solution):
        """Fitness of a solution"""
        return -torch.norm(solution - self.goal)

    def update_pheromones(self):
        """Update pheromone levels"""
        for i, solution in enumerate(self.solutions):
            self.pheromones[i] = (1 - self.evaporation_rate) * self.pheromones[
                i
            ] + self.fitness(solution)

    def choose_next_path(self):
        """Choose the next path based on the pheromone levels"""
        probabilities = (self.pheromones**self.alpha) * (
            (1.0 / (1 + self.pheromones)) ** self.beta
        )

        probabilities /= probabilities.sum()

        return torch.multinomial(probabilities, num_samples=1).item()

    def optimize(self):
        """Optimize the goal string"""
        for iteration in range(self.num_iterations):
            self.solutions = []
            for _ in range(self.num_ants):
                # This is a placeholder. Actual implementation will define how
                # ants traverse the search space.
                solution = torch.randint(
                    32, 127, (len(self.goal),), dtype=torch.float32
                )  # Random characters.
                self.solutions.append(solution)
            self.update_pheromones()

        best_solution_index = self.pheromones.argmax().item()
        best_solution = self.solutions[best_solution_index]
        return "".join([chr(int(c)) for c in best_solution])

swarms_torch/ant_colony_swarm.py
kyegomez/swarms-pytorch
from math import ceil, log
from typing import Callable, Optional, Tuple, Union

import torch
import torch.nn.functional as F
from einops import pack, rearrange, unpack
from torch import Tensor
from torch.nn import Module


def exists(val):
    return val is not None


def default(val, d):
    return val if exists(val) else d


def identity(t, *args, **kwargs):
    return t


def cast_tuple(t, length=1):
    return t if isinstance(t, tuple) else (t,) * length


def eval_decorator(fn):
    def inner(self, *args, **kwargs):
        was_training = self.training
        self.eval()
        out = fn(self, *args, **kwargs)
        self.train(was_training)
        return out

    return inner


# for variable lengthed prefixes


def align_right(t, lens, pad_id=0):
    batch, seq_len, device, dtype = *t.shape, t.device, t.dtype

    assert lens.ndim == 1 and lens.shape[0] == batch
    assert lens.amax() <= seq_len

    pad_lens = seq_len - lens
    max_pad_len = pad_lens.amax()

    batch_arange = torch.arange(batch, device=device, dtype=torch.long)[..., None]
    prompt_len_arange = torch.arange(seq_len, device=device, dtype=torch.long)

    t = F.pad(t, (max_pad_len, 0), value=0)
    offset = max_pad_len - pad_lens

    aligned = t[batch_arange, prompt_len_arange + offset[..., None]]
    return aligned


# nucleus


def top_p(logits, thres=0.9):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

    sorted_indices_to_remove = cum_probs > thres
    sorted_indices_to_remove = F.pad(sorted_indices_to_remove, (1, -1), value=False)

    sorted_logits[sorted_indices_to_remove] = float("-inf")
    return sorted_logits.scatter(1, sorted_indices, sorted_logits)


# topk


def top_k(logits, frac_num_tokens=0.1, k=None):
    num_tokens = logits.shape[-1]

    k = default(k, ceil(frac_num_tokens * num_tokens))
    k = min(k, num_tokens)

    val, ind = torch.topk(logits, k)
    probs = torch.full_like(logits, float("-inf"))
    probs.scatter_(1, ind, val)
    return probs


# top_a


def top_a(logits, min_p_pow=2.0, min_p_ratio=0.02):
    probs = F.softmax(logits, dim=-1)
    max_probs = torch.amax(probs, dim=-1, keepdim=True)
    limit = torch.pow(max_probs, min_p_pow) * min_p_ratio
    return torch.where(probs < limit, float("-inf"), logits)


# contrastive decoding function


def contrastive_decode_fn(expert_logits, amateur_logits, alpha=0.1, beta=0.5):
    """
    Appendix A Algorithm 2
    https://arxiv.org/abs/2309.09117
    """

    cutoff = log(alpha) + expert_logits.amax(dim=-1, keepdim=True)
    diffs = (1 + beta) * expert_logits - beta * amateur_logits
    contrastive_decode_logits = diffs.masked_fill(
        expert_logits < cutoff, -torch.finfo(expert_logits.dtype).max
    )
    return contrastive_decode_logits


# autoregressive wrapper class


class AutoregressiveWrapper(Module):
    def __init__(
        self, net, ignore_index=-100, pad_value=0, mask_prob=0.0, add_attn_z_loss=False
    ):
        super().__init__()
        self.pad_value = pad_value
        self.ignore_index = ignore_index

        self.net = net
        self.max_seq_len = net.max_seq_len

        # paper shows masking (MLM) in conjunction with autoregressive
        # decoder-only training leads to big improvements
        # https://arxiv.org/abs/2210.13432
        assert mask_prob < 1.0
        self.mask_prob = mask_prob

        # whether to add router z-loss
        self.add_attn_z_loss = add_attn_z_loss

    @torch.no_grad()
    @eval_decorator
    def generate(
        self,
        prompts,
        seq_len,
        eos_token=None,
        temperature=1.0,
        prompt_lens: Optional[Tensor] = None,
        filter_logits_fn: Callable = top_k,
        restrict_to_max_seq_len=True,
        amateur_model: Optional[Union[Module, Tuple[Module]]] = None,
        filter_kwargs: dict = dict(),
        contrastive_decode_kwargs: Union[dict, Tuple[dict]] = dict(beta=0.5, alpha=0.1),
        cache_kv=True,
        **kwargs
    ):
        max_seq_len, device = self.max_seq_len, prompts.device

        prompts, ps = pack([prompts], "* n")

        b, t = prompts.shape

        # handle variable lengthed prompts (prefixes)

        seq_start_pos = None
        if exists(prompt_lens):
            prompts = align_right(prompts, prompt_lens, pad_id=self.pad_value)
            seq_start_pos = t - prompt_lens

        # output from which sampled tokens appended to

        out = prompts

        # kv caches

        cache = None

        # if doing contrastive decoding, turn off filter automatically

        if exists(amateur_model):
            amateur_model = cast_tuple(amateur_model)
            contrastive_decode_kwargs = cast_tuple(contrastive_decode_kwargs)

            assert len(amateur_model) == len(contrastive_decode_kwargs)

            amateur_caches = [None] * len(amateur_model)
            filter_logits_fn = identity

            for i, module in enumerate(amateur_model):
                if isinstance(module, AutoregressiveWrapper):
                    amateur_model[i] = module.net

                module.eval()

        # sampling up to seq_len

        for _ in range(seq_len):
            if restrict_to_max_seq_len:
                x = out[:, -max_seq_len:]

                if exists(cache):
                    for inter in cache.attn_intermediates:
                        inter.cached_kv = [
                            t[..., -(max_seq_len - 1) :, :] for t in inter.cached_kv
                        ]

            logits, new_cache = self.net(
                x,
                return_intermediates=True,
                cache=cache,
                seq_start_pos=seq_start_pos,
                **kwargs
            )

            if cache_kv and self.net.can_cache_kv:
                cache = new_cache

            logits = logits[:, -1]

            # handle contrastive decoding, Li et al.
            # https://arxiv.org/abs/2210.15097

            if exists(amateur_model):
                for i, (
                    amateur,
                    amateur_cache,
                    amateur_contrastive_decode_kwargs,
                ) in enumerate(
                    zip(amateur_model, amateur_caches, contrastive_decode_kwargs)
                ):
                    amateur_logits, next_amateur_cache = amateur(
                        x,
                        return_intermediates=True,
                        cache=amateur_cache,
                        seq_start_pos=seq_start_pos,
                        **kwargs
                    )

                    amateur_logits = amateur_logits[:, -1]

                    assert (
                        amateur_logits.shape == logits.shape
                    ), "logits dimension are not the same between amateur and expert model"
                    logits = contrastive_decode_fn(
                        logits, amateur_logits, **amateur_contrastive_decode_kwargs
                    )

                    if cache_kv and amateur.can_cache_kv:
                        amateur_caches[i] = next_amateur_cache

            # filter by top_k, top_p (nucleus), top_a, or custom

            filtered_logits = filter_logits_fn(logits, **filter_kwargs)

            probs = F.softmax(filtered_logits / temperature, dim=-1)

            sample = torch.multinomial(probs, 1)

            out = torch.cat((out, sample), dim=-1)

            if exists(eos_token):
                is_eos_tokens = out == eos_token

                if is_eos_tokens.any(dim=-1).all():
                    # mask out everything after the eos tokens
                    shifted_is_eos_tokens = F.pad(is_eos_tokens, (1, -1))
                    mask = shifted_is_eos_tokens.float().cumsum(dim=-1) >= 1
                    out = out.masked_fill(mask, self.pad_value)
                    break

        out = out[:, t:]

        (out,) = unpack(out, ps, "* n")

        return out

    def forward(self, x, **kwargs):
        seq, ignore_index, add_attn_z_loss = (
            x.shape[1],
            self.ignore_index,
            self.add_attn_z_loss,
        )

        inp, target = x[:, :-1], x[:, 1:]
        inp = torch.where(inp == ignore_index, self.pad_value, inp)

        if self.mask_prob > 0.0:
            rand = torch.randn(inp.shape, device=x.device)
            rand[:, 0] = -torch.finfo(
                rand.dtype
            ).max  # first token should not be masked out
            num_mask = min(int(seq * self.mask_prob), seq - 1)
            indices = rand.topk(num_mask, dim=-1).indices
            mask = ~torch.zeros_like(inp).scatter(1, indices, 1.0).bool()
            kwargs.update(self_attn_kv_mask=mask)

        logits, cache = self.net(
            inp, return_intermediates=True, return_attn_z_loss=add_attn_z_loss, **kwargs
        )

        loss = F.cross_entropy(
            rearrange(logits, "b n c -> b c n"), target, ignore_index=ignore_index
        )

        if add_attn_z_loss:
            loss = loss + cache.attn_z_loss

        return loss

swarms_torch/autoregressive.py
kyegomez/swarms-pytorch
import torch
from torch import nn


class TransformerCell(nn.Module):
    def __init__(
        self,
        input_dim,
        nhead,
        num_layers=1,
        neighborhood_size=3,
    ):
        super(TransformerCell, self).__init__()
        self.transformer = nn.Transformer(
            input_dim, nhead=nhead, num_encoder_layers=num_layers
        )
        self.neighborhood_size = neighborhood_size

    def forward(self, x, neigbors):
        x = self.transformer(x, x)

        out = torch.cat([x] + neigbors, dim=0)

        return out


class CellularSwarm(nn.Module):
    """
    CellularSwarm

    Architecture:
        - Input -> TransformerCell -> TransformerCell -> ... -> Output

    Overview:
    CellularSwarm is a cellular neural network that uses a transformer cell
    to process the input.

    Args:
        cell_count (int): Number of transformer cells
        input_dim (int): Input dimension
        nhead (int): Number of heads in the transformer cell
        time_steps (int): Number of time steps to run the network

    Returns:
        torch.Tensor: Output tensor

    Usage:
        >>> x = torch.randn(10, 32, 512)
        >>> model = CellularSwarm(cell_count=5, input_dim=512, nhead=8)
        >>> output = model(x)
        >>> print(output)


    """

    def __init__(self, cell_count, input_dim, nhead, time_steps=4):
        super(CellularSwarm, self).__init__()
        self.cells = nn.ModuleList(
            [TransformerCell(input_dim, nhead) for _ in range(cell_count)]
        )
        self.time_steps = time_steps

    def forward(self, x):
        for _ in range(self.time_steps):
            for i, cell in enumerate(self.cells):
                # get neighboring cells states
                start_idx = max(0, i - cell.neighborhood_size)

                end_idx = min(len(self.cells), i + cell.neighborhood_size + 1)

                neighbors = [
                    self.cells[j].transformer(x, x)
                    for j in range(start_idx, end_idx)
                    if j != i
                ]

                x = cell(x, neighbors)
        return x

swarms_torch/cellular_transformer.py
kyegomez/swarms-pytorch
import torch
from torch import nn
from torch.nn import CrossEntropyLoss, Transformer
from torch.optim import Adam
from torch.optim.lr_scheduler import ReduceLROnPlateau


class Fish(nn.Module):
    """
    A fish is a transformer model with a negative loss as food.

    Parameters
    ----------
    dim : int
        The number of expected features in the input (required).
    heads : int
        The number of heads in the multiheadattention models (required).
    depth : int
        The number of sub-encoder-layers in the encoder (required).

    Attributes

    model : torch.nn.Transformer
        The transformer model.
    food : float
        The fish's food, which is the negative loss of the model.

    Methods
    =======
    train(src, tgt, labels)
        Train the model with the given source, target, and labels.


    Usage:
    >>> fish = Fish(512, 8, 6)
    >>> fish.train(src, tgt, labels)
    >>> fish.food
    -0.123456789



    Example2
    # Create random source and target sequences
    src = torch.randn(10, 32, 512)
    tgt = torch.randn(10, 32, 512)

    # Create random labels
    labels = torch.randint(0, 512, (10, 32))

    # Create a fish and train it on the random data
    fish = Fish(512, 8, 6)
    fish.train(src, tgt, labels)
    print(fish.food)  # Print the fish's food

    # Create a fish school and optimize it on the random data
    school = FishSchool(10, 512, 8, 6, 100)
    school.forward(src, tgt, labels)
    print(school.fish[0].food)  # Print the first fish's food


    """

    def __init__(
        self,
        dim,
        heads,
        depth,
        dynamic_learning_rate=False,
        early_stopping=False,
        complexity_regularization=False,
        max_patience=None,
        alpha=0.1,
    ):
        super().__init__()
        self.model = Transformer(
            d_model=dim, nhead=heads, num_encoder_layers=depth, num_decoder_layers=depth
        )
        self.optimizer = Adam(self.parameters())
        self.scheduler = ReduceLROnPlateau(self.optimizer, "min")

        self.complexity_regularization = complexity_regularization
        self.dynamic_learning_rate = dynamic_learning_rate
        self.early_stopping = early_stopping
        self.food = 0
        self.best_food = float("inf")
        self.patience = 0
        self.max_patience = max_patience
        self.alpha = alpha

    def train(self, src, tgt, labels):
        """Trains the fish school"""
        self.model.train()

        self.optimizer.zero_grad()

        outputs = self.model(src, tgt)

        # cross entropy loss
        loss = CrossEntropyLoss()(outputs.view(-1, outputs.size(-1)), labels.view(-1))

        # complexity regularization by adding the sum of the squares of the
        # weights
        if self.complexity_regularization:
            # complexity regularization
            loss += self.alpha * sum(p.pow(2.0).sum() for p in self.model.parameters())

        # backpropagation
        loss.backward()

        # dynamic learning rate
        if self.dynamic_learning_rate:
            self.scheduler.step(loss)
        else:
            forward = self.optimizer
            forward.step()

        # # use negative loss as food
        self.food = -loss.item()

        # early stopping if the fish is not improving
        if self.early_stopping:
            if loss < self.best_food:
                self.best_food = loss
                self.patience = 0
            else:
                self.patience += 1

    def forward(self, src, tgt):
        """Forward pass of the fish school"""
        return self.model(src, tgt)

    def early_stopping(self, food):
        """
        Early stopping if the fish is not improving.
        """
        if food < self.best_food:
            self.best_food = food
            self.patience = 0
        else:
            self.patience += 1
            if self.patience > 5:
                return True
        return False

    def generate(self, src, tgt):
        """
        Generate a sequence using the fish's model.
        """
        return self.model.generate(src, tgt)

    def save(self, path):
        """
        Save the fish's model.
        """
        torch.save(self.model.state_dict(), path)


class FishSchool(nn.Module):
    """
    Fish School is a collection of fish.

    Parameters
    ----------
    num_fish : int
        The number of fish in the school.
    dim : int
        The number of expected features in the input (required).
    heads : int
        The number of heads in the multiheadattention models (required).
    depth : int
        The number of sub-encoder-layers in the encoder (required).
    num_iter : int
        The number of iterations to train the fish school.


    Usage:
    >>> school = FishSchool(10, 512, 8, 6, 100)
    >>> school.train(src, tgt, labels)
    >>> school.fish[0].food

    """

    def __init__(
        self,
        num_fish,
        dim,
        heads,
        depth,
        num_iter,
        num_top_fish=None,
        complex_school=False,
        max_seq_len=8192,
    ):
        super().__init__()
        self.fish = [Fish(dim, heads, depth) for _ in range(num_fish)]
        self.num_iter = num_iter
        self.num_top_fish = num_top_fish
        self.complex_school = complex_school
        self.max_seq_len = max_seq_len

    def forward(self, src, tgt, labels):
        for _ in range(self.num_iter):
            total_food = 0
            for fish in self.fish:
                fish.train(src, tgt, labels)
                total_food += fish.food
            # adjust schoold behavior on total food
            avg_food = total_food / len(self.fish)

            # complex school behavior => fish with lower food learn from fish
            # with higher food
            if self.complex_school:
                for fish in self.fish:
                    neighbor = self.fish[torch.randint(0, len(self.fish), (1,)).item()]
                    if neighbor.food > fish.food:
                        fish.model.load_state_dict(neighbor.model.state_dict())

            # adjust schoold behavior on average food
            for fish in self.fish:
                if fish.food < avg_food:
                    # transformer weights from the best performing fish
                    best_fish = max(self.fish, key=lambda f: f.food)
                    fish.model.load_state_dict(best_fish.model.state_dict())

    def generate(self, src, tgt):
        """
        Generate a sequence using the fish school's model.
        """
        return self.fish[0].generate(src, tgt)

    def predict(self, src, tgt):
        """
        Ensemble learning => enseble prediction of top performing models

        averages outputs of the top peforming models
        """
        top_fish = sorted(self.fish, key=lambda f: f.food, reverse=True)[
            : self.num_top_fish
        ]

        self.model.eval()

        with torch.no_grad():
            outputs = torch.stack([fish.model(src, tgt) for fish in top_fish])

        return sum(outputs) / len(outputs)

    def save(self, path):
        """
        Save the fish school's models.
        """
        for i, fish in enumerate(self.fish):
            fish.save(path + f"fish_{i}.pt")

    def load(self, path):
        """
        Load the fish school's models.
        """
        for i, fish in enumerate(self.fish):
            fish.model.load_state_dict(torch.load(path + f"fish_{i}.pt"))

    def early_stopping(self):
        """
        Early stopping if the fish school is not improving.
        """
        for fish in self.fish:
            if fish.early_stopping(fish.food):
                return True
        return False

    def dynamic_learning_rate(self):
        """
        Dynamic learning rate for the fish school.
        """
        for fish in self.fish:
            fish.dynamic_learning_rate = True

    def complexity_regularization(self):
        """
        Complexity regularization for the fish school.
        """
        for fish in self.fish:
            fish.complexity_regularization = True

    def reset(self):
        """
        Reset the fish school's food.
        """
        for fish in self.fish:
            fish.food = 0

    def __getitem__(self, index):
        """Get the fish at the given index"""
        return self.fish[index]

    def __len__(self):
        """Get the number of fish in the school"""
        return len(self.fish)

    def __iter__(self):
        """Iterate over the fish in the school"""
        return iter(self.fish)

    def __next__(self):
        """Get the next fish in the school"""
        return next(self.fish)

    def __str__(self):
        """Get the string representation of the fish school"""
        return str(self.fish)

swarms_torch/fish_school.py
kyegomez/swarms-pytorch
import torch
import torch.nn as nn


class GraphCellularAutomata(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(GraphCellularAutomata, self).__init__()

        self.mlp = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim),
        )

    def forward(self, x):
        return self.mlp(x)


class ReplicationModel(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(ReplicationModel, self).__init__()

        self.mlp = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid(),  # for binary classification
        )

    def forward(self, x):
        return self.mlp(x)


class WeightUpdateModel(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(WeightUpdateModel, self).__init__()

        self.mlp = nn.Sequential(
            nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1)
        )

    def forward(self, x):
        return self.mlp(x)


class NDP(nn.Module):
    def __init__(self, embedding_dim, hidden_dim):
        super(NDP, self).__init__()

        self.gc_automata = GraphCellularAutomata(
            embedding_dim, hidden_dim, embedding_dim
        )
        self.replication_model = ReplicationModel(embedding_dim, hidden_dim)
        self.weight_update_model = WeightUpdateModel(2 * embedding_dim, hidden_dim)

    def forward(self, node_embeddings, adjacency_matrix):
        # Update node embeddings using Graph Cellular Automata
        updated_embeddings = self.gc_automata(node_embeddings)

        # Check which nodes need to replicate
        replication_decisions = self.replication_model(updated_embeddings)

        # Weight update (assuming weighted network)
        num_nodes = node_embeddings.shape[0]
        edge_weights = torch.zeros((num_nodes, num_nodes))

        for i in range(num_nodes):
            for j in range(num_nodes):
                combined_embedding = torch.cat(
                    (updated_embeddings[i], updated_embeddings[j])
                )

                edge_weights[i, j] = self.weight_update_model(combined_embedding)

        return updated_embeddings, replication_decisions, edge_weights


# # Usage examples
# embedding_dim = 16
# hidden_dim = 32
# node_embeddings = torch.rand((10, embedding_dim))  # For 10 nodes
# adjacency_matrix = torch.rand((10, 10))  # Dummy adjacency matrix for 10
# nodes

# model = NDP(embedding_dim, hidden_dim)
# updated_embeddings, replication_decisions, edge_weights = model(
#     node_embeddings, adjacency_matrix
# )

# print(updated_embeddings.shape)
# print(replication_decisions.shape)
# print(edge_weights.shape)

swarms_torch/graph_cellular_automa.py
kyegomez/swarms-pytorch
import torch
import torch.nn as nn
import torch.optim as optim
import gym


class MAgent:
    class Agent(nn.Module):
        def __init__(self, input_dim, output_dim):
            super().__init__()
            self.policy = nn.Sequential(
                nn.Linear(input_dim, 128),
                nn.ReLU(),
                nn.Linear(128, output_dim),
                nn.Softmax(dim=-1),
            )

        def forward(self, state):
            return self.policy(state)

    class MultiGymEnvironment:
        def __init__(self, env_name, num_agents):
            self.envs = [gym.make(env_name) for _ in range(num_agents)]
            self.agents = [
                MAgent.Agent(
                    self.envs[0].observation_space.shape[0], self.envs[0].action_space.n
                )
                for _ in range(num_agents)
            ]
            self.optimizers = [optim.Adam(agent.parameters()) for agent in self.agents]

        def step(self, agent_actions):
            rewards = []
            for env, action in zip(self.envs, agent_actions):
                _, reward, _, _ = env.step(action)
                rewards.append(reward)
            return rewards

        def get_states(self):
            states = [env.reset() for env in self.envs]
            return states

        def train(self, epochs=1000):
            for epoch in range(epochs):
                states = self.get_states()
                actions = [
                    torch.argmax(agent(torch.FloatTensor(state))).item()
                    for agent, state in zip(self.agents, states)
                ]
                rewards = self.step(actions)

                for agent, optimizer, reward in zip(
                    self.agents, self.optimizers, rewards
                ):
                    loss = (
                        -torch.log(agent(torch.FloatTensor(states))) * reward
                    )  # Example loss function
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()


# Example usage:
num_agents = 5
simulation = MAgent.MultiGymEnvironment("CartPole-v1", num_agents)
simulation.train()

swarms_torch/ma_agent.py
kyegomez/swarms-pytorch
import torch
import random
import string


class MultiSwarmPSO:
    """
    Multi-Swarm PSO Algorithm

    Parameters
    ----------
    target_string : str
        The target string to be generated
    num_sub_swarms : int
        The number of sub-swarms
    num_particles_per_swarm : int
        The number of particles per sub-swarm
    max_iterations : int
        The maximum number of iterations to run the algorithm

    Attributes
    ----------
    target_string : str
        The target string to be generated
    num_sub_swarms : int
        The number of sub-swarms
    num_particles_per_swarm : int
        The number of particles per sub-swarm
    num_dimensions : int
        The number of dimensions in the search space
    max_iterations : int
        The maximum number of iterations to run the algorithm

    Methods
    -------
    generate_random_string()
        Generates a random string of length num_dimensions
    fitness_function(position)
        Calculates the fitness of a given position
    diversification_method(sub_swarms)
        Adds a new sub-swarm if the number of sub-swarms is less than the maximum
    optimize()
        Runs the Multi-Swarm PSO algorithm

    References
    ----------
    .. [1] https://www.researchgate.net/publication/221172800_Multi-swarm_Particle_Swarm_Optimization


    Usage:
    ------
    target_string = "hello world"
    multi_swarm = MultiSwarm(target_string)
    multi_swarm.optimize()



    """

    def __init__(
        self,
        target_string,
        num_sub_swarms=5,
        num_particles_per_swarm=20,
        max_iterations=100,
    ):
        self.target_string = target_string
        self.num_sub_swarms = num_sub_swarms
        self.num_particles_per_swarm = num_particles_per_swarm
        self.num_dimensions = len(target_string)
        self.max_iterations = max_iterations

    def generate_random_string(self):
        """
        Generates a random string of length num_dimensions

        """
        return "".join(
            random.choice(string.ascii_lowercase + " ")
            for _ in range(self.num_dimensions)
        )

    def fitness_function(self, position):
        """Fitness function to be maximized"""
        fitness = sum(a == b for a, b in zip(position, self.target_string))
        return fitness

    def diversification_method(self, sub_swarms):
        """Diversification method to add a new sub-swarm if the number of sub-swarms is less than the maximum"""
        if len(sub_swarms) < self.num_sub_swarms:
            new_sub_swarm = [
                self.generate_random_string()
                for _ in range(self.num_particles_per_swarm)
            ]
            sub_swarms.append(new_sub_swarm)

    def optimize(self):
        """Optimizes the fitness function"""
        sub_swarms = [
            [self.generate_random_string() for _ in range(self.num_particles_per_swarm)]
            for _ in range(self.num_sub_swarms)
        ]

        for iteration in range(self.max_iterations):
            for sub_swarm in sub_swarms:
                for particle in sub_swarm:
                    fitness = self.fitness_function(particle)
                    if fitness > 0:
                        index_to_change = random.randint(0, self.num_dimensions - 1)
                        new_char = random.choice(string.ascii_lowercase + " ")
                        new_position = list(particle)
                        new_position[index_to_change] = new_char
                        new_position = "".join(new_position)
                        particle = new_position

            self.diversification_method(sub_swarms)

            global_best_fitness = max(
                self.fitness_function(particle)
                for sub_swarm in sub_swarms
                for particle in sub_swarm
            )
            global_best_position = [
                particle
                for sub_swarm in sub_swarms
                for particle in sub_swarm
                if self.fitness_function(particle) == global_best_fitness
            ][0]
            print(
                f"Iteration {iteration}: Global Best Fitness = {global_best_fitness}, Global Best Position = {global_best_position}"
            )

        global_best_fitness = max(
            self.fitness_function(particle)
            for sub_swarm in sub_swarms
            for particle in sub_swarm
        )
        global_best_position = [
            particle
            for sub_swarm in sub_swarms
            for particle in sub_swarm
            if self.fitness_function(particle) == global_best_fitness
        ][0]
        print(
            f"Final Result: Global Best Fitness = {global_best_fitness}, Global Best Position = {global_best_position}"
        )


# Example usage
if __name__ == "__main__":
    target_string = "hello world"
    multi_swarm = MultiSwarmPSO(target_string)
    multi_swarm.optimize()

swarms_torch/multi_swarm_pso.py
kyegomez/swarms-pytorch
import torch


class Particle:
    def __init__(self, dim, minx, maxx):
        self.position = torch.rand(dim) * (maxx - minx) + minx
        self.velocity = torch.rand(dim) * (maxx - minx) + minx
        self.best_position = self.position.clone()
        self.best_score = float("inf")

    def update_velocity(self, global_best, w=0.7, c1=1.5, c2=1.5):
        r1 = torch.rand(self.position.size())
        r2 = torch.rand(self.position.size())
        self.velocity = (
            w * self.velocity
            + c1 * r1 * (self.best_position - self.position)
            + c2 * r2 * (global_best - self.position)
        )

    def update_position(self, minx, maxx):
        self.position += self.velocity
        self.position = torch.clamp(self.position, minx, maxx)


class Swarm:
    def __init__(self, num_particles, dim, minx, maxx):
        self.particles = [Particle(dim, minx, maxx) for _ in range(num_particles)]
        self.global_best = None
        self.global_best_score = float("inf")

    def update_global_best(self):
        for particle in self.particles:
            if particle.best_score < self.global_best_score:
                self.global_best = particle.best_position.clone()
                self.global_best_score = particle.best_score

    def move_particles(self, minx, maxx):
        for particle in self.particles:
            particle.update_velocity(self.global_best)
            particle.update_position(minx, maxx)


class MultiSwarm:
    def __init__(self, num_swarms, num_particles, dim, minx, maxx):
        self.swarms = [Swarm(num_particles, dim, minx, maxx) for _ in range(num_swarms)]
        self.minx = minx
        self.maxx = maxx

    def optimize(self, func, max_iter):
        for _ in range(max_iter):
            for swarm in self.swarms:
                swarm.update_global_best()
                swarm.move_particles(self.minx, self.maxx)
        best_swarm = min(self.swarms, key=lambda s: s.global_best_score)
        return best_swarm.global_best, best_swarm.global_best_score


def rosenbrock(x, a=1, b=100):
    return (a - x[0]) ** 2 + b * (x[1] - x[0] ** 2) ** 2


num_swarms = 5
num_particles = 20
dim = 2
minx = -5
maxx = 5
max_iter = 100

multi_swarm = MultiSwarm(num_swarms, num_particles, dim, minx, maxx)

best_position, best_score = multi_swarm.optimize(rosenbrock, max_iter)

print(f"Best position: {best_position}")
print(f"Best score: {best_score}")

swarms_torch/multi_swarm_pso2.py
kyegomez/swarms-pytorch
# import torch
# import torch.nn as nn
# from copy import deepcopy
# from swarms_torch.transformer_pso import Particle, TransformerParticleSwarmOptimization


# class MultiSwarm(nn.Module):
#     def __init__(
#         self,
#         num_swarms,
#         *args,
#         **kwargs
#     ):
#         #create multiple instances of the transformerparticle swarm optimization
#         self.swarms = [TransformerParticleSwarmOptimization(*args, **kwargs) for _ in range(num_swarms)]
#         self.num_swarms = num_swarms

#     def optimize(self, iterations):
#         for _ in range(iterations):
#             #update each swarm
#             for swarm in self.swarms:
#                 swarm.update()

#             #apply diversification strategy
#             self.diversification_method()
        
#     def diversification_strategy(self):
#         for i in range(self.num_swarms):
#             for j in range(i + 1, self.num_swarms):
#                 if self.is_collided(self.swarms[i].global_best, self.swarms[j].global_best):
#                     #handle collision by launching a new swarm or re init one of the swarms
#                     self.handle_collision(i, j)
    
#     def is_collided(self, global_best_1, global_best_2):
#         #Check if difference between the global bests or 2 swarms is below a threshold
#         diff = sum((global_best_1[key] - global_best_2[key]).abs().sum() for key in global_best_1.keys())
#         COLLISION_THRESHOLD = 0.1

#         return diff < COLLISION_THRESHOLD
    
#     def handle_collision(self, idx1, idx2):
#         #for simplicity re init 2nd swarm
#         self.swarms[idx2] = TransformerParticleSwarmOptimization(*self.swarms[idx2].model_args, **self.swarms[idx2].kwargs)



# import torch
# from torch.utils.data import DataLoader, TensorDataset

# # Generate random data
# num_samples = 1000
# input_dim = 50  # Length of input sequence
# num_classes = 2

# inputs = torch.randint(0, 1000, (num_samples, input_dim))
# targets = torch.randint(0, num_classes, (num_samples,))

# dataset = TensorDataset(inputs, targets)
# data_loader = DataLoader(dataset, batch_size=32, shuffle=True)

# # Define hyperparameters and model arguments
# model_args = (1000, 512, 8, 6, 2)  # (input_dim, d_model, nhead, num_layers, output_dim)
# optimizer_args = {
#     "model_constructor": Particle,
#     "model_args": model_args,
#     "device": "cpu",
#     "criterion": torch.nn.CrossEntropyLoss(),
#     "data_loader": data_loader
# }

# # Create MultiSwarmOptimizer
# num_swarms = 3
# mso = MultiSwarm(num_swarms, **optimizer_args)

# # Optimize
# mso.optimize(iterations=10)

# # Get the best model from the best-performing swarm
# best_swarm = max(mso.swarms, key=lambda s: s.compute_fitness(s.global_best))
# best_model = best_swarm.get_best_model()

import torch
import torch.nn as nn
from copy import deepcopy

class Particle(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(Particle, self).__init__()
        self.transformer = nn.Transformer(input_dim, hidden_dim)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.transformer(x)
        x = self.fc(x)
        return x

class MultiSwarmOptimizer:
    def __init__(self, particle, num_particles, num_subswarms, fitness_func, bounds, num_epochs):
        self.particle = particle
        self.num_particles = num_particles
        self.num_subswarms = num_subswarms
        self.fitness_func = fitness_func
        self.bounds = bounds
        self.num_epochs = num_epochs

        self.subswarms = []
        for _ in range(num_subswarms):
            self.subswarms.append([deepcopy(particle) for _ in range(num_particles)])

    def optimize(self):
        for epoch in range(self.num_epochs):
            for subswarm in self.subswarms:
                for particle in subswarm:
                    fitness = self.fitness_func(particle)
                    if fitness > particle.best_fitness:
                        particle.best_fitness = fitness
                        particle.best_position = deepcopy(particle.position)

                best_particle = max(subswarm, key=lambda p: p.best_fitness)
                for particle in subswarm:
                    particle.velocity = (particle.velocity + 
                                         0.5 * (particle.best_position - particle.position) +
                                         0.5 * (best_particle.best_position - particle.position))
                    particle.position = particle.position + particle.velocity
                    particle.position = torch.clamp(particle.position, *self.bounds)

            best_subswarm = max(self.subswarms, key=lambda s: max(p.best_fitness for p in s))
            best_particle = max(best_subswarm, key=lambda p: p.best_fitness)
            print(f'Epoch {epoch+1}/{self.num_epochs}, Best Fitness: {best_particle.best_fitness}')

        best_subswarm = max(self.subswarms, key=lambda s: max(p.best_fitness for p in s))
        best_particle = max(best_subswarm, key=lambda p: p.best_fitness)
        return best_particle
    
    def get_best_model(self):
        return self.particle
    


import torch
import torch.nn as nn
from random import random

# Define the fitness function
def fitness_func(particle):
    # This is a dummy fitness function. Replace it with your own.
    return random()

# Define the bounds for the particle positions
bounds = (-1.0, 1.0)

# Define the number of particles, sub-swarms, and epochs
num_particles = 10
num_subswarms = 5
num_epochs = 100

# Define the dimensions for the transformer model
input_dim = 10
hidden_dim = 20
output_dim = 2

# Create a particle (transformer model)
particle = Particle(input_dim, hidden_dim, output_dim)

# Create the multi-swarm optimizer
optimizer = MultiSwarmOptimizer(particle, num_particles, num_subswarms, fitness_func, bounds, num_epochs)

# Run the optimization
best_particle = optimizer.optimize()

# The best_particle is the model with the highest fitness score
print(best_particle)
swarms_torch/multi_swarm_pso_transformer.py
kyegomez/swarms-pytorch
"""

Cellular neural network

Architecture:
    - Input -> Linear -> ReLU -> Linear -> ReLU -> Output
    - Neuron states are updated after each synapse
    - Softmax is applied after each synapse
    - Layer normalization is applied after each synapse

"""

import torch
from torch import nn


class TransformerLayer(nn.Module):
    """
    Transformer Layer

    Architecture:
        - Input -> Linear -> ReLU -> Linear -> ReLU -> Output

    Args:
        input_dim (int): Input dimension
        output_dim (int): Output dimension

    Returns:
        torch.Tensor: Output tensor

    Usage

    """

    def __init__(
        self,
        input_dim,
        output_dim,
        nhead: int,
    ):
        super(TransformerLayer, self).__init__()
        self.transformer = nn.TransformerEncoderLayer(
            d_model=input_dim,
            nhead=nhead,
        )
        self.fc = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return self.fc(self.transformer(x))


class Neuron(nn.Module):
    def __init__(self, num_states):
        super(Neuron, self).__init__()
        self.states = nn.Parameter(torch.randn(num_states))


class SynapseTransformer(nn.Module):
    def __init__(self, input_dim, output_dim, nhead: int):
        super(SynapseTransformer, self).__init__()
        self.transformer = TransformerLayer(input_dim, output_dim, nhead)

    def forward(self, x):
        return self.transformer(x)


class NNTransformer(nn.Module):
    """
    Neural Network NNTransformer

    Args:
        neuron_count (int): Number of neurons
        num_states (int): Number of states
        input_dim (int): Input dimension
        output_dim (int): Output dimension
        nhead (int): Number of heads in transformer layer

    Returns:
        torch.Tensor: Output tensor

    Architecture:
        - Input -> Linear -> ReLU -> Linear -> ReLU -> Output
        - Neuron states are updated after each synapse
        - Softmax is applied after each synapse
        - Layer normalization is applied after each synapse

    Usage:
        network = CellularNN(5, 10, 10, 10, 2)
        output = network(torch.randn(1, 10))
        print(output)


    Training:
    network = NNTransformer(5, 10, 10, 10, 2)
    output = network(torch.randn(1, 10))
    print(output)


    # Test the network
    import torch
    import torch.optim as optim
    import torch.nn.functional as F

    # Random dataset
    batch_size = 64
    input_size = 10
    output_size = 10

    x = torch.randn(batch_size, input_size)  # Random inputs
    y = torch.randn(batch_size, output_size)  # Random targets

    # Hyperparameters
    neuron_count = 5
    num_states = 10
    input_dim = input_size
    output_dim = output_size
    n_head = 2

    # Initialize the network
    network = CellularNN(neuron_count, num_states, input_dim, output_dim, n_head)

    # Define the loss function and optimizer
    criterion = nn.MSELoss()
    optimizer = optim.Adam(network.parameters(), lr=0.001)

    # Training loop
    epochs = 1000
    for epoch in range(epochs):
        # Forward pass
        outputs = network(x)

        # Compute loss
        loss = criterion(outputs, y)

        # Backward pass and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Print loss every 100 epochs
        if (epoch+1) % 100 == 0:
            print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

    # Test the network with a new random input
    test_input = torch.randn(1, input_size)
    test_output = network(test_input)
    print(test_output)


    """

    def __init__(self, neuron_count, num_states, input_dim, output_dim, nhead):
        super(NNTransformer, self).__init__()

        # Initialize neurons and synapses
        self.neurons = nn.ModuleList([Neuron(num_states) for _ in range(neuron_count)])
        self.synapses = nn.ModuleList(
            [
                SynapseTransformer(input_dim, output_dim, nhead)
                for _ in range(neuron_count)
            ]
        )

        self.norm = nn.LayerNorm(output_dim)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        for neuron, synapse in zip(self.neurons[:-1], self.synapses):
            # norm before synapse
            x = self.norm(x)

            # synapse
            x = synapse(x)

            # softmax after synapse
            x = self.softmax(x)

            neuron.states.data = x
        return self.neurons[-1].states

swarms_torch/neuronal_transformer.py
kyegomez/swarms-pytorch
import torch


class ParticleSwarmOptimization:
    """
    Particle Swarm Optimization
    Overview: https://en.wikipedia.org/wiki/Particle_swarm_optimization

    How does it work?
    1. Initialize particles with random positions and velocities
    2. For each particle, compute the fitness value
    3. Update the personal best and global best
    4. Update the velocity and position of each particle
    5. Repeat step 2 to 4 until the maximum number of iterations is reached



    Parameters
    ----------
    goal: str
        The goal string to be optimized
    n_particles: int
        Number of particles
    inertia: float
        Inertia weight
    personal_best_weight: float
        Personal best weight
    global_best_weight: float
        Global best weight

    Usage
    -----
    pso = ParticleSwarmOptimization(goal="Attention is all you need", n_particles=100)
    pso.optimize(iterations=1000)


    Future Improvements
    -------------------
    1. Add a stopping criterion
    2. Add a callback function to track the progress
    3. Add a function to plot the fitness value
    4. Add a function to plot the particles
    5. Add a function to plot the velocity
    6. Add a function to plot the position
    7. Add a function to plot the personal best
    8. Add a function to plot the global best
    9. Add a function to plot the personal best weight
    10. Add a function to plot the global best weight



    """

    def __init__(
        self,
        goal: str = None,
        n_particles: int = 100,
        inertia: float = 0.5,
        personal_best_weight: float = 1.5,
        global_best_weight: float = 1.5,
        dim: int = 1,
    ):
        self.goal = torch.tensor([ord(c) for c in goal])
        self.n_particles = n_particles
        self.inertia = inertia
        self.personal_best_weight = personal_best_weight
        self.global_best_weight = global_best_weight

        self.particles = torch.randint(0, 255, (n_particles, len(goal)))
        self.velocities = torch.zeros((n_particles, len(goal)))

        self.personal_best = self.particles.clone()
        self.global_best = self.particles[0].clone()

    def compute_fitness(
        self,
        particle,
    ):
        return 1.0 / (1.0 + torch.norm((particle - self.goal).float()))

    def update(
        self,
    ):
        """Update the particles"""
        for i in range(self.n_particles):
            fitness = self.compute_fitness(
                self.particles[i],
            )

            personal_best_fitness = self.compute_fitness(
                self.personal_best[i],
            )

            if fitness > personal_best_fitness:
                self.personal_best[i] = self.particles[i]

            global_best_fitness = self.compute_fitness(self.global_best)

            if fitness > global_best_fitness:
                self.global_best = self.particles[i]

            # update velocity
            personal_attraction = (
                self.personal_best_weight
                * torch.rand(self.goal.size())
                * (self.personal_best[i] - self.particles[i])
            )

            global_attraction = (
                self.global_best_weight
                * torch.rand(self.goal.size())
                * (self.global_best - self.particles[i])
            )

            self.velocities[i] = (
                self.inertia * self.velocities[i]
                + personal_attraction
                + global_attraction
            )

            # Update position
            self.particles[i] += self.velocities[i].int()
            self.particles[i].clamp_(0, 255)

    def optimize(
        self,
        iterations: int = 1000,
    ):
        """Optimize the goal string"""
        for _ in range(iterations):
            self.update()
            best_particle = self.global_best
            print("Best Particle: ", "".join([chr(int(i)) for i in best_particle]))

swarms_torch/particle_swarm.py
kyegomez/swarms-pytorch
import torch


class QueenBeeGa:
    """
    Queen Bee evolution for genetic algos

    Inspired by the evolution of bees, the fittest solution is designated
    and the rest of the population contends to mate with it.

    The strong exploitation is balanced by a higher than a normal mutation rate.

    Reference:
    ---------
    https://www.researchgate.net/publication/228961729_A_Queen_Bee_GA_for_optimization

    Usage
    -----
    optimizer = QueenBeeGa(
        goal="Attention is all you need",
        pop_size=100,
        mutation_prob=0.04,
        strong_mutation_rate=0.1,
        strong_mutation_prob=0.25,
        num_tournament_participants=25
    )
    optimizer.run(max_generations=100)
    """

    def __init__(
        self,
        goal: str = "Attention is all you need",
        pop_size: int = 100,
        mutation_prob: float = 0.04,
        strong_mutation_rate: float = 0.1,
        strong_mutation_prob: float = 0.25,
        num_tournament_participants: int = 25,
    ):
        """
        QueenBeeGa with params and initial configs

        Parameters
        ----------
        goal: str
            The goal string to be optimized
        pop_size: int
            Number of ants
        mutation_prob: float
            Mutation rate
        strong_mutation_rate: float
            Strong mutation rate
        strong_mutation_prob: float
            Strong mutation probability
        num_tournament_participants: int
            Number of tournament participants
        """
        self.goal = goal
        self.pop_size = pop_size
        self.mutation_prob = mutation_prob
        self.strong_mutation_rate = strong_mutation_rate
        self.strong_mutation_prob = strong_mutation_prob
        self.num_tournament_participants = num_tournament_participants

        self.gene_length = len(goal)
        self.gene_midpoint = self.gene_length // 2
        self.target_gene = self.encode(goal)

        self.strong_mutate_pool_size = strong_mutation_rate * pop_size
        self.num_code_mutate = mutation_prob * self.gene_length
        self.strong_num_code_mutate = strong_mutation_prob * self.gene_length

        self.pool = torch.randint(0, 255, (pop_size, self.gene_length))
        self.queen = None
        self.queen_fitness = None
        self.generation = 0

    @staticmethod
    def encode(s):
        """Convert string to it's values"""
        return torch.tensor([ord(c) for c in s])

    @staticmethod
    def decode(t):
        """Convert ASCII values tensor back to string"""
        return "".join([chr(i) for i in t.tolist()])

    def run(self, max_generations: int = 1000):
        """
        Run the queen genertic algorithm evolution

        Parameters:
        -----------
        max_generations: int
            Maximum number of generations
        """
        for _ in range(max_generations):
            self.generation += 1
            print(f"Generation: {self.generation}")
            self._evolve()
            if self._check_convergence():
                pass

    def _evolve(self):
        """
        Execute one step of the evolution process.
        """

        # Sort population by fitness
        fitnesses = 1.0 / torch.square(self.pool - self.target_gene).sum(dim=-1)
        indices = fitnesses.sort(descending=True).indices
        self.pool, fitnesses = self.pool[indices], fitnesses[indices]

        # Display every generation
        if self.queen is not None:
            print("queen:")
            print(f"{self.decode(self.queen)} ({self.queen_fitness.item():.3f})\n")
        for gene, fitness in zip(self.pool, fitnesses):
            print(f"{self.decode(gene)} ({fitness.item():.3f})")

        # If one of the children has a better fitness than queen, that child becomes the new queen
        # and the queen replaces the worst bee in the population
        if self.queen is not None and self.queen_fitness < fitnesses[0]:
            self.pool = torch.cat((self.pool, self.queen[None, :]), dim=0)
            fitnesses = torch.cat((fitnesses, self.queen_fitness[None]), dim=0)
            self.queen = self.queen_fitness = None

        # Separate the queen bee from the rest of the population
        if self.queen is None:
            self.queen, self.pool = self.pool[0], self.pool[1:]
            self.queen_fitness, fitnesses = fitnesses[0], fitnesses[1:]

        # Deterministic tournament selection
        contender_ids = torch.randn((self.pop_size - 1, self.pop_size - 1)).argsort(
            dim=-1
        )[..., : self.num_tournament_participants]
        participants, tournaments = self.pool[contender_ids], fitnesses[contender_ids]
        top_winner = tournaments.topk(1, dim=-1, largest=True, sorted=False).indices
        top_winner = top_winner.unsqueeze(-1).expand(-1, -1, self.gene_length)
        parents = participants.gather(1, top_winner).squeeze(1)

        # Cross over all chosen drones with the queen
        queen_parents = self.queen.unsqueeze(0).expand(
            self.pop_size - 1, self.gene_length
        )
        self.pool = torch.cat(
            (queen_parents[:, : self.gene_midpoint], parents[:, self.gene_midpoint :]),
            dim=-1,
        )

        # Mutate genes in population
        mutate_mask = (
            torch.randn(self.pool.shape).argsort(dim=-1) < self.num_code_mutate
        )
        noise = torch.randint(0, 2, self.pool.shape) * 2 - 1
        mutated_pool = torch.where(mutate_mask, self.pool + noise, self.pool)

        strong_mutate_mask = (
            torch.randn(self.pool.shape).argsort(dim=-1) < self.strong_num_code_mutate
        )
        noise = torch.randint(0, 2, self.pool.shape) * 2 - 1
        strong_mutated_pool = torch.where(
            strong_mutate_mask, self.pool + noise, self.pool
        )

        strong_mutate_pool_mask = (
            torch.randn(self.pop_size - 1).argsort(dim=-1)
            < self.strong_mutate_pool_size
        )
        self.pool = torch.where(
            strong_mutate_pool_mask[:, None], strong_mutated_pool, mutated_pool
        )
        self.pool.clamp_(0, 255)

    def _check_convergence(self):
        """
        Check if any of the solutions has achieved the goal
        """
        fitnesses = 1.0 / torch.square(self.pool - self.target_gene).sum(dim=-1)
        return (fitnesses == float("inf")).any().item()


# # Usage:
# optimizer = QueenBeeGa()
# optimizer.run(max_generations=100)

swarms_torch/queen_bee.py
kyegomez/swarms-pytorch
import torch


class SPO:
    """
    Spiral Optimization (SPO) Algorithm in PyTorch.

    Implements the SPO algorithm for optimization towards a target string.

    How does it work?
    ----------
    1. Initialize the search points randomly
    2. Initialize the center randomly
    3. Update the search points based on the spiral model
    4. Find the best search point and set as the new center
    5. Repeat step 3 to 4 until the maximum number of iterations is reached

    Usage
    -----
    from swarms_torch import SPO

    goaling = "Hello SPO"
    spo = SPO(goaling, m=100, k_max=1000)
    spo.optimize()

    print("Best Matched String:", spo.best_string())

    Future Features to implement
    --------
    1. Add a stopping criterion
    2. Add a callback function to track the progress
    3. Add a function to plot the search points
    4. Add a function to plot the best solution

    """

    def __init__(self, goal: str = None, m: int = 10, k_max: int = 1000):
        """
        Initialize the SPO class.

        Args:
        - goal: The target string.
        - m: Number of search points (strings).
        - k_max: Maximum number of iterations.
        """
        self.goal = torch.tensor(
            [ord(c) for c in goal], dtype=torch.float32
        )  # ASCII representation

        self.m = m
        self.k_max = k_max
        self.n_dim = len(goal)

        # Initializing the search points and center randomly
        # Note: 32-126 is the ASCII range for all printable characters
        self.points = torch.randint(32, 127, (self.m, self.n_dim), dtype=torch.float32)
        self.center = torch.randint(32, 127, (self.n_dim,), dtype=torch.float32)

    def _step_rate(self, k):
        """
        Define the step rate function.

        Args:
        - k: Current iteration.

        Returns: Step rate for the current iteration.
        """
        return 1 / (1 + k)

    def _update_points(self, k):
        """Update the search points based on the spiral model."""
        r = self._step_rate(k)
        R = torch.eye(self.n_dim)  # Identity for simplicity in n-dimensions
        for i in range(self.m):
            self.points[i] = self.center + r * torch.mv(
                R, (self.points[i] - self.center)
            )

    def _update_center(self):
        """Find the best search point and set as the new center."""
        fitnesses = torch.norm(self.points - self.goal, dim=1)
        best_idx = torch.argmin(fitnesses)
        self.center = self.points[best_idx]

    def optimize(self):
        """Run the optimization loop."""
        for k in range(self.k_max):
            self._update_points(k)
            self._update_center()
            if (
                torch.norm(self.center - self.goal) < 1e-5
            ):  # Example convergence condition
                break

    def best_string(self):
        """Convert the best found point to its string representation"""
        return "".join([chr(int(c)) for c in self.center.round()])


# # Example Usage
# goal = "Attention is all you need"
# optimizer = SPO(goal)
# optimizer.optimize()
# print(f"Optimized String: {optimizer.best_string()}")

swarms_torch/spiral_optimization.py
kyegomez/swarms-pytorch
import torch
import torch.nn as nn
from copy import deepcopy


class Particle(nn.Module):
    """
    Simple Transformer model for classification.

    Parameters
    ----------
    input_dim : int
        The number of expected features in the input (required).
    d_model : int
        The number of expected features in the encoder/decoder inputs (required).
    nhead : int
        The number of heads in the multiheadattention models (required).
    num_layers : int
        The number of sub-encoder-layers in the encoder (required).
    output_dim : int
        The number of classes to predict (required).

    Usage:
    >>> model = SimpleTransformer(1000, 512, 8, 6, 10)
    >>> model(x)


    """

    def __init__(self, input_dim, d_model, nhead, num_layers, output_dim):
        super(Particle, self).__init__()
        self.embedding = nn.Embedding(input_dim, d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_layers, num_layers)
        self.fc = nn.Linear(d_model, output_dim)

    def forward(self, x):
        """
        Forward pass through the model.

        """
        x = self.embedding(x)
        x = self.transformer(x, x)
        return self.fc(x[-1])


class TransformerParticleSwarmOptimization(nn.Module):
    """
    Transformer Particle Swarm Optimization.

    Parameters
    ----------
    model_constructor : function
        Function to create a new model instance.
    model_args : tuple
        Arguments for the model constructor.
    device : str
        'cuda' or 'cpu'.
    criterion : nn.Module
        Loss function.
    data_loader : torch.utils.data.DataLoader
        Data loader.
    n_particles : int
        Number of particles.
    inertia : float
        Inertia weight.
    personal_best_weight : float
        Personal best weight.
    global_best_weight : float
        Global best weight.

    Usage:
    >>> pso = TransformerParticleSwarmOptimization(
    ...     SimpleTransformer,
    ...     (1000, 512, 8, 6, 10),
    ...     device="cuda",
    ...     criterion=nn.CrossEntropyLoss(),
    ...     data_loader=your_dataloader
    ... )

    """

    def __init__(
        self,
        model_constructor,  # Function to create a new model instance
        model_args,  # Arguments for the model constructor
        device,  # 'cuda' or 'cpu'
        criterion,
        data_loader,
        n_particles=10,
        inertia=0.5,
        personal_best_weight=1.5,
        global_best_weight=1.5,
    ):
        super(TransformerParticleSwarmOptimization, self).__init__()
        self.model_constructor = model_constructor
        self.model_args = model_args
        self.criterion = criterion
        self.data_loader = data_loader
        self.device = device

        self.n_particles = n_particles
        self.inertia = inertia
        self.personal_best_weight = personal_best_weight
        self.global_best_weight = global_best_weight

        # Representing particles using model parameters
        param_size = sum(p.numel() for p in model_constructor(*model_args).parameters())
        self.particles = [
            self.model_constructor(*model_args).to(device) for _ in range(n_particles)
        ]
        self.velocities = [
            torch.zeros((param_size,)).to(device) for _ in range(n_particles)
        ]
        self.personal_best = [deepcopy(p.state_dict()) for p in self.particles]
        self.global_best = deepcopy(self.particles[0].state_dict())

    def compute_fitness(self, model_state):
        """
        Compute the fitness of a model.
        """
        model = self.model_constructor(*self.model_args).to(self.device)
        model.load_state_dict(model_state)
        model.eval()

        total_loss = 0.0
        with torch.no_grad():
            for inputs, targets in self.data_loader:
                outputs = model(inputs.to(self.device))
                loss = self.criterion(outputs, targets.to(self.device))
                total_loss += loss.item()
        return 1.0 / (1.0 + total_loss)

    def update(self):
        """
        Update particles.
        """
        # Update particles
        for idx, particle in enumerate(self.particles):
            fitness = self.compute_fitness(particle.state_dict())

            # Update personal best
            if fitness > self.compute_fitness(self.personal_best[idx]):
                self.personal_best[idx] = deepcopy(particle.state_dict())

            # Update global best
            if fitness > self.compute_fitness(self.global_best):
                self.global_best = deepcopy(particle.state_dict())

            # Update velocities and positions
            for name, param in particle.named_parameters():
                delta = self.personal_best_weight * torch.rand_like(param) * (
                    self.personal_best[idx][name].to(self.device) - param.data
                ) + self.global_best_weight * torch.rand_like(param) * (
                    self.global_best[name].to(self.device) - param.data
                )
                self.velocities[idx] += self.inertia * self.velocities[idx] + delta
                param.data += self.velocities[idx]

    def optimize(self, iterations=1000):
        """Optimize the model."""
        for _ in range(iterations):
            self.update()
            best_particle_score = self.compute_fitness(self.global_best)
            print(
                f"Iteration {_ + 1}/{iterations} - Best Particle Fitness: {best_particle_score}"
            )

    def get_best_model(self):
        """Get the best model."""
        best_model = self.model_constructor(*self.model_args).to(self.device)
        best_model.load_state_dict(self.global_best)
        return best_model


# # Define model and optimization parameters
# input_dim = 1000
# d_model = 512
# nhead = 8
# num_layers = 3
# output_dim = 10

# batch_size = 32
# sequence_length = 50

# # Instantiate the optimizer
# pso = ParticleSwarmOptimization(
#     SimpleTransformer,
#     (input_dim, d_model, nhead, num_layers, output_dim),
#     device="cuda",  # or 'cpu'
#     criterion=nn.CrossEntropyLoss(),
#     # data_loader=your_dataloader  # replace with your dataloader
# )

# # Run optimization
# pso.optimize(iterations=100)

# # Get the best model
# best_model = pso.get_best_model()

# # Generate a random input tensor
# x = torch.randint(0, input_dim, (batch_size, sequence_length)).to(
#     "cuda"
# )  # ensure it's on the same device as your model

# # Pass the tensor through the model
# output = best_model(x)
# print(output.shape)  # should be [batch_size, output_dim]

swarms_torch/transformer_pso.py
kyegomez/swarms-pytorch
import unittest
import torch

from swarms_torch import AntColonyOptimization  # Import your class


class TestAntColonyOptimization(unittest.TestCase):
    def setUp(self):
        self.aco = AntColonyOptimization(
            goal="Hello ACO", num_ants=1000, num_iterations=10
        )

    def test_initialization(self):
        self.assertEqual(self.aco.goal.tolist(), [ord(c) for c in "Hello ACO"])
        self.assertEqual(self.aco.pheromones.size(), torch.Size([1000]))
        self.assertEqual(self.aco.pheromones.tolist(), [1.0] * 1000)

    def test_fitness(self):
        solution = torch.tensor([ord(c) for c in "Hello ACO"], dtype=torch.float32)
        self.assertEqual(
            self.aco.fitness(solution).item(), 0
        )  # Should be maximum fitness

    def test_update_pheromones(self):
        initial_pheromones = self.aco.pheromones.clone()
        self.aco.solutions = [
            torch.tensor([ord(c) for c in "Hello ACO"], dtype=torch.float32)
            for _ in range(1000)
        ]
        self.aco.update_pheromones()
        # After updating, pheromones should not remain the same
        self.assertFalse(torch.equal(initial_pheromones, self.aco.pheromones))

    def test_choose_next_path(self):
        path = self.aco.choose_next_path()
        # Path should be an integer index within the number of ants
        self.assertIsInstance(path, int)
        self.assertGreaterEqual(path, 0)
        self.assertLess(path, 1000)

    def test_optimize(self):
        solution = self.aco.optimize()
        self.assertIsInstance(solution, str)
        # Given enough iterations and ants, the solution should approach the goal. For short runs, this might not hold.
        # self.assertEqual(solution, "Hello ACO")

    def test_invalid_parameters(self):
        with self.assertRaises(ValueError):
            _ = AntColonyOptimization(num_ants=-5)
        with self.assertRaises(ValueError):
            _ = AntColonyOptimization(evaporation_rate=1.5)


if __name__ == "__main__":
    unittest.main()

tests/ant_colony.py
kyegomez/swarms-pytorch
import pytest
import torch
from swarms_torch import TransformerCell, CellularSwarm


def test_transformercell_initialization():
    transformercell = TransformerCell(input_dim=512, nhead=8)
    assert isinstance(transformercell, TransformerCell)
    assert transformercell.neighborhood_size == 3


def test_transformercell_forward():
    transformercell = TransformerCell(input_dim=512, nhead=8)
    x = torch.randn(10, 32, 512)
    neighbors = [torch.randn(10, 32, 512)]
    output = transformercell(x, neighbors)
    assert output.shape == torch.Size([20, 32, 512])


def test_cellularswarm_initialization():
    cellularswarm = CellularSwarm(cell_count=5, input_dim=512, nhead=8)
    assert isinstance(cellularswarm, CellularSwarm)
    assert len(cellularswarm.cells) == 5
    assert cellularswarm.time_steps == 4


def test_cellularswarm_forward():
    cellularswarm = CellularSwarm(cell_count=5, input_dim=512, nhead=8)
    x = torch.randn(10, 32, 512)
    output = cellularswarm(x)
    assert output.shape == torch.Size([10, 32, 512])

tests/cellular_swarm.py
kyegomez/swarms-pytorch
import pytest
import torch
from swarms_torch.fish_school import Fish, FishSchool


def test_fish_initialization():
    fish = Fish(dim=512, heads=8, depth=6)
    assert isinstance(fish, Fish)


def test_fish_train():
    fish = Fish(dim=512, heads=8, depth=6)
    src = torch.randn(10, 32, 512)
    tgt = torch.randn(10, 32, 512)
    labels = torch.randint(0, 512, (10, 32))
    fish.train(src, tgt, labels)
    assert isinstance(fish.food, float)


def test_fishschool_initialization():
    fishschool = FishSchool(num_fish=10, dim=512, heads=8, depth=6, num_iter=100)
    assert isinstance(fishschool, FishSchool)
    assert len(fishschool.fish) == 10


def test_fishschool_forward():
    fishschool = FishSchool(num_fish=10, dim=512, heads=8, depth=6, num_iter=100)
    src = torch.randn(10, 32, 512)
    tgt = torch.randn(10, 32, 512)
    labels = torch.randint(0, 512, (10, 32))
    fishschool.forward(src, tgt, labels)
    assert isinstance(fishschool.fish[0].food, float)

tests/fish_school.py
kyegomez/swarms-pytorch
import pytest
import torch
from swarms_torch.neuronal_transformer import (
    TransformerLayer,
    Neuron,
    SynapseTransformer,
    NNTransformer,
)


def test_transformerlayer_initialization():
    transformerlayer = TransformerLayer(input_dim=512, output_dim=256, nhead=8)
    assert isinstance(transformerlayer, TransformerLayer)


def test_transformerlayer_forward():
    transformerlayer = TransformerLayer(input_dim=512, output_dim=256, nhead=8)
    x = torch.randn(10, 32, 512)
    output = transformerlayer(x)
    assert output.shape == torch.Size([10, 32, 256])


def test_neuron_initialization():
    neuron = Neuron(num_states=10)
    assert isinstance(neuron, Neuron)
    assert neuron.states.shape == torch.Size([10])


def test_synapsetransformer_initialization():
    synapsetransformer = SynapseTransformer(input_dim=512, output_dim=256, nhead=8)
    assert isinstance(synapsetransformer, SynapseTransformer)


def test_synapsetransformer_forward():
    synapsetransformer = SynapseTransformer(input_dim=512, output_dim=256, nhead=8)
    x = torch.randn(10, 32, 512)
    output = synapsetransformer(x)
    assert output.shape == torch.Size([10, 32, 256])


def test_nntransformer_initialization():
    nntransformer = NNTransformer(
        neuron_count=5, num_states=10, input_dim=512, output_dim=256, nhead=8
    )
    assert isinstance(nntransformer, NNTransformer)
    assert len(nntransformer.neurons) == 5
    assert len(nntransformer.synapses) == 5


def test_nntransformer_forward():
    nntransformer = NNTransformer(
        neuron_count=5, num_states=10, input_dim=512, output_dim=256, nhead=8
    )
    x = torch.randn(1, 10)
    output = nntransformer(x)
    assert output.shape == torch.Size([10])

tests/neuronal_transformer.py
kyegomez/swarms-pytorch
import unittest
import torch

from swarms_torch import ParticleSwarmOptimization  # Import your class here


class TestParticleSwarmOptimization(unittest.TestCase):
    def setUp(self):
        self.pso = ParticleSwarmOptimization(goal="Hello", n_particles=10)

    def test_initialization(self):
        self.assertEqual(self.pso.goal.tolist(), [ord(c) for c in "Hello"])
        self.assertEqual(self.pso.particles.size(), (10, 5))
        self.assertEqual(self.pso.velocities.size(), (10, 5))

    def test_compute_fitness(self):
        particle = torch.tensor([ord(c) for c in "Hello"])
        fitness = self.pso.compute_fitness(particle)
        self.assertEqual(fitness.item(), 1.0)

    def test_update(self):
        initial_particle = self.pso.particles.clone()
        self.pso.update()
        # After updating, particles should not remain the same (in most cases)
        self.assertFalse(torch.equal(initial_particle, self.pso.particles))

    def test_optimize(self):
        initial_best_particle = self.pso.global_best.clone()
        self.pso.optimize(iterations=10)
        # After optimization, global best should be closer to the goal
        initial_distance = torch.norm(
            (initial_best_particle - self.pso.goal).float()
        ).item()
        final_distance = torch.norm(
            (self.pso.global_best - self.pso.goal).float()
        ).item()
        self.assertLess(final_distance, initial_distance)


if __name__ == "__main__":
    unittest.main()

tests/particle_swarm.py
kyegomez/swarms-pytorch
import unittest
import torch
from swarms_torch.queen_bee import QueenBeeGa  # Import the class


class TestQueenBeeGa(unittest.TestCase):
    def setUp(self):
        self.optimizer = QueenBeeGa(goal="Hello QBGA", pop_size=50)

    def test_initialization(self):
        self.assertEqual(self.optimizer.goal, "Hello QBGA")
        self.assertEqual(self.optimizer.gene_length, len("Hello QBGA"))
        self.assertIsNone(self.optimizer.queen)
        self.assertIsNone(self.optimizer.queen_fitness)

    def test_encode_decode(self):
        encoded = QueenBeeGa.encode("Hello")
        decoded = QueenBeeGa.decode(encoded)
        self.assertEqual(decoded, "Hello")

    def test_evolution(self):
        initial_population = self.optimizer.pool.clone()
        self.optimizer._evolve()
        self.assertFalse(torch.equal(initial_population, self.optimizer.pool))

    def test_run(self):
        initial_population = self.optimizer.pool.clone()
        self.optimizer.run(max_generations=10)
        self.assertNotEqual(
            QueenBeeGa.decode(self.optimizer.queen),
            QueenBeeGa.decode(initial_population[0]),
        )

    def test_check_convergence(self):
        self.optimizer.pool = torch.stack([self.optimizer.target_gene] * 50)
        self.assertTrue(self.optimizer._check_convergence())

    def test_invalid_parameters(self):
        with self.assertRaises(ValueError):
            _ = QueenBeeGa(mutation_prob=1.5)
        with self.assertRaises(ValueError):
            _ = QueenBeeGa(strong_mutation_rate=-0.5)


if __name__ == "__main__":
    unittest.main()

tests/queen_bee.py
kyegomez/swarms-pytorch
import pytest
import torch
from swarms_torch.spiral_optimization import SPO


def test_spo_initialization():
    spo = SPO(goal="Hello SPO", m=100, k_max=1000)
    assert isinstance(spo, SPO)
    assert spo.goal.shape == torch.Size([9])
    assert spo.points.shape == torch.Size([100, 9])
    assert spo.center.shape == torch.Size([9])


def test_spo_step_rate():
    spo = SPO(goal="Hello SPO", m=100, k_max=1000)
    step_rate = spo._step_rate(1)
    assert step_rate == 0.5


def test_spo_update_points():
    spo = SPO(goal="Hello SPO", m=100, k_max=1000)
    spo._update_points(1)
    assert spo.points.shape == torch.Size([100, 9])


def test_spo_update_center():
    spo = SPO(goal="Hello SPO", m=100, k_max=1000)
    spo._update_center()
    assert spo.center.shape == torch.Size([9])


def test_spo_optimize():
    spo = SPO(goal="Hello SPO", m=100, k_max=1000)
    spo.optimize()
    assert spo.center.shape == torch.Size([9])


def test_spo_best_string():
    spo = SPO(goal="Hello SPO", m=100, k_max=1000)
    spo.optimize()
    best_string = spo.best_string()
    assert isinstance(best_string, str)
    assert len(best_string) == 9

tests/spiral_optimization.py
kyegomez/swarms-pytorch
import pytest
import torch
from torch.utils.data import DataLoader
from swarms_torch.transformer_pso import (
    SimpleTransformer,
    TransformerParticleSwarmOptimization,
)


def test_simpletransformer_initialization():
    simpletransformer = SimpleTransformer(
        input_dim=10, d_model=512, nhead=8, num_layers=1, output_dim=2
    )
    assert isinstance(simpletransformer, SimpleTransformer)


def test_simpletransformer_forward():
    simpletransformer = SimpleTransformer(
        input_dim=10, d_model=512, nhead=8, num_layers=1, output_dim=2
    )
    x = torch.randint(0, 10, (10, 32))
    output = simpletransformer(x)
    assert output.shape == torch.Size([32, 2])


def test_TransformerParticleSwarmOptimization_initialization():
    model_constructor = SimpleTransformer
    model_args = (10, 512, 8, 1, 2)
    device = "cpu"
    criterion = torch.nn.CrossEntropyLoss()
    data_loader = DataLoader(
        [(torch.randint(0, 10, (10,)), torch.tensor(1)) for _ in range(100)],
        batch_size=32,
    )
    pso = TransformerParticleSwarmOptimization(
        model_constructor, model_args, device, criterion, data_loader
    )
    assert isinstance(pso, TransformerParticleSwarmOptimization)
    assert len(pso.particles) == 10
    assert len(pso.velocities) == 10
    assert len(pso.personal_best) == 10


def test_TransformerParticleSwarmOptimization_compute_fitness():
    model_constructor = SimpleTransformer
    model_args = (10, 512, 8, 1, 2)
    device = "cpu"
    criterion = torch.nn.CrossEntropyLoss()
    data_loader = DataLoader(
        [(torch.randint(0, 10, (10,)), torch.tensor(1)) for _ in range(100)],
        batch_size=32,
    )
    pso = TransformerParticleSwarmOptimization(
        model_constructor, model_args, device, criterion, data_loader
    )
    fitness = pso.compute_fitness(pso.particles[0].state_dict())
    assert isinstance(fitness, float)


def test_TransformerParticleSwarmOptimization_update():
    model_constructor = SimpleTransformer
    model_args = (10, 512, 8, 1, 2)
    device = "cpu"
    criterion = torch.nn.CrossEntropyLoss()
    data_loader = DataLoader(
        [(torch.randint(0, 10, (10,)), torch.tensor(1)) for _ in range(100)],
        batch_size=32,
    )
    pso = TransformerParticleSwarmOptimization(
        model_constructor, model_args, device, criterion, data_loader
    )
    pso.update()
    assert len(pso.particles) == 10
    assert len(pso.velocities) == 10
    assert len(pso.personal_best) == 10


def test_TransformerParticleSwarmOptimization_optimize():
    model_constructor = SimpleTransformer
    model_args = (10, 512, 8, 1, 2)
    device = "cpu"
    criterion = torch.nn.CrossEntropyLoss()
    data_loader = DataLoader(
        [(torch.randint(0, 10, (10,)), torch.tensor(1)) for _ in range(100)],
        batch_size=32,
    )
    pso = TransformerParticleSwarmOptimization(
        model_constructor, model_args, device, criterion, data_loader
    )
    pso.optimize(iterations=10)
    assert len(pso.particles) == 10
    assert len(pso.velocities) == 10
    assert len(pso.personal_best) == 10


def test_TransformerParticleSwarmOptimization_get_best_model():
    model_constructor = SimpleTransformer
    model_args = (10, 512, 8, 1, 2)
    device = "cpu"
    criterion = torch.nn.CrossEntropyLoss()
    data_loader = DataLoader(
        [(torch.randint(0, 10, (10,)), torch.tensor(1)) for _ in range(100)],
        batch_size=32,
    )
    pso = TransformerParticleSwarmOptimization(
        model_constructor, model_args, device, criterion, data_loader
    )
    pso.optimize(iterations=10)
    best_model = pso.get_best_model()
    assert isinstance(best_model, SimpleTransformer)

tests/transformer_pso.py
kyegomez/swarms-pytorch
import gzip
import random

import numpy as np
import torch
import torch.optim as optim
import tqdm
from torch.nn import functional as F
from torch.utils.data import DataLoader, Dataset
from swarms_torch.fish_school import FishSchool
from swarms_torch.autoregressive import AutoregressiveWrapper

# constants

NUM_BATCHES = int(1e5)
BATCH_SIZE = 4
GRADIENT_ACCUMULATE_EVERY = 4
LEARNING_RATE = 2e-4
VALIDATE_EVERY = 100
GENERATE_EVERY = 500
GENERATE_LENGTH = 512
SEQ_LEN = 1024

# helpers


def cycle(loader):
    while True:
        for data in loader:
            yield data


def decode_token(token):
    return str(chr(max(32, token)))


def decode_tokens(tokens):
    return "".join(list(map(decode_token, tokens)))


# init model
model = FishSchool(10, 512, 8, 6, 100)
model = AutoregressiveWrapper(model)

# model = AutoregressiveWrapper(model)
# model.cuda()

# prepare enwik8 data

with gzip.open("./data/enwik8.gz") as file:
    X = np.fromstring(file.read(int(95e6)), dtype=np.uint8)
    trX, vaX = np.split(X, [int(90e6)])
    data_train, data_val = torch.from_numpy(trX), torch.from_numpy(vaX)


class TextSamplerDataset(Dataset):
    def __init__(self, data, seq_len):
        super().__init__()
        self.data = data
        self.seq_len = seq_len

    def __getitem__(self, index):
        rand_start = torch.randint(0, self.data.size(0) - self.seq_len, (1,))
        full_seq = self.data[rand_start : rand_start + self.seq_len + 1].long()
        return full_seq.cuda()

    def __len__(self):
        return self.data.size(0) // self.seq_len


train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
val_dataset = TextSamplerDataset(data_val, SEQ_LEN)
train_loader = cycle(DataLoader(train_dataset, batch_size=BATCH_SIZE))
val_loader = cycle(DataLoader(val_dataset, batch_size=BATCH_SIZE))

# optimizer

optim = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

# training

for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10.0, desc="training"):
    model.train()

    for __ in range(GRADIENT_ACCUMULATE_EVERY):
        loss = model(next(train_loader))
        loss.backward()

    print(f"training loss: {loss.item()}")
    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
    optim.step()
    optim.zero_grad()

    if i % VALIDATE_EVERY == 0:
        model.eval()
        with torch.no_grad():
            loss = model(next(val_loader))
            print(f"validation loss: {loss.item()}")

    if i % GENERATE_EVERY == 0:
        model.eval()
        inp = random.choice(val_dataset)[:-1]
        prime = decode_tokens(inp)
        print(f"%s \n\n %s", (prime, "*" * 100))

        sample = model.generate(inp[None, ...], GENERATE_LENGTH)
        output_str = decode_tokens(sample[0])
        print(output_str)

train.py
kyegomez/swarms-pytorch