adaptible 1.0.0a1__tar.gz

adaptible-1.0.0a1/.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+uv.lock
+.venv
+.gitignore

adaptible-1.0.0a1/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: adaptible
+Version: 1.0.0a1
+Summary: Stateful LLM serving instances that self-reflect and learn from their mistakes during periods of low server utilization.
+Requires-Python: >=3.13
+Requires-Dist: absl-py
+Requires-Dist: fastapi[all]
+Requires-Dist: immutabledict
+Requires-Dist: lm-eval
+Requires-Dist: mlx
+Requires-Dist: mlx-lm
+Requires-Dist: optax
+Requires-Dist: torch
+Requires-Dist: tqdm
+Requires-Dist: transformers
+Requires-Dist: uvicorn[standard]
+Requires-Dist: vizible

adaptible-1.0.0a1/adaptible/__init__.py

@@ -0,0 +1,3 @@
+"""Adaptible - LLMs that can wander."""
+
+from ._src import app, InteractionHistory, StatefulLLM, MutableHostedLLM

adaptible-1.0.0a1/adaptible/_src/__init__.py

@@ -0,0 +1,6 @@
+"""Adaptible - LLMs that can wander."""
+
+from ._api import app
+from ._classes import InteractionHistory
+from ._llm import StatefulLLM
+from ._server import MutableHostedLLM

adaptible-1.0.0a1/adaptible/_src/_api.py

@@ -0,0 +1,142 @@
+"""Standard Interaction Logic for a Stateful LLM"""
+
+import asyncio
+from asyncio import log
+import collections
+import os
+import time
+import threading
+from typing import Any, List
+
+
+from fastapi import FastAPI, HTTPException
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import StreamingResponse
+import tqdm
+import vizible
+
+from ._classes import (
+    InteractionHistory,
+    InteractionRequest,
+    InteractionResponse,
+    ReviewResponse,
+    SyncResponse,
+)
+from ._llm import StatefulLLM
+
+app = FastAPI(
+    title="Stateful Self-Improving LLM Server",
+    description="An API for a stateful LLM that tries to improve itself over time.",
+)
+app.mount(
+    "/static",
+    StaticFiles(
+        directory=os.path.join(os.path.dirname(__file__), "static"),
+        html=True,
+    ),
+    name="static",
+)
+
+# In-memory store or interaction history from the current session.
+interaction_history: List[InteractionHistory] = []
+unreviewed_interaction_history_indices: List[int] = []
+
+outstanding_tasks: collections.deque[Any] = collections.deque([])
+
+# Instantiate the stateful model. This is a singleton for the lifecycle of the app.
+model = StatefulLLM()
+
+
+@app.post("/interact", response_model=InteractionResponse)
+async def interact_with_model(request: InteractionRequest):
+    """
+    Main endpoint for interacting with the LLM (Forward-pass).
+    """
+    if not request.prompt:
+        raise HTTPException(status_code=400, detail="Prompt cannot be empty.")
+    print(f"Received request: {request}")
+    # Generate the response using the current state of the model.
+    response_text = model.generate_response(request.prompt)
+
+    # Store the interaction for later review
+    interaction_idx = len(interaction_history)
+    unreviewed_interaction_history_indices.append(len(interaction_history))
+    interaction_history.append(
+        InteractionHistory(
+            idx=interaction_idx,
+            user_input=request.prompt,
+            llm_response=response_text,
+            reviewed=False,
+            timestamp=time.time(),
+        )
+    )
+    return {"response": response_text, "interaction_idx": interaction_idx}
+
+
+@app.post("/stream_interact")
+async def stream_interact_with_model(request: InteractionRequest):
+    """Stream interaction with model."""
+    # TODO - enable logging.
+    return StreamingResponse(
+        model.stream_response(request.prompt), media_type="text/plain"
+    )
+
+@app.post("/trigger_review", response_model=ReviewResponse)
+async def trigger_review_cycle():
+    """Manually triggers the self-correction and training cycle."""
+    unreviewed_count = len(unreviewed_interaction_history_indices)
+    if unreviewed_count == 0:
+        return {
+            "message": "No unreviewed interactions to process.",
+            "unreviewed_count": unreviewed_count,
+        }
+    unreviewed_interaction_history = [
+        interaction_history[idx] for idx in unreviewed_interaction_history_indices
+    ]
+    outstanding_tasks.append(
+        asyncio.to_thread(model.self_correct_and_train, unreviewed_interaction_history)
+    )
+    return {
+        "message": "Self-correction and training cycle has been initiated in the background.",
+        "unreviewed_count": unreviewed_count,
+    }
+
+
+@app.get("/sync", response_model=SyncResponse)
+async def sync_server_for_background_tasks():
+    """Waits for any tasks to background complete."""
+    start_time = time.time()
+    num_tasks = len(outstanding_tasks)
+    print("Waiting for model state to stabilize...")
+    lock = threading.Lock()
+    with (
+        lock,
+        tqdm.tqdm(desc="Waiting for server to sync.", unit=" Seconds") as server_pbar,
+    ):
+        vizible.green("Finished background tasks")
+        while not model.ok:
+            log.logger.info(
+                "Waiting for model server to sync. Is model is %s ok.",
+                "" if model.ok else "not ",
+            )
+            server_pbar.update(1)
+            await asyncio.sleep(1)
+    vizible.green("Model has reached a stable state!!!")
+    elapsed_time = time.time() - start_time
+    return {
+        "message": "Sync'd all background tasks.",
+        "tasks_count": num_tasks,
+        "elapsed_time": elapsed_time,
+    }
+
+
+@app.get("/history")
+async def get_history():
+    """Returns the full interaction history."""
+    return {"history": interaction_history}
+
+
+@app.get("/status")
+async def check_is_running():
+    """Basic response to signal that the server is operational."""
+    return {"status": "up"}

