init

.env.example

@@ -0,0 +1 @@
+REDIS_URL=""

.gitignore

@@ -0,0 +1,17 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+.env
+.venv
+
+/models/
+
+frontend/node_modules
+frontend/build
+
+.DS_Store

.python-version

@@ -0,0 +1 @@
+3.12

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.12-slim
+WORKDIR /code 
+
+RUN pip install uv
+
+COPY pyproject.toml uv.lock ./
+
+RUN uv export --no-dev | uv pip install --system -r -
+
+COPY evolutiontransformer/ ./evolutiontransformer/

docker-compose.yaml

@@ -0,0 +1,18 @@
+services:
+  api:
+    build: .
+    command: >
+      sh -c "gunicorn evolutiontransformer.api:app -w 4 -k uvicorn.workers.UvicornWorker -b 0.0.0.0:8000"
+    ports:
+      - "8000:8000"
+    depends_on:
+      - worker
+    env_file:
+      - .env
+
+  worker:
+    build: .
+    command: >
+      sh -c "/usr/local/bin/celery -A evolutiontransformer.worker.celery_app worker --loglevel=info -c 1"
+    env_file:
+      - .env

evolutiontransformer/api.py

@@ -0,0 +1,98 @@
+import os
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+import uuid
+from typing import List, Tuple
+from fastapi import FastAPI, Depends, HTTPException, Request, Response
+from pydantic import BaseModel
+from celery import Celery
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
+
+
+celery_app = Celery("tasks", broker=REDIS_URL, backend=REDIS_URL)
+
+app = FastAPI()
+
+
+class GenerateRequest(BaseModel):
+    model_name: str
+    prompt: str
+    max_new_tokens: int = 512
+    temperature: float = 0.7
+
+
+class MergeRequest(BaseModel):
+    model1_name: str
+    model2_name: str
+    layer_recipe: List[List[Tuple[int, int, float]]]
+    embedding_lambdas: List[float] = [0.5, 0.5]
+    linear_lambdas: List[float] = [0.5, 0.5]
+    merged_name: str = "merged"
+
+
+def get_session_id(request: Request, response: Response):
+    session_id = request.cookies.get("session_id")
+
+    if not session_id:
+        session_id = str(uuid.uuid4())
+        response.set_cookie(key="session_id", value=session_id)
+
+    return session_id
+
+
+@app.post("/generate")
+def generate(request: GenerateRequest, session_id: str = Depends(get_session_id)):
+    task = celery_app.send_task(
+        "tasks.inference",
+        args=[
+            session_id,
+            request.model_name,
+            request.prompt,
+            request.max_new_tokens,
+            request.temperature,
+        ],
+    )
+    return {"task_id": task.id}
+
+
+@app.post("/merge")
+def merge(request: MergeRequest, session_id: str = Depends(get_session_id)):
+    task = celery_app.send_task(
+        "tasks.merge_models",
+        args=[
+            session_id,
+            request.model1_name,
+            request.model2_name,
+            request.layer_recipe,
+            request.embedding_lambdas,
+            request.linear_lambdas,
+            request.merged_name,
+        ],
+    )
+
+    return {"task_id": task.id}
+
+
+@app.post("/list_models")
+def list_models(session_id: str = Depends(get_session_id)):
+    task = celery_app.send_task("tasks.get_all_models", args=[session_id])
+    return {"task_id": task.id}
+
+
+@app.get("/tasks/{task_id}")
+def get_task_status(task_id: str):
+    task_result = celery_app.AsyncResult(task_id)
+
+    if task_result.ready():
+        if task_result.status == "FAILURE":
+            raise HTTPException(status_code=500, detail=str(task_result.result))
+        else:
+            return {"status": task_result.status, "result": task_result.result}
+    else:
+        return {"status": task_result.status}

evolutiontransformer/redis.py

@@ -0,0 +1,45 @@
+import os
+from redis import Redis
+import json
+
+
+REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
+redis_client = Redis.from_url(REDIS_URL, decode_responses=True)
+
+
+def add_model_to_session(session_id: str, model_name: str, ttl_seconds: int = 3600):
+    session_key = f"session:{session_id}:models"
+    redis_client.sadd(session_key, model_name)
+    redis_client.expire(session_key, ttl_seconds)
+
+
+def get_session_models(session_id: str):
+    session_key = f"session:{session_id}:models"
+    models = redis_client.smembers(session_key)
+    return list(models)
+
+
+def save_model_recipe(
+    session_id: str, model_name: str, recipe: dict, ttl_seconds: int = 3600
+):
+    recipe_key = f"model:{session_id}:{model_name}"
+    serialized_recipe = json.dumps(recipe)
+    redis_client.setex(recipe_key, ttl_seconds, serialized_recipe)
+
+
+def get_model_recipe(session_id: str, model_name: str):
+    recipe_key = f"model:{session_id}:{model_name}"
+    serialized_recipe = redis_client.get(recipe_key)
+    if serialized_recipe:
+        return json.loads(serialized_recipe)
+    return None
+
+
+def delete_session(session_id: str):
+    model_names = get_session_models(session_id)
+
+    for model_name in model_names:
+        recipe_key = f"model:{session_id}:{model_name}"
+        redis_client.delete(recipe_key)
+
+    redis_client.delete(f"session:{session_id}:models")

evolutiontransformer/worker.py

