rough frontend + fix backend

.gitignore

@@ -6,6 +6,8 @@ dist/
 wheels/
 *.egg-info
 
+.pytest_cache/
+
 .env
 .venv
 

README.md

@@ -6,3 +6,5 @@ colorTo: pink
 sdk: docker
 app_port: 7860
 ---
+
+Uses Huggingface Spaces for the backend, Upstart for the database, and Vercel for the frontend.

evolutiontransformer/worker.py

@@ -12,6 +12,7 @@ from evolutiontransformer.redis import (
     get_session_models,
     save_model_recipe,
     get_model_recipe,
+    delete_session,
 )
 
 
@@ -256,16 +257,13 @@ def merge_models_task(
 
 @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()),
-    }
+    base_models = BASE_MODELS_NAMES
+    session_models = get_session_models(session_id)
+    all_models = list(set(base_models + session_models))
+    return {"response": all_models}
 
 
 @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"}
+    delete_session(session_id)
+    return {"response": ""}

frontend/src/App.css

@@ -1,44 +1,39 @@
 @import "tailwindcss";
 
-#root {
-  max-width: 1280px;
-  margin: 0 auto;
-  padding: 2rem;
-  text-align: center;
+input[type="range"] {
+  -webkit-appearance: none;
+  appearance: none;
+  height: 6px;
+  border-radius: 3px;
+  background: #e2e8f0;
+  outline: none;
 }
 
-.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);
-  }
+input[type="range"]::-webkit-slider-thumb {
+  -webkit-appearance: none;
+  appearance: none;
+  width: 18px;
+  height: 18px;
+  border-radius: 50%;
+  background: #0ea5e9;
+  cursor: pointer;
+  box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
 }
 