adaptible-1.0.0a1/adaptible/_src/_classes.py

@@ -0,0 +1,80 @@
+"""Common class definitions."""
+
+import dataclasses
+import mlx.core as mx
+from pydantic import BaseModel
+
+
+@dataclasses.dataclass
+class InteractionHistory:
+    """Event turn during user-LLM dialog
+
+    Attributes:
+        idx: Index of current turn amongst all global turns.
+        user_input: User-provided prompt.
+        llm_response: LLM response.
+        reviewed: Whether this interaction has been reviewed already.
+        timestamp: When the interaction took place, measured in seconds.
+    """
+    idx: int
+    user_input: str
+    llm_response: str = ""
+    reviewed: bool = False
+    timestamp: float = 0.0
+
+@dataclasses.dataclass
+class TrainingExample:
+    """Pre-tokenized training data
+
+    Attributes:
+        user_input: Actual model response.
+        label: Target model response.
+        mask: Mask of model response that dictates which parts are used in training.
+    """
+    input: mx.array
+    label: mx.array
+    mask: mx.array
+
+
+class InteractionRequest(BaseModel):
+    """User prompt to be sent to the LLM.
+
+    Attributes:
+        prompt: User-provided input.
+    """
+    prompt: str
+
+
+class InteractionResponse(BaseModel):
+    """LLM response to user-provided prompt
+    
+    Attributes:
+        response: LLM-generated text response.
+        interaction_id: Index of the current response within the context of the current session.
+    """
+    response: str
+    interaction_idx: int
+
+
+class ReviewResponse(BaseModel):
+    """Response to initiating asynchronous review of entire unreviewed interaction history.
+
+    Attributes:
+        message: Human-readable output after completion of review.
+        unreviewed_count: the number of unreviewed interactions handled by this operation.
+    """
+    message: str
+    unreviewed_count: int
+
+
+class SyncResponse(BaseModel):
+    """Response after server completes all unfinished background tasks.
+
+    Attributes:
+        message: Human-readable message.
+        tasks_count: Number of tasks waited on and successfully finished.
+        elapsed_time: Amount of time background tasks took to finish.
+    """
+    message: str
+    tasks_count: int
+    elapsed_time: float

adaptible-1.0.0a1/adaptible/_src/_llm.py