@@ -0,0 +1,271 @@
+import os
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+from celery import Celery
+from celery.exceptions import InvalidTaskError
+import torch
+import torch.nn as nn
+from dotenv import load_dotenv
+from evolutiontransformer.redis import (
+    add_model_to_session,
+    get_session_models,
+    save_model_recipe,
+    get_model_recipe,
+)
+
+
+from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
+from typing import List, Tuple
+from tqdm import tqdm
+
+load_dotenv()
+
+BASE_MODELS_NAMES = ["svamp", "tinystories"]
+BASE_MODELS = {}
+TOKENIZER = None
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
+
+celery_app = Celery("tasks", broker=REDIS_URL, backend=REDIS_URL)
+
+
+def load_base_models_if_needed():
+    global BASE_MODELS
+    if not BASE_MODELS:
+        print("WORKER: Loading base models into memory...")
+        for model_name in BASE_MODELS_NAMES:
+            model_path = f"tcmmichaelb139/gpt2-medium-{model_name}"
+            model = AutoModelForCausalLM.from_pretrained(model_path)
+            BASE_MODELS[model_name] = model.to(DEVICE)
+
+            if get_model_recipe("default", model_name) is None:
+                add_model_to_session("default", model_name)
+                save_model_recipe(
+                    "default",
+                    model_name,
+                    {
+                        "layer_recipe": [[(i, model_name, 1.0)] for i in range(24)],
+                        "embedding_lambdas": [1.0, 1.0],
+                        "linear_lambdas": [1.0, 1.0],
+                    },
+                )
+
+        print("WORKER: Base models loaded.")
+
+
+def get_tokenizer():
+    global TOKENIZER
+    if TOKENIZER is None:
+        print("WORKER: Initializing Tokenizer...")
+        TOKENIZER = AutoTokenizer.from_pretrained("gpt2-medium")
+    return TOKENIZER
+
+
+def inference(model, prompt, max_new_tokens=512, temperature=0.7):
+    global DEVICE
+
+    do_sample = temperature > 0
+    model = model.to(DEVICE)
+    model.eval()
+    tokenizer = get_tokenizer()
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+            temperature=temperature,
+        ).to(DEVICE)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+def merge_model_recipe(
+    model1_recipe: dict,
+    model2_recipe: dict,
+    layer_recipe: List[List[Tuple[int, int, float]]],
+    embedding_lambdas: List[float] = [0.5, 0.5],
+    linear_lambdas: List[float] = [0.5, 0.5],
+) -> dict:
+    models = [model1_recipe, model2_recipe]
+    result_layer_recipe = []
+    for makeup in layer_recipe:
+        layer_result = {}
+        for comb in makeup:
+            idx, model_i, alpha = comb
+
+            for orig_i, orig_model, orig_a in models[model_i]["layer_recipe"][idx]:
+                if (orig_i, orig_model) in layer_result:
+                    layer_result[(orig_i, orig_model)] += alpha * orig_a
+                else:
+                    layer_result[(orig_i, orig_model)] = alpha * orig_a
+
+        final_layer_result = []
+        for k in layer_result:
+            final_layer_result.append((k[0], k[1], layer_result[k]))
+
+        result_layer_recipe.append(final_layer_result)
+
+    result_embedding_lambdas = [
+        embedding_lambdas[0] * model1_recipe["embedding_lambdas"][0]
+        + (1 - embedding_lambdas[0]) * model2_recipe["embedding_lambdas"][0],
+        embedding_lambdas[1] * model1_recipe["embedding_lambdas"][1]
+        + (1 - embedding_lambdas[1]) * model2_recipe["embedding_lambdas"][1],
+    ]
+    result_linear_lambdas = [
+        linear_lambdas[0] * model1_recipe["linear_lambdas"][0]
+        + (1 - linear_lambdas[0]) * model2_recipe["linear_lambdas"][0],
+        linear_lambdas[1] * model1_recipe["linear_lambdas"][1]
+        + (1 - linear_lambdas[1]) * model2_recipe["linear_lambdas"][1],
+    ]
+
+    return {
+        "layer_recipe": result_layer_recipe,
+        "embedding_lambdas": result_embedding_lambdas,
+        "linear_lambdas": result_linear_lambdas,
+    }
+
+
+def merge_models(
+    model_recipe: dict,
+    base_model="gpt2-medium",
+) -> nn.Module:
+    """Merge two models based on a given recipe."""
+
+    model1_name = "svamp"
+    model2_name = "tinystories"
+
+    load_base_models_if_needed()
+
+    def get_model_layer(layer, model):
+        return model.transformer.h[layer].state_dict()
+
+    def merge_layer(recipe: List[Tuple[int, str, float]]):
+        base = get_model_layer(recipe[0][0], BASE_MODELS[recipe[0][1]])
+        for key in base.keys():
+            base[key] = recipe[0][2] * base[key]
+        for layer in recipe[1:]:
+            layer_data = get_model_layer(layer[0], BASE_MODELS[layer[1]])
+            for key in base.keys():
+                base[key] += layer[2] * layer_data[key]
+        return base
+
+    print("### Merging models... ###")
+
+    layer_recipe = model_recipe["layer_recipe"]
+    embedding_lambdas = model_recipe["embedding_lambdas"]
+    linear_lambdas = model_recipe["linear_lambdas"]
+
+    config = AutoConfig.from_pretrained(base_model)
+    config.n_layer = len(layer_recipe)
+
+    child_model = AutoModelForCausalLM.from_config(config).to(DEVICE)
+    child_model.eval()
+
+    print("Merging embeddings and lm_head...")
+    child_model.transformer.wte.weight.data = (
+        embedding_lambdas[0] * BASE_MODELS[model1_name].transformer.wte.weight.data
+        + (1 - embedding_lambdas[0])
+        * BASE_MODELS[model2_name].transformer.wte.weight.data
+    )
+    child_model.transformer.wpe.weight.data = (
+        embedding_lambdas[1] * BASE_MODELS[model1_name].transformer.wpe.weight.data
+        + (1 - embedding_lambdas[1])
+        * BASE_MODELS[model2_name].transformer.wpe.weight.data
+    )
+    child_model.lm_head.weight.data = (
+        linear_lambdas[0] * BASE_MODELS[model1_name].lm_head.weight.data
+        + (1 - linear_lambdas[0]) * BASE_MODELS[model2_name].lm_head.weight.data
+    )
+    child_model.transformer.ln_f.weight.data = (
+        linear_lambdas[1] * BASE_MODELS[model1_name].transformer.ln_f.weight.data
+        + (1 - linear_lambdas[1])
+        * BASE_MODELS[model2_name].transformer.ln_f.weight.data
+    )
+    child_model.transformer.ln_f.bias.data = (
+        linear_lambdas[1] * BASE_MODELS[model1_name].transformer.ln_f.bias.data
+        + (1 - linear_lambdas[1]) * BASE_MODELS[model2_name].transformer.ln_f.bias.data
+    )
+
+    for i, layer in tqdm(enumerate(layer_recipe), desc="Merging layers..."):
+        merged_layer = merge_layer(layer)
+        child_model.transformer.h[i].load_state_dict(merged_layer)
+
+    return child_model
+
+
+def get_model_recipe_default(session_id: str, model_name: str) -> dict:
+    if model_name in BASE_MODELS_NAMES:
+        return get_model_recipe("default", model_name)
+    return get_model_recipe(session_id, model_name)
+
+
+@celery_app.task(name="tasks.inference")
+def inference_task(
+    session_id: str, model_name, prompt, max_new_tokens=512, temperature=0.7
+):
+    try:
+        model_recipe = get_model_recipe_default(session_id, model_name)
+        print("WORKER: Creating merged model...")
+        model = merge_models(model_recipe)
+        print("WORKER: Model loaded.")
+        output = inference(model, prompt, max_new_tokens, temperature)
+        return {"response": output}
+    except Exception as e:
+        raise InvalidTaskError(f"Inference failed: {e}")
+
+
+@celery_app.task(name="tasks.merge_models")
+def merge_models_task(
+    session_id: str,
+    model1_name: str,
+    model2_name: str,
+    layer_recipe: List[List[Tuple[int, int, float]]],
+    embedding_lambdas: List[float] = [0.5, 0.5],
+    linear_lambdas: List[float] = [0.5, 0.5],
+    merged_name: str = "merged",
+):
+    if len(layer_recipe) > 48:
+        raise InvalidTaskError("Layer recipe too long. Max 48 layers supported.")
+
+    session_models = get_session_models(session_id)
+
+    model1_recipe = get_model_recipe_default(session_id, model1_name)
+    model2_recipe = get_model_recipe_default(session_id, model2_name)
+    if model1_recipe is None or model2_recipe is None:
+        raise InvalidTaskError("One of the models does not exist.")
+
+    merged_recipe = merge_model_recipe(
+        model1_recipe,
+        model2_recipe,
+        layer_recipe,
+        embedding_lambdas,
+        linear_lambdas,
+    )
+
+    for i in range(20):
+        full_merged_name = f"{merged_name}_{i}"
+        if full_merged_name not in session_models:
+            add_model_to_session(session_id, full_merged_name)
+            save_model_recipe(session_id, full_merged_name, merged_recipe)
+            return {"response": full_merged_name}
+
+    raise InvalidTaskError("Could not find a unique model name.")
+
+
+@celery_app.task(name="tasks.get_all_models")
+def get_all_models_task(session_id: str) -> List[str]:
+    global SESSION_MODELS
+    return {
+        "response": list((BASE_MODELS | SESSION_MODELS[session_id]).keys()),
+    }
+
+
+@celery_app.task(name="tasks.clear_session_models")
+def clear_session_models_task(session_id: str) -> str:
+    global SESSION_MODELS
+    if session_id in SESSION_MODELS:
+        del SESSION_MODELS[session_id]
+    del SESSION_MODELS[session_id]
+    return {"response": "SUCCESS"}