-@media (prefers-reduced-motion: no-preference) {
-  a:nth-of-type(2) .logo {
-    animation: logo-spin infinite 20s linear;
-  }
+input[type="range"]::-moz-range-thumb {
+  width: 18px;
+  height: 18px;
+  border-radius: 50%;
+  background: #0ea5e9;
+  cursor: pointer;
+  border: none;
+  box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
 }
 
-.card {
-  padding: 2em;
+input[type="range"].accent-primary-600::-webkit-slider-thumb {
+  background: #0284c7;
 }
 
-.read-the-docs {
-  color: #888;
+input[type="range"].accent-primary-600::-moz-range-thumb {
+  background: #0284c7;
 }

frontend/src/App.jsx

@@ -1,35 +1,132 @@
-import { useState } from 'react'
-import reactLogo from './assets/react.svg'
-import viteLogo from '/vite.svg'
-import './App.css'
+import { useState, useEffect } from "react";
+import "./App.css";
+import Header from "./components/Header";
+import ModelSelection from "./components/ModelSelection";
+import GlobalControls from "./components/GlobalControls";
+import LayerRecipe from "./components/LayerRecipe";
+import { useAPI } from "./hooks/useAPI";
 
 function App() {
-  const [count, setCount] = useState(0)
+  const [models, setModels] = useState([]);
+  const [selectedModel1, setSelectedModel1] = useState("");
+  const [selectedModel2, setSelectedModel2] = useState("");
+  const [layerRecipe, setLayerRecipe] = useState([]);
+  const [embeddingLambdas, setEmbeddingLambdas] = useState([0.5, 0.5]);
+  const [linearLambdas, setLinearLambdas] = useState([0.5, 0.5]);
+  const [mergedName, setMergedName] = useState("merged");
+  const [isLoading, setIsLoading] = useState(false);
+  const [result, setResult] = useState(null);
+  const [numLayers, setNumLayers] = useState(12);
+
+  const { checkTaskStatus, fetchModels, mergeModels } = useAPI();
+
+  useEffect(() => {
+    const loadModels = async () => {
+      const taskId = await fetchModels();
+      if (taskId) {
+        setTimeout(
+          () =>
+            checkTaskStatus(taskId, (result) => {
+              if (result && Array.isArray(result)) {
+                setModels(result);
+              }
+            }),
+          1000
+        );
+      }
+    };
+
+    loadModels();
+  }, [fetchModels, checkTaskStatus]);
+
+  const initializeLayerRecipe = () => {
+    const recipe = [];
+    for (let i = 0; i < numLayers; i++) {
+      recipe.push([[i, i, 0.5]]);
+    }
+    setLayerRecipe(recipe);
+  };
+
+  const updateLayerWeight = (layerIndex, weight) => {
+    const newRecipe = [...layerRecipe];
+    if (!newRecipe[layerIndex]) {
+      newRecipe[layerIndex] = [[layerIndex, layerIndex, weight]];
+    } else {
+      newRecipe[layerIndex][0][2] = weight;
+    }
+    setLayerRecipe(newRecipe);
+  };
+
+  const handleMerge = async () => {
+    if (!selectedModel1 || !selectedModel2 || layerRecipe.length === 0) {
+      alert("Please select both models and configure layer recipe");
+      return;
+    }
+
+    setIsLoading(true);
+    setResult(null);
+
+    const mergeData = {
+      model1_name: selectedModel1,
+      model2_name: selectedModel2,
+      layer_recipe: layerRecipe,
+      embedding_lambdas: embeddingLambdas,
+      linear_lambdas: linearLambdas,
+      merged_name: mergedName,
+    };
+
+    const taskId = await mergeModels(mergeData);
+    if (taskId) {
+      checkTaskStatus(taskId, (result) => {
+        setResult(result);
+        setIsLoading(false);
+      });
+    } else {
+      setIsLoading(false);
+    }
+  };
+
+  const isMergeDisabled =
+    isLoading || !selectedModel1 || !selectedModel2 || layerRecipe.length === 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 className="min-h-screen bg-gradient-to-br from-primary-50 to-secondary-50 p-6">
+      <div className="max-w-6xl mx-auto">
+        <Header />
+
+        <div className="grid lg:grid-cols-2 gap-8">
+          <ModelSelection
+            models={models}
+            selectedModel1={selectedModel1}
+            setSelectedModel1={setSelectedModel1}
+            selectedModel2={selectedModel2}
+            setSelectedModel2={setSelectedModel2}
+            numLayers={numLayers}
+            setNumLayers={setNumLayers}
+            mergedName={mergedName}
+            setMergedName={setMergedName}
+            onInitializeRecipe={initializeLayerRecipe}
+          />
+
+          <GlobalControls
+            embeddingLambdas={embeddingLambdas}
+            setEmbeddingLambdas={setEmbeddingLambdas}
+            linearLambdas={linearLambdas}
+            setLinearLambdas={setLinearLambdas}
+            onMerge={handleMerge}
+            isLoading={isLoading}
+            isDisabled={isMergeDisabled}
+            result={result}
+          />
+        </div>
+
+        <LayerRecipe
+          layerRecipe={layerRecipe}
+          onUpdateLayerWeight={updateLayerWeight}
+        />
       </div>
-      <p className="read-the-docs">
-        Click on the Vite and React logos to learn more
-      </p>
-    </>
-  )
+    </div>
+  );
 }
 
-export default App
+export default App;

frontend/src/components/GlobalControls.jsx

@@ -0,0 +1,135 @@
+const GlobalControls = ({
+  embeddingLambdas,
+  setEmbeddingLambdas,
+  linearLambdas,
+  setLinearLambdas,
+  onMerge,
+  isLoading,
+  isDisabled,
+  result,
+}) => {
+  return (
+    <div className="bg-white rounded-2xl shadow-lg p-8 border border-secondary-100">
+      <h2 className="text-2xl font-semibold text-secondary-800 mb-6">
+        Global Controls
+      </h2>
+
+      <div className="space-y-6">
+        <div>
+          <label className="block text-sm font-medium text-secondary-700 mb-2">
+            Embedding Weights
+          </label>
+          <div className="grid grid-cols-2 gap-4">
+            <div>
+              <span className="text-xs text-secondary-600">Model 1</span>
+              <input
+                type="range"
+                min="0"
+                max="1"
+                step="0.05"
+                value={embeddingLambdas[0]}
+                onChange={(e) =>
+                  setEmbeddingLambdas([
+                    parseFloat(e.target.value),
+                    1 - parseFloat(e.target.value),
+                  ])
+                }
+                className="w-full"
+              />
+              <span className="text-sm text-secondary-600">
+                {embeddingLambdas[0].toFixed(2)}
+              </span>
+            </div>
+            <div>
+              <span className="text-xs text-secondary-600">Model 2</span>
+              <input
+                type="range"
+                min="0"
+                max="1"
+                step="0.05"
+                value={embeddingLambdas[1]}
+                onChange={(e) =>
+                  setEmbeddingLambdas([
+                    1 - parseFloat(e.target.value),
+                    parseFloat(e.target.value),
+                  ])
+                }
+                className="w-full"
+              />
+              <span className="text-sm text-secondary-600">
+                {embeddingLambdas[1].toFixed(2)}
+              </span>
+            </div>
+          </div>
+        </div>
+
+        <div>
+          <label className="block text-sm font-medium text-secondary-700 mb-2">
+            Linear Weights
+          </label>
+          <div className="grid grid-cols-2 gap-4">
+            <div>
+              <span className="text-xs text-secondary-600">Model 1</span>
+              <input
+                type="range"
+                min="0"
+                max="1"
+                step="0.05"
+                value={linearLambdas[0]}
+                onChange={(e) =>
+                  setLinearLambdas([
+                    parseFloat(e.target.value),
+                    1 - parseFloat(e.target.value),
+                  ])
+                }
+                className="w-full"
+              />
+              <span className="text-sm text-secondary-600">
+                {linearLambdas[0].toFixed(2)}
+              </span>
+            </div>
+            <div>
+              <span className="text-xs text-secondary-600">Model 2</span>
+              <input
+                type="range"
+                min="0"
+                max="1"
+                step="0.05"
+                value={linearLambdas[1]}
+                onChange={(e) =>
+                  setLinearLambdas([
+                    1 - parseFloat(e.target.value),
+                    parseFloat(e.target.value),
+                  ])
+                }
+                className="w-full"
+              />
+              <span className="text-sm text-secondary-600">
+                {linearLambdas[1].toFixed(2)}
+              </span>
+            </div>
+          </div>
+        </div>
+
+        <button
+          onClick={onMerge}
+          disabled={isDisabled}
+          className="w-full py-4 bg-gradient-to-r from-primary-600 to-accent-600 text-white font-semibold rounded-lg hover:from-primary-700 hover:to-accent-700 transition-all disabled:opacity-50 disabled:cursor-not-allowed"
+        >
+          {isLoading ? "Merging Models..." : "Merge Models"}
+        </button>
+
+        {result && (
+          <div className="mt-4 p-4 bg-accent-50 border border-accent-200 rounded-lg">
+            <h3 className="font-medium text-accent-800 mb-2">
+              Merge Complete!
+            </h3>
+            <p className="text-sm text-accent-700">{JSON.stringify(result)}</p>
+          </div>
+        )}
+      </div>
+    </div>
+  );
+};
+
+export default GlobalControls;

frontend/src/components/Header.jsx

@@ -0,0 +1,15 @@
+const Header = () => {
+  return (
+    <header className="text-center mb-12">
+      <h1 className="text-5xl font-bold text-secondary-800 mb-4">
+        Evolution Transformer
+      </h1>
+      <p className="text-xl text-secondary-600 max-w-2xl mx-auto">
+        Combine and evolve transformer models by blending layers with precise
+        control
+      </p>
+    </header>
+  );
+};
+
+export default Header;

frontend/src/components/LayerRecipe.jsx

@@ -0,0 +1,46 @@
+const LayerRecipe = ({ layerRecipe, onUpdateLayerWeight }) => {
+  if (layerRecipe.length === 0) return null;
+
+  return (
+    <div className="mt-8 bg-white rounded-2xl shadow-lg p-8 border border-secondary-100">
+      <h2 className="text-2xl font-semibold text-secondary-800 mb-6">
+        Layer Recipe
+      </h2>
+
+      <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 xl:grid-cols-4 gap-4">
+        {layerRecipe.map((layer, index) => (
+          <div
+            key={index}
+            className="bg-secondary-50 rounded-lg p-4 border border-secondary-200"
+          >
+            <div className="text-sm font-medium text-secondary-700 mb-2">
+              Layer {index}
+            </div>
+            <div className="space-y-2">
+              <input
+                type="range"
+                min="0"
+                max="1"
+                step="0.05"
+                value={layer[0]?.[2] || 0.5}
+                onChange={(e) =>
+                  onUpdateLayerWeight(index, parseFloat(e.target.value))
+                }
+                className="w-full accent-primary-600"
+              />
+              <div className="flex justify-between text-xs text-secondary-600">
+                <span>Model 1</span>
+                <span className="font-medium">
+                  {(layer[0]?.[2] || 0.5).toFixed(2)}
+                </span>
+                <span>Model 2</span>
+              </div>
+            </div>
+          </div>
+        ))}
+      </div>
+    </div>
+  );
+};
+
+export default LayerRecipe;

frontend/src/components/ModelSelection.jsx

@@ -0,0 +1,93 @@
+const ModelSelection = ({
+  models,
+  selectedModel1,
+  setSelectedModel1,
+  selectedModel2,
+  setSelectedModel2,
+  numLayers,
+  setNumLayers,
+  mergedName,
+  setMergedName,
+  onInitializeRecipe,
+}) => {
+  return (
+    <div className="bg-white rounded-2xl shadow-lg p-8 border border-secondary-100">
+      <h2 className="text-2xl font-semibold text-secondary-800 mb-6">
+        Model Selection
+      </h2>
+
+      <div className="space-y-6">
+        <div>
+          <label className="block text-sm font-medium text-secondary-700 mb-2">
+            Base Model (Model 1)
+          </label>
+          <select
+            value={selectedModel1}
+            onChange={(e) => setSelectedModel1(e.target.value)}
+            className="w-full p-3 border border-secondary-300 rounded-lg focus:ring-2 focus:ring-primary-500 focus:border-transparent"
+          >
+            <option value="">Select a model</option>
+            {models.map((model) => (
+              <option key={model} value={model}>
+                {model}
+              </option>
+            ))}
+          </select>
+        </div>
+
+        <div>
+          <label className="block text-sm font-medium text-secondary-700 mb-2">
+            Target Model (Model 2)
+          </label>
+          <select
+            value={selectedModel2}
+            onChange={(e) => setSelectedModel2(e.target.value)}
+            className="w-full p-3 border border-secondary-300 rounded-lg focus:ring-2 focus:ring-primary-500 focus:border-transparent"
+          >
+            <option value="">Select a model</option>
+            {models.map((model) => (
+              <option key={model} value={model}>
+                {model}
+              </option>
+            ))}
+          </select>
+        </div>
+
+        <div>
+          <label className="block text-sm font-medium text-secondary-700 mb-2">
+            Number of Layers
+          </label>
+          <input
+            type="number"
+            value={numLayers}
+            onChange={(e) => setNumLayers(parseInt(e.target.value))}
+            className="w-full p-3 border border-secondary-300 rounded-lg focus:ring-2 focus:ring-primary-500 focus:border-transparent"
+            min="1"
+            max="48"
+          />
+          <button
+            onClick={onInitializeRecipe}
+            className="mt-2 px-4 py-2 bg-secondary-600 text-white rounded-lg hover:bg-secondary-700 transition-colors"
+          >
+            Initialize Recipe
+          </button>
+        </div>
+
+        <div>
+          <label className="block text-sm font-medium text-secondary-700 mb-2">
+            Merged Model Name
+          </label>
+          <input
+            type="text"
+            value={mergedName}
+            onChange={(e) => setMergedName(e.target.value)}
+            className="w-full p-3 border border-secondary-300 rounded-lg focus:ring-2 focus:ring-primary-500 focus:border-transparent"
+            placeholder="Enter merged model name"
+          />
+        </div>
+      </div>
+    </div>
+  );
+};
+
+export default ModelSelection;

frontend/src/hooks/useAPI.js

@@ -0,0 +1,55 @@
+import { useCallback } from "react";
+
+const API_BASE = "https://tcmmichaelb139-evolutiontransformer.hf.space";
+
+export const useAPI = () => {
+  const checkTaskStatus = useCallback(async (taskId, callback) => {
+    try {
+      const response = await fetch(`${API_BASE}/tasks/${taskId}`);
+      const data = await response.json();
+
+      if (data.status === "SUCCESS") {
+        callback(data.result);
+      } else if (data.status === "PENDING") {
+        setTimeout(() => checkTaskStatus(taskId, callback), 1000);
+      }
+    } catch (error) {
+      console.error("Task check failed:", error);
+    }
+  }, []);
+
+  const fetchModels = useCallback(async () => {
+    try {
+      const response = await fetch(`${API_BASE}/list_models`, {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+      });
+      const data = await response.json();
+      return data.task_id;
+    } catch (error) {
+      console.error("Failed to fetch models:", error);
+      return null;
+    }
+  }, []);
+
+  const mergeModels = useCallback(async (mergeData) => {
+    try {
+      const response = await fetch(`${API_BASE}/merge`, {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify(mergeData),
+      });
+      const data = await response.json();
+      return data.task_id;
+    } catch (error) {
+      console.error("Merge failed:", error);
+      return null;
+    }
+  }, []);
+
+  return {
+    checkTaskStatus,
+    fetchModels,
+    mergeModels,
+  };
+};

frontend/src/index.css

@@ -1,68 +1,19 @@
+@import "tailwindcss";
+
 :root {
-  font-family: system-ui, Avenir, Helvetica, Arial, sans-serif;
+  font-family: system-ui, -apple-system, "Segoe UI", Roboto, Oxygen, Ubuntu,
+    Cantarell, sans-serif;
   line-height: 1.5;
   font-weight: 400;
-
-  color-scheme: light dark;
-  color: rgba(255, 255, 255, 0.87);
-  background-color: #242424;
-
+  color-scheme: light;
   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;
+  padding: 0;
   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/tailwind.config.js

@@ -0,0 +1,47 @@
+/** @type {import('tailwindcss').Config} */
+export default {
+  content: ["./index.html", "./src/**/*.{js,ts,jsx,tsx}"],
+  theme: {
+    extend: {
+      colors: {
+        primary: {
+          50: "#f0f9ff",
+          100: "#e0f2fe",
+          200: "#bae6fd",
+          300: "#7dd3fc",
+          400: "#38bdf8",
+          500: "#0ea5e9",
+          600: "#0284c7",
+          700: "#0369a1",
+          800: "#075985",
+          900: "#0c4a6e",
+        },
+        secondary: {
+          50: "#f8fafc",
+          100: "#f1f5f9",
+          200: "#e2e8f0",
+          300: "#cbd5e1",
+          400: "#94a3b8",
+          500: "#64748b",
+          600: "#475569",
+          700: "#334155",
+          800: "#1e293b",
+          900: "#0f172a",
+        },
+        accent: {
+          50: "#f0fdf4",
+          100: "#dcfce7",
+          200: "#bbf7d0",
+          300: "#86efac",
+          400: "#4ade80",
+          500: "#22c55e",
+          600: "#16a34a",
+          700: "#15803d",
+          800: "#166534",
+          900: "#14532d",
+        },
+      },
+    },
+  },
+  plugins: [],
+};

tests/test_hf_api.py

@@ -279,6 +279,20 @@ def test_merge_two_children_then_merge(session):
     output_text = final_result["response"]
     answer = get_final_answer(output_text)
 
+    number_of_models = session.post(f"{BASE_URL}/list_models")
+
+    assert number_of_models.status_code == 200
+    number_of_models_data = number_of_models.json()
+    assert "task_id" in number_of_models_data
+    number_of_models_task_id = number_of_models_data["task_id"]
+
+    number_of_models_result = await_task_completion(session, number_of_models_task_id)
+
+    assert "response" in number_of_models_result
+    models = number_of_models_result["response"]
+    print(models)
+    assert len(models) >= 5
+
     assert answer == 14