@@ -0,0 +1,302 @@
+"""Stateful LLM."""
+
+from typing import AsyncIterable, List, Tuple
+
+import functools
+import threading
+import tqdm
+
+import immutabledict
+from mlx import optimizers
+from mlx import nn
+import mlx.core as mx
+from mlx_lm.generate import generate, stream_generate
+from mlx_lm.utils import load
+from mlx_lm.tuner.utils import linear_to_lora_layers
+from transformers.tokenization_utils import PreTrainedTokenizer
+import vizible
+
+from .libs import revise
+from ._classes import InteractionHistory, TrainingExample
+
+
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+#                         Default constants.                          #
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+_MODEL_NAME = "lmstudio-community/DeepSeek-R1-0528-Qwen3-8B-MLX-4bit"
+# _MODEL_NAME = "mlx-community/DeepSeek-R1-Distill-Qwen-1.5B-4bit"
+_MAX_TOKENS = 32768
+_LEARNING_RATE = 5e-5
+_EPOCHS = 5
+_NUM_LORA_LAYERS = 16
+_LORA_PARAMETERS = immutabledict.immutabledict(
+    {"rank": 8, "dropout": 0.0, "scale": 20.0}
+)
+_USE_DORA = False
+
+
+def _load(
+    model_name: str,
+    num_lora_layers: int,
+    use_dora: bool,
+    lora_parameters: dict | None = None,
+) -> Tuple[nn.Module, PreTrainedTokenizer]:
+    """Load model parameters and tokenizer.
+
+    Args:
+        model_name: Path or Huggingface name.
+        num_lora_layers: Number of LORA layers, if LORA is enabled.
+        use_dora: Whether to use DORA, if LORA is enabled.
+        lora_parameters: LORA hyperparameters. If not None, LORA will be enabled.
+
+    Returns:
+        Model and tokenizer.
+    """
+    model, wrapped_tokenizer = load(model_name)
+    print("Freezing all non-Lora model parameters.")
+    model.freeze()
+    if lora_parameters is not None:
+        linear_to_lora_layers(
+            model=model,
+            num_layers=num_lora_layers,
+            config=lora_parameters,
+            use_dora=use_dora,
+        )
+    return model, wrapped_tokenizer._tokenizer  # pylint: disable=protected-access
+
+
+def _loss_fn(
+    model: nn.Module,
+    inputs: mx.array,
+    targets: mx.array,
+    mask: mx.array,
+) -> mx.array:
+    logits: mx.array = model(inputs, mx.ones_like(inputs))
+    loss = nn.losses.cross_entropy(logits, targets, reduction="mean") * mask
+    normalized_loss = loss.sum() / mask.sum()
+    return normalized_loss
+
+
+class StatefulLLM:
+    """Model container that bundles revision, learning, and serving logic."""
+
+    def __init__(
+        self,
+        model_name: str = _MODEL_NAME,
+        learning_rate: float = _LEARNING_RATE,
+        max_tokens: int = _MAX_TOKENS,
+        epochs: int = _EPOCHS,
+        num_lora_layers: int = _NUM_LORA_LAYERS,
+        lora_parameters: dict | None = _LORA_PARAMETERS,
+        use_dora: bool = _USE_DORA,
+    ) -> None:
+        """Initializes the StatefulLLM
+
+        Args:
+            model_name: Path or Huggingface name.
+            learning_rate: Backpropagation hyperparameter.
+            max_tokens: Maximum number of tokens to decode in a single turn.
+            epochs: Number of training epochs to perform on self-reflective model revisions.
+            num_lora_layers: Number of LORA layers, if LORA is enabled.
+            lora_parameters: LORA hyperparameters. If not None, LORA will be enabled.
+            use_dora: Whether to use DORA, if LORA is enabled.
+
+        Returns: None
+        """
+        self._lock = threading.Lock()
+        with self._lock:
+            self._messages: list[dict[str, str]] = []
+            self._model, self._tokenizer = _load(
+                model_name,
+                num_lora_layers=num_lora_layers,
+                lora_parameters=lora_parameters,
+                use_dora=use_dora,
+            )
+        self._model_name = model_name
+        self._optimizer = optimizers.AdamW(learning_rate=learning_rate)
+        self._epochs = epochs
+        self._max_tokens = max_tokens
+        self._bos_token = self._tokenizer.special_tokens_map.get("bos_token", "")
+        self._eos_token = self._tokenizer.special_tokens_map.get("eos_token", "")
+
+        self._model_is_stable = True
+        self._response_stream = None
+
+    @property
+    def ok(self) -> bool:
+        """Checks if the model is in a stable state (i.e. not in the middle of backprop)."""
+        return self._model_is_stable
+
+    def generate_response(
+        self,
+        prompt: str,
+        use_history: bool = True,
+        max_tokens: int | None = None,
+    ) -> str:
+        """Generates model response, including handling tokenization and de-tokenization.
+
+        Args:
+            prompt: User input.
+            use_history: Whether to include previous interaction history within the prompt.
+            max_tokens: Maximum number of tokens to decode in a single turn.
+
+        Returns:
+            Model-generated response.
+        """
+        vizible.blue(f"Generating response for {prompt = }")
+        message = {"role": "user", "content": prompt}
+        if max_tokens is None:
+            max_tokens = self._max_tokens
+        if use_history:
+            self._messages.append(message)
+            messages = self._messages
+        else:
+            messages = [message]
+        tokenized_prompt = self._tokenizer.apply_chat_template(messages)
+        model_response = generate(
+            self._model,
+            self._tokenizer,
+            prompt=tokenized_prompt,
+            verbose=True,
+            max_tokens=max_tokens,
+        )
+        return model_response.strip()
+
+    async def stream_response(
+        self,
+        prompt: str,
+        use_history: bool = True,
+        max_tokens: int | None = None,
+    ) -> AsyncIterable[str]:
+        """Stream generated response"""
+        message = {"role": "user", "content": prompt}
+        if max_tokens is None:
+            max_tokens = self._max_tokens
+        print(f'{self._messages = }')
+        if use_history:
+            self._messages.append(message)
+            messages = self._messages
+        else:
+            messages = [message]
+        tokenized_prompt = self._tokenizer.apply_chat_template(messages)
+        responses = []
+        try:
+            for response in stream_generate(
+                self._model,
+                self._tokenizer,
+                prompt=tokenized_prompt,
+                max_tokens=max_tokens,
+            ):
+                text = response.text
+                responses.append(text)
+                yield text
+        finally:
+            self._response_stream = "".join(responses)
+
+    def _tokenize(self, inp: str, dtype: mx.Dtype = mx.int32) -> mx.array:
+        return mx.array(
+            self._tokenizer.encode(inp, add_special_tokens=False), dtype=dtype
+        )
+
+    def _self_correct(
+        self,
+        interaction_history: List[InteractionHistory],
+        indices_to_review: List[int] | None,
+        verbose: bool,
+    ) -> TrainingExample:
+        self._model_is_stable = False
+        vizible.green("\n--- Starting Self-Correction and Training Cycle ---")
+        if indices_to_review is None:
+            indices_to_review = list(range(len(interaction_history)))
+        if not interaction_history or not indices_to_review:
+            raise ValueError("No unreviewed interactions to process.")
+        if verbose:
+            vizible.magenta(
+                f"Found {len(interaction_history)} unreviewed interactions."
+            )
+        interactions_to_review = []
+        for idx in indices_to_review:
+            # Mark as reviewed to skip re-processing in the next cycle.
+            interaction_history[idx].reviewed = True
+            interactions_to_review.append(interaction_history[idx])
+
+        # 1. Have the model re-evaluate its past responses and try to improve upon one of its turns.
+        review_prompt = revise.make_revision_prompt(
+            interactions_to_review, self._tokenizer
+        )
+        llm_rewrite_response = self.generate_response(review_prompt, use_history=False)
+        if verbose:
+            vizible.blue(f"  - Response: {llm_rewrite_response}")
+
+        # 2. Prepare training data to train the model on how it should have responded in this
+        #    situation.
+        example = revise.make_collated_training_example(
+            llm_rewrite_response, interactions_to_review, self._tokenizer
+        )
+        return example
+
+    def _train(self, example: TrainingExample, verbose: bool):
+        state = [self._model.state, self._optimizer.state, mx.random.state]
+        mx.eval(state)
+        loss_and_grad_fn = nn.value_and_grad(self._model, _loss_fn)
+
+        @functools.partial(mx.compile, inputs=state, outputs=state)
+        def _step(inputs, labels, mask):
+            loss, grads = loss_and_grad_fn(self._model, inputs, labels, mask)
+            self._optimizer.update(self._model, grads)
+            return loss
+
+        # 3. Train the model on the new, improved examples (backward-pass)
+        losses = []
+
+        if verbose:
+            vizible.green(f"During training: {mx.metal.device_info() = }")
+        mx.set_wired_limit(mx.metal.device_info()["max_recommended_working_set_size"])
+        mx.set_cache_limit(mx.metal.device_info()["max_buffer_length"])
+        world = mx.distributed.init()
+        world_size = world.size()
+        rank = world.rank()
+        if world_size > 1:
+            tqdm.tqdm.write(f"Node {rank} of {world_size}")
+
+        self._model.train(True)
+        for epoch in tqdm.tqdm(
+            range(self._epochs), desc="Training", total=self._epochs
+        ):
+            loss = _step(example.input, example.label, example.mask)
+            mx.eval(state, loss)
+            if verbose:
+                vizible.green(f"Epoch: {epoch}\tLoss: {loss = }")
+            losses.append(loss)
+
+        self._model.train(False)
+        if verbose:
+            vizible.cyan(f"{losses = }")
+
+    def self_correct_and_train(
+        self,
+        interaction_history: List[InteractionHistory],
+        indices_to_review: List[int] | None = None,
+        verbose: bool = False,
+    ) -> bool:
+        """Cycle in which the model revises a previously unreviewed prompt and trains from its rewrite.
+
+        Args:
+            interaction_history: Past user and model messages.
+            indices_to_review: Optional indices of relevant interactions to revise within
+                               interaction_history. If not set, all interactions will be reviewed.
+            verbose: Enable verbose logging.
+
+        Returns:
+            Whether the process completed successfully.
+        """
+        self._model_is_stable = False
+
+        # Prepare training example from self-reflective revision of past dialog.
+        example = self._self_correct(interaction_history, indices_to_review, verbose)
+
+        # Train the model on the new, improved examples (backward-pass)
+        self._train(example, verbose)
+
+        self._model_is_stable = True
+        return True