finetuning/finetuning.ipynb

@@ -0,0 +1,336 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "%pip install evaluate"
+      ],
+      "metadata": {
+        "id": "aqcbe-No3r2r"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "lOwXY3N4tmbr"
+      },
+      "outputs": [],
+      "source": [
+        "import numpy as np\n",
+        "import matplotlib\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "from torch.utils.data import Dataset, DataLoader\n",
+        "from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments\n",
+        "from datasets import load_dataset\n",
+        "import evaluate\n",
+        "from copy import deepcopy\n",
+        "\n",
+        "SEED=42\n",
+        "MODEL=\"gpt2-medium\"\n",
+        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def tokenize(x, tokenizer):\n",
+        "  output = tokenizer(x[\"text\"], padding=\"max_length\", truncation=True, max_length=512)\n",
+        "  output[\"label\"] = output[\"input_ids\"].copy()\n",
+        "  return output\n",
+        "\n",
+        "def gen_tokenizer(model_name):\n",
+        "  tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+        "  tokenizer.pad_token = tokenizer.eos_token\n",
+        "  return tokenizer\n",
+        "\n",
+        "def finetune(config):\n",
+        "  ds = config[\"ds\"]\n",
+        "  preprocess_function = config[\"datasets_preprocess\"][config[\"dataset\"]]\n",
+        "  tokenizer = gen_tokenizer(config[\"model\"])\n",
+        "\n",
+        "  train_dataset = ds[\"train\"].select(range(config[\"max_train_size\"])).map(\n",
+        "    lambda x: preprocess_function(x, tokenizer),\n",
+        "  )\n",
+        "\n",
+        "\n",
+        "  train_dataset = train_dataset.map(lambda x: tokenize(x, tokenizer), batched=True)\n",
+        "\n",
+        "  model = AutoModelForCausalLM.from_pretrained(config[\"model\"])\n",
+        "  orig_model = deepcopy(model)\n",
+        "\n",
+        "  trainer = Trainer(\n",
+        "    model=model,\n",
+        "    args=config[\"training_args\"],\n",
+        "    train_dataset=train_dataset,\n",
+        "    processing_class=tokenizer,\n",
+        "  )\n",
+        "\n",
+        "  print(\"Starting training\")\n",
+        "  trainer.train()\n",
+        "  print(\"Training complete\")\n",
+        "\n",
+        "  return orig_model, model"
+      ],
+      "metadata": {
+        "id": "B3XugMEV5vZF"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def gsm8k_preprocess(x, tokenizer):\n",
+        "  return {\"text\": f\"Question: {x['question']}\\nAnswer: {x['answer']}\" + tokenizer.eos_token}\n",
+        "\n",
+        "def svamp_preprocess(x, tokenizer):\n",
+        "  return {\"text\": f\"{x['question_concat']}\\nAnswer: {x['Answer']}\" + tokenizer.eos_token}\n",
+        "\n",
+        "def tinystories_preprocess(x, tokenizer):\n",
+        "  return {\"text\": x[\"text\"] + tokenizer.eos_token}\n",
+        "\n",
+        "datasets_finetune = {\n",
+        "  \"openai/gsm8k\": gsm8k_preprocess,\n",
+        "  \"ChilleD/SVAMP\": svamp_preprocess,\n",
+        "  \"roneneldan/TinyStories\": tinystories_preprocess\n",
+        "}\n",
+        "\n",
+        "def preprocess_test_gsm8k(x):\n",
+        "  return {\"text\": f\"Question: {x['question']}\\nAnswer:\" }\n",
+        "\n",
+        "def preprocess_test_svamp(x):\n",
+        "  return {\"text\": f\"{x['question_concat']}\\nAnswer:\"}\n",
+        "\n",
+        "def preprocess_test_tinystories(x):\n",
+        "  return {\"text\": x[\"text\"]}\n",
+        "\n",
+        "datasets_finetune_test = {\n",
+        "  \"openai/gsm8k\": preprocess_test_gsm8k,\n",
+        "  \"ChilleD/SVAMP\": preprocess_test_svamp,\n",
+        "  \"roneneldan/TinyStories\": preprocess_test_tinystories\n",
+        "}"
+      ],
+      "metadata": {
+        "id": "Y09qs3FFxwx1"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def test_finetune(dataset, ds, orig_model, model, datasets_preprocess, first_x):\n",
+        "  tokenizer = gen_tokenizer(MODEL)\n",
+        "  preprocess_function = datasets_preprocess[dataset]\n",
+        "  if \"validation\" in ds:\n",
+        "    ds[\"test\"] = deepcopy(ds[\"validation\"])\n",
+        "\n",
+        "  test_dataset = ds[\"test\"].map(\n",
+        "    lambda x: preprocess_function(x),\n",
+        "  )\n",
+        "\n",
+        "  model = model.to(device)\n",
+        "  orig_model = orig_model.to(device)\n",
+        "\n",
+        "  model.eval()\n",
+        "  orig_model.eval()\n",
+        "  xi = 0\n",
+        "  with torch.no_grad():\n",
+        "    for x in test_dataset:\n",
+        "      input_tensor = tokenizer(x[\"text\"], return_tensors=\"pt\")\n",
+        "      input_tensor[\"input_ids\"] = input_tensor[\"input_ids\"].to(device)\n",
+        "      input_tensor[\"attention_mask\"] = input_tensor[\"attention_mask\"].to(device)\n",
+        "\n",
+        "      output = orig_model.generate(**input_tensor, max_new_tokens=512)\n",
+        "\n",
+        "      print(\"Original model output\")\n",
+        "      print(tokenizer.decode(output[0], skip_special_tokens=True))\n",
+        "\n",
+        "      finetuned_output = model.generate(**input_tensor, max_new_tokens=512)\n",
+        "\n",
+        "      print(\"Finetuned model output\")\n",
+        "      print(tokenizer.decode(finetuned_output[0], skip_special_tokens=True))\n",
+        "\n",
+        "      xi += 1\n",
+        "      if xi > first_x:\n",
+        "        break\n"
+      ],
+      "metadata": {
+        "id": "zQqD3dHWDj6H"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def generate_config(dataset):\n",
+        "  return config\n"
+      ],
+      "metadata": {
+        "id": "2hlh0GERDBdC"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "dataset = \"ChilleD/SVAMP\"\n",
+        "ds = load_dataset(dataset, \"default\")\n",
+        "ds_1 = dataset.split('/')[1]\n",
+        "\n",
+        "config = {\n",
+        "  \"ds\": ds,\n",
+        "  \"dataset\": dataset,\n",
+        "  \"datasets_preprocess\": datasets_finetune,\n",
+        "  \"model\": MODEL,\n",
+        "  \"max_train_size\": 700,\n",
+        "  \"training_args\": TrainingArguments(\n",
+        "    output_dir=f\"./results_{ds_1}\",\n",
+        "    report_to=\"none\",\n",
+        "    num_train_epochs=10,\n",
+        "    per_device_train_batch_size=4,\n",
+        "    warmup_steps=200,\n",
+        "    learning_rate=5e-5,\n",
+        "    weight_decay=0.01,\n",
+        "    logging_steps=200,\n",
+        "    save_strategy=\"steps\",\n",
+        "    metric_for_best_model=\"loss\",\n",
+        "    greater_is_better=False,\n",
+        "    seed=SEED,\n",
+        "  ),\n",
+        "}\n",
+        "\n",
+        "orig_model, model = finetune(config)"
+      ],
+      "metadata": {
+        "id": "WSGsa3Xtx04j"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "test_finetune(dataset, ds, orig_model, model, datasets_finetune_test, 3)"
+      ],
+      "metadata": {
+        "id": "2Z6kyEGqL7zN"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "dataset = \"roneneldan/TinyStories\"\n",
+        "ds = load_dataset(dataset, \"default\")\n",
+        "ds_1 = dataset.split('/')[1]\n",
+        "\n",
+        "\n",
+        "config = {\n",
+        "  \"ds\": ds,\n",
+        "  \"dataset\": dataset,\n",
+        "  \"datasets_preprocess\": datasets_finetune,\n",
+        "  \"model\": MODEL,\n",
+        "  \"max_train_size\": 7000,\n",
+        "  \"training_args\": TrainingArguments(\n",
+        "    output_dir=f\"./results_{ds_1}\",\n",
+        "    report_to=\"none\",\n",
+        "    num_train_epochs=1,\n",
+        "    per_device_train_batch_size=4,\n",
+        "    warmup_steps=200,\n",
+        "    learning_rate=5e-5,\n",
+        "    weight_decay=0.01,\n",
+        "    logging_steps=200,\n",
+        "    save_strategy=\"steps\",\n",
+        "    metric_for_best_model=\"loss\",\n",
+        "    greater_is_better=False,\n",
+        "    seed=SEED,\n",
+        "  ),\n",
+        "}\n",
+        "\n",
+        "orig_model, model = finetune(generate_config(dataset))"
+      ],
+      "metadata": {
+        "id": "mOzvvJWP_PL1"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "test_finetune(dataset, ds, orig_model, model, datasets_finetune_test, 3)"
+      ],
+      "metadata": {
+        "id": "X6WryZ6p3xGm"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from google.colab import files\n",
+        "files.download('/content/results_TinyStories/TinyStories-checkpoint-1750.zip')"
+      ],
+      "metadata": {
+        "id": "LBJxFu5oVP29"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "xpFlk05UW87Q"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from google.colab import files\n",
+        "files.download('/content/results_SVAMP/SVAMP-checkpoint-1750.zip')"
+      ],
+      "metadata": {
+        "id": "jxnCHVVDVO6j"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "TygO0jjlVWG_"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

frontend/.gitignore

@@ -0,0 +1,24 @@
+# Logs
+logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+pnpm-debug.log*
+lerna-debug.log*
+
+node_modules
+dist
+dist-ssr
+*.local
+
+# Editor directories and files
+.vscode/*
+!.vscode/extensions.json
+.idea
+.DS_Store
+*.suo
+*.ntvs*
+*.njsproj
+*.sln
+*.sw?

frontend/README.md

@@ -0,0 +1,12 @@
+# React + Vite
+
+This template provides a minimal setup to get React working in Vite with HMR and some ESLint rules.
+
+Currently, two official plugins are available:
+
+- [@vitejs/plugin-react](https://github.com/vitejs/vite-plugin-react/blob/main/packages/plugin-react) uses [Babel](https://babeljs.io/) for Fast Refresh
+- [@vitejs/plugin-react-swc](https://github.com/vitejs/vite-plugin-react/blob/main/packages/plugin-react-swc) uses [SWC](https://swc.rs/) for Fast Refresh
+
+## Expanding the ESLint configuration
+
+If you are developing a production application, we recommend using TypeScript with type-aware lint rules enabled. Check out the [TS template](https://github.com/vitejs/vite/tree/main/packages/create-vite/template-react-ts) for information on how to integrate TypeScript and [`typescript-eslint`](https://typescript-eslint.io) in your project.

frontend/eslint.config.js

@@ -0,0 +1,29 @@
+import js from '@eslint/js'
+import globals from 'globals'
+import reactHooks from 'eslint-plugin-react-hooks'
+import reactRefresh from 'eslint-plugin-react-refresh'
+import { defineConfig, globalIgnores } from 'eslint/config'
+
+export default defineConfig([
+  globalIgnores(['dist']),
+  {
+    files: ['**/*.{js,jsx}'],
+    extends: [
+      js.configs.recommended,
+      reactHooks.configs['recommended-latest'],
+      reactRefresh.configs.vite,
+    ],
+    languageOptions: {
+      ecmaVersion: 2020,
+      globals: globals.browser,
+      parserOptions: {
+        ecmaVersion: 'latest',
+        ecmaFeatures: { jsx: true },
+        sourceType: 'module',
+      },
+    },
+    rules: {
+      'no-unused-vars': ['error', { varsIgnorePattern: '^[A-Z_]' }],
+    },
+  },
+])

frontend/index.html

@@ -0,0 +1,13 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <link rel="icon" type="image/svg+xml" href="/vite.svg" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>Vite + React</title>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/main.jsx"></script>
+  </body>
+</html>

frontend/package.json

@@ -0,0 +1,31 @@
+{
+  "name": "frontend",
+  "private": true,
+  "version": "0.0.0",
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "vite build",
+    "lint": "eslint .",
+    "preview": "vite preview"
+  },
+  "dependencies": {
+    "@tailwindcss/vite": "^4.1.13",
+    "react": "^19.1.1",
+    "react-dom": "^19.1.1"
+  },
+  "devDependencies": {
+    "@eslint/js": "^9.35.0",
+    "@types/react": "^19.1.13",
+    "@types/react-dom": "^19.1.9",
+    "@vitejs/plugin-react": "^5.0.2",
+    "autoprefixer": "^10.4.21",
+    "eslint": "^9.35.0",
+    "eslint-plugin-react-hooks": "^5.2.0",
+    "eslint-plugin-react-refresh": "^0.4.20",
+    "globals": "^16.4.0",
+    "postcss": "^8.5.6",
+    "tailwindcss": "^4.1.13",
+    "vite": "^7.1.6"
+  }
+}

frontend/src/App.css

@@ -0,0 +1,44 @@
+@import "tailwindcss";
+
+#root {
+  max-width: 1280px;
+  margin: 0 auto;
+  padding: 2rem;
+  text-align: center;
+}
+
+.logo {
+  height: 6em;
+  padding: 1.5em;
+  will-change: filter;
+  transition: filter 300ms;
+}
+.logo:hover {
+  filter: drop-shadow(0 0 2em #646cffaa);
+}
+.logo.react:hover {
+  filter: drop-shadow(0 0 2em #61dafbaa);
+}
+
+@keyframes logo-spin {
+  from {
+    transform: rotate(0deg);
+  }
+  to {
+    transform: rotate(360deg);
+  }
+}
+
+@media (prefers-reduced-motion: no-preference) {
+  a:nth-of-type(2) .logo {
+    animation: logo-spin infinite 20s linear;
+  }
+}
+
+.card {
+  padding: 2em;
+}
+
+.read-the-docs {
+  color: #888;
+}

frontend/src/App.jsx

@@ -0,0 +1,35 @@
+import { useState } from 'react'
+import reactLogo from './assets/react.svg'
+import viteLogo from '/vite.svg'
+import './App.css'
+
+function App() {
+  const [count, setCount] = useState(0)
+
+  return (
+    <>
+      <div>
+        <a href="https://vite.dev" target="_blank">
+          <img src={viteLogo} className="logo" alt="Vite logo" />
+        </a>
+        <a href="https://react.dev" target="_blank">
+          <img src={reactLogo} className="logo react" alt="React logo" />
+        </a>
+      </div>
+      <h1>Vite + React</h1>
+      <div className="card">
+        <button onClick={() => setCount((count) => count + 1)}>
+          count is {count}
+        </button>
+        <p>
+          Edit <code>src/App.jsx</code> and save to test HMR
+        </p>
+      </div>
+      <p className="read-the-docs">
+        Click on the Vite and React logos to learn more
+      </p>
+    </>
+  )
+}
+
+export default App

frontend/src/index.css

@@ -0,0 +1,68 @@
+:root {
+  font-family: system-ui, Avenir, Helvetica, Arial, sans-serif;
+  line-height: 1.5;
+  font-weight: 400;
+
+  color-scheme: light dark;
+  color: rgba(255, 255, 255, 0.87);
+  background-color: #242424;
+
+  font-synthesis: none;
+  text-rendering: optimizeLegibility;
+  -webkit-font-smoothing: antialiased;
+  -moz-osx-font-smoothing: grayscale;
+}
+
+a {
+  font-weight: 500;
+  color: #646cff;
+  text-decoration: inherit;
+}
+a:hover {
+  color: #535bf2;
+}
+
+body {
+  margin: 0;
+  display: flex;
+  place-items: center;
+  min-width: 320px;
+  min-height: 100vh;
+}
+
+h1 {
+  font-size: 3.2em;
+  line-height: 1.1;
+}
+
+button {
+  border-radius: 8px;
+  border: 1px solid transparent;
+  padding: 0.6em 1.2em;
+  font-size: 1em;
+  font-weight: 500;
+  font-family: inherit;
+  background-color: #1a1a1a;
+  cursor: pointer;
+  transition: border-color 0.25s;
+}
+button:hover {
+  border-color: #646cff;
+}
+button:focus,
+button:focus-visible {
+  outline: 4px auto -webkit-focus-ring-color;
+}
+
+@media (prefers-color-scheme: light) {
+  :root {
+    color: #213547;
+    background-color: #ffffff;
+  }
+  a:hover {
+    color: #747bff;
+  }
+  button {
+    background-color: #f9f9f9;
+  }
+}

frontend/src/main.jsx

@@ -0,0 +1,10 @@
+import { StrictMode } from 'react'
+import { createRoot } from 'react-dom/client'
+import './index.css'
+import App from './App.jsx'
+
+createRoot(document.getElementById('root')).render(
+  <StrictMode>
+    <App />
+  </StrictMode>,
+)

frontend/vite.config.js

@@ -0,0 +1,8 @@
+import { defineConfig } from "vite";
+import react from "@vitejs/plugin-react";
+import tailwindcss from "@tailwindcss/vite";
+
+// https://vite.dev/config/
+export default defineConfig({
+  plugins: [react(), tailwindcss()],
+});

main.py

@@ -0,0 +1,8 @@
+import time
+from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
+
+start = time.time()
+model = AutoModelForCausalLM.from_pretrained("tcmmichaelb139/gpt2-medium-tinystories")
+print(model)
+
+print("Loaded model in", time.time() - start)

pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "evolutiontransformer"
+version = "0.1.0"
+description = "Simulating evolution among LLMs"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "accelerate>=1.10.1",
+    "celery>=5.5.3",
+    "datasets>=4.1.1",
+    "evaluate>=0.4.6",
+    "fastapi>=0.116.2",
+    "gradio>=5.46.0",
+    "gunicorn>=23.0.0",
+    "matplotlib>=3.10.6",
+    "numpy>=2.3.3",
+    "pytest>=8.4.2",
+    "redis>=6.4.0",
+    "torch>=2.8.0",
+    "transformers>=4.56.1",
+    "uvicorn[standard]>=0.35.0",
+]

tests/test_api.py

@@ -0,0 +1,310 @@
+import pytest
+from fastapi.testclient import TestClient
+import time
+import re
+
+from evolutiontransformer.api import app
+
+
+def get_final_answer(text: str) -> int | None:
+    numbers = re.findall(r"\d+", text)
+    return int(numbers[-1]) if numbers else None
+
+
+@pytest.fixture
+def client():
+    with TestClient(app) as c:
+        yield c
+
+
+def await_task_completion(client, task_id, timeout=60):
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        status_response = client.get(f"/tasks/{task_id}")
+
+        print(status_response.json())
+
+        if status_response.status_code == 500:
+            return {"error": status_response.json().get("detail", "Unknown error")}
+        assert status_response.status_code == 200
+        status_data = status_response.json()
+
+        if status_data["status"] == "SUCCESS":
+            return status_data["result"]
+
+        time.sleep(2)
+    else:
+        pytest.fail(
+            f"Task {task_id} did not complete within the {timeout}-second timeout."
+        )
+
+    return None
+
+
+def test_generate_endpoint_svamp(client):
+    """
+    Tests inference on svamp
+    """
+    response = client.post(
+        "/generate",
+        json={
+            "model_name": "svamp",
+            "prompt": "A spider has 8 legs. A fly has 6 legs. How many legs do they have in total?\nAnswer:",
+            "max_new_tokens": 50,
+            "temperature": 0.7,
+        },
+    )
+
+    assert response.status_code == 200
+    data = response.json()
+
+    assert "task_id" in data
+    task_id = data["task_id"]
+
+    final_result = await_task_completion(client, task_id)
+
+    assert "response" in final_result
+    output_text = final_result["response"]
+
+    answer = get_final_answer(output_text)
+    assert answer == 14
+
+
+def test_merge_then_inference_svamp_1(client):
+    """
+    Tests merging then inference for svamp dataset
+    """
+
+    merge_response = client.post(
+        "/merge",
+        json={
+            "model1_name": "svamp",
+            "model2_name": "tinystories",
+            "layer_recipe": [[(i, 0, 1.0)] for i in range(24)],
+            "embedding_lambdas": [1.0, 1.0],
+            "linear_lambdas": [1.0, 1.0],
+            "merged_name": "svamp_merged",
+        },
+    )
+
+    assert merge_response.status_code == 200
+    merge_data = merge_response.json()
+    assert "task_id" in merge_data
+    merge_task_id = merge_data["task_id"]
+
+    merge_status_data = await_task_completion(client, merge_task_id)
+    model_name = merge_status_data["response"]
+
+    time.sleep(5)
+
+    generate_response = client.post(
+        "/generate",
+        json={
+            "model_name": model_name,
+            "prompt": "A spider has 8 legs. A fly has 6 legs. How many legs do they have in total?\nAnswer:",
+            "max_new_tokens": 50,
+            "temperature": 0.7,
+        },
+    )
+
+    assert generate_response.status_code == 200
+    generate_data = generate_response.json()
+    assert "task_id" in generate_data
+    generate_task_id = generate_data["task_id"]
+
+    final_result = await_task_completion(client, generate_task_id)
+
+    assert "response" in final_result
+    output_text = final_result["response"]
+    answer = get_final_answer(output_text)
+
+    assert answer == 14
+
+
+def test_merge_then_inference_svamp_2(client):
+    """
+    Tests merging then inference for svamp dataset
+    """
+
+    merge_repsonse = client.post(
+        "/merge",
+        json={
+            "model1_name": "svamp",
+            "model2_name": "tinystories",
+            "layer_recipe": [[(i % 24, 0, 1.0 if i < 24 else 0.5)] for i in range(48)],
+            "embedding_lambdas": [1.0, 1.0],
+            "linear_lambdas": [1.0, 1.0],
+            "merged_name": "svamp_merged",
+        },
+    )
+
+    assert merge_repsonse.status_code == 200
+    merge_data = merge_repsonse.json()
+    assert "task_id" in merge_data
+    merge_task_id = merge_data["task_id"]
+
+    merge_status_data = await_task_completion(client, merge_task_id)
+
+    model_name = merge_status_data["response"]
+
+    merge_response2 = client.post(
+        "/merge",
+        json={
+            "model1_name": model_name,
+            "model2_name": "tinystories",
+            "layer_recipe": [[(i, 1, 0.25)] for i in range(24)],
+            "embedding_lambdas": [0.0, 0.0],
+            "linear_lambdas": [0.0, 0.0],
+            "merged_name": "svamp_merged",
+        },
+    )
+
+    assert merge_response2.status_code == 200
+    merge_data2 = merge_response2.json()
+    assert "task_id" in merge_data2
+    merge_task_id2 = merge_data2["task_id"]
+    merge_status_data2 = await_task_completion(client, merge_task_id2)
+    model_name2 = merge_status_data2["response"]
+
+    time.sleep(5)
+
+    generate_response = client.post(
+        "/generate",
+        json={
+            "model_name": model_name2,
+            "prompt": "A spider has 8 legs. A fly has 6 legs. How many legs do they have in total?\nAnswer:",
+            "max_new_tokens": 50,
+            "temperature": 0.7,
+        },
+    )
+
+    assert generate_response.status_code == 200
+    generate_data = generate_response.json()
+    assert "task_id" in generate_data
+    generate_task_id = generate_data["task_id"]
+
+    final_result = await_task_completion(client, generate_task_id)
+
+    assert "response" in final_result
+    output_text = final_result["response"]
+    answer = get_final_answer(output_text)
+
+    assert answer == 14
+
+
+def test_merge_two_children_then_merge(client):
+    """
+    Tests creating two children and merging them
+    """
+
+    merge_response1 = client.post(
+        "/merge",
+        json={
+            "model1_name": "svamp",
+            "model2_name": "tinystories",
+            "layer_recipe": [[(i, 0, 0.8)] for i in range(12)]
+            + [[(i, 1, 0.6)] for i in range(12)],
+            "embedding_lambdas": [0.7, 0.3],
+            "linear_lambdas": [0.8, 0.2],
+            "merged_name": "child1",
+        },
+    )
+
+    assert merge_response1.status_code == 200
+    merge_data1 = merge_response1.json()
+    assert "task_id" in merge_data1
+    merge_task_id1 = merge_data1["task_id"]
+    merge_status_data1 = await_task_completion(client, merge_task_id1)
+    child1_name = merge_status_data1["response"]
+
+    merge_response2 = client.post(
+        "/merge",
+        json={
+            "model1_name": "svamp",
+            "model2_name": "tinystories",
+            "layer_recipe": [[(i, 1, 0.9)] for i in range(8)]
+            + [[(i, 0, 0.4)] for i in range(16)],
+            "embedding_lambdas": [0.2, 0.9],
+            "linear_lambdas": [0.3, 0.7],
+            "merged_name": "child2",
+        },
+    )
+
+    assert merge_response2.status_code == 200
+    merge_data2 = merge_response2.json()
+    assert "task_id" in merge_data2
+    merge_task_id2 = merge_data2["task_id"]
+    merge_status_data2 = await_task_completion(client, merge_task_id2)
+    child2_name = merge_status_data2["response"]
+
+    merge_response3 = client.post(
+        "/merge",
+        json={
+            "model1_name": child1_name,
+            "model2_name": child2_name,
+            "layer_recipe": [[(i, 0, 0.6), (i, 1, 0.4)] for i in range(24)],
+            "embedding_lambdas": [0.5, 0.5],
+            "linear_lambdas": [0.6, 0.4],
+            "merged_name": "final_merged",
+        },
+    )
+
+    assert merge_response3.status_code == 200
+    merge_data3 = merge_response3.json()
+    assert "task_id" in merge_data3
+    merge_task_id3 = merge_data3["task_id"]
+    merge_status_data3 = await_task_completion(client, merge_task_id3)
+    final_model_name = merge_status_data3["response"]
+
+    time.sleep(5)
+
+    generate_response = client.post(
+        "/generate",
+        json={
+            "model_name": final_model_name,
+            "prompt": "A spider has 8 legs. A fly has 6 legs. How many legs do they have in total?\nAnswer:",
+            "max_new_tokens": 50,
+            "temperature": 0.7,
+        },
+    )
+
+    assert generate_response.status_code == 200
+    generate_data = generate_response.json()
+    assert "task_id" in generate_data
+    generate_task_id = generate_data["task_id"]
+
+    final_result = await_task_completion(client, generate_task_id)
+
+    assert "response" in final_result
+    output_text = final_result["response"]
+    answer = get_final_answer(output_text)
+
+    assert answer == 14
+
+
+def test_merge_fail(client):
+    """
+    Tests merging with too many layers
+    """
+
+    merge_repsonse = client.post(
+        "/merge",
+        json={
+            "model1_name": "svamp",
+            "model2_name": "tinystories",
+            "layer_recipe": [[(i, 0, 1.0)] for i in range(50)],
+            "embedding_lambdas": [1.0, 1.0],
+            "linear_lambdas": [1.0, 1.0],
+            "merged_name": "svamp_merged",
+        },
+    )
+
+    assert merge_repsonse.status_code == 200
+    merge_data = merge_repsonse.json()
+    assert "task_id" in merge_data
+    merge_task_id = merge_data["task_id"]
+
+    merge_status_data = await_task_completion(client, merge_task_id)
+    assert "response" not in merge_status_data
+    assert "error" in merge_status_data
+    assert "Layer recipe too long" in merge_status_data["error"]

tests/test_model_actions.py

@@ -0,0 +1,83 @@
+import torch
+from transformers import AutoModelForCausalLM
+import re
+
+
+from evolutiontransformer.worker import (
+    load_base_models_if_needed,
+    BASE_MODELS,
+    inference,
+    inference_task,
+    merge_models,
+)
+
+
+def get_final_answer(text: str) -> int | None:
+    numbers = re.findall(r"\d+", text)
+    return int(numbers[-1]) if numbers else None
+
+
+def test_inference():
+    session_id = "test_session"
+
+    print("### Testing inference on SVAMP model...")
+    prompt = "If there are 3 cars and 2 bikes, how many vehicles are there in total?\nAnswer:"
+    output = inference_task(session_id, "svamp", prompt)
+    assert get_final_answer(output["response"]) == 5
+
+
+def test_merge_models():
+    load_base_models_if_needed()
+
+    model_recipe = {
+        "layer_recipe": [[(i, "svamp", 1.0)] for i in range(24)],
+        "embedding_lambdas": [1.0, 1.0],
+        "linear_lambdas": [1.0, 1.0],
+    }
+
+    merged_model = merge_models(model_recipe)
+
+    for (name1, param1), (name2, param2) in zip(
+        BASE_MODELS["svamp"].named_parameters(), merged_model.named_parameters()
+    ):
+        assert torch.allclose(param1, param2)
+
+
+def test_merge_models_with_inference1():
+    load_base_models_if_needed()
+
+    model_recipe = {
+        "layer_recipe": [
+            [(i % 24, "svamp", 1.0 if i < 24 else 0.5)] for i in range(48)
+        ],
+        "embedding_lambdas": [1.0, 1.0],
+        "linear_lambdas": [1.0, 1.0],
+    }
+
+    merged_model = merge_models(model_recipe)
+
+    print(
+        inference(
+            merged_model,
+            "A spider has 8 legs. A fly has 6 legs. How many legs do they have in total?\nAnswer:",
+        )
+    )
+
+
+def test_merge_models_with_inference2():
+    load_base_models_if_needed()
+
+    model_recipe = {
+        "layer_recipe": [[(i, "tinystories", 1.0)] for i in range(24)],
+        "embedding_lambdas": [0.0, 0.0],
+        "linear_lambdas": [0.0, 0.0],
+    }
+
+    merged_model = merge_models(model_recipe)
+
+    print(
+        inference(
+            merged_model,
+            "A spider has 8 legs. A fly has 6 legs. How many legs do they have in total?\nAnswer:",
+        )
+    )