adaptible-1.0.0a1/adaptible/_src/_server.py

@@ -0,0 +1,40 @@
+"""Server definition for spinning up a StatefulLLM instance."""
+
+import asyncio
+import logging
+import socket
+import sys
+
+from fastapi import FastAPI
+import uvicorn
+
+from ._api import app
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', stream=sys.stdout)
+logger = logging.getLogger(__name__) # Get a logger for this module
+
+
+class MutableHostedLLM(uvicorn.Server):
+    """Host server to interact with a stateful LLM"""
+
+    def __init__(self, uvicorn_app: FastAPI = app, host: str = '127.0.0.1', port: int = 8000):
+        super().__init__(config=uvicorn.Config(uvicorn_app, host=host, port=port))
+        self._startup_done = asyncio.Event()
+        self._serve_task = None
+        self.should_exit = False
+
+    async def startup(self, sockets: list[socket.socket] | None = None) -> None:
+        """Override uvicorn startup"""
+        await super().startup(sockets=sockets)
+        self.config.setup_event_loop()
+        self._startup_done.set()
+
+    async def up(self) -> None:
+        """Start up server asynchronously"""
+        self._serve_task = asyncio.create_task(self.serve())
+        await self._startup_done.wait()
+
+    async def down(self) -> None:
+        """Shut down server asynchronously"""
+        self.should_exit = True
+        await self._serve_task

adaptible-1.0.0a1/adaptible/_src/libs/__init__.py

@@ -0,0 +1 @@
+from . import revise