diff --git a/gpu_container/Dockerfile b/gpu_container/Dockerfile
new file mode 100644
index 000000000..50531e633
--- /dev/null
+++ b/gpu_container/Dockerfile
@@ -0,0 +1,30 @@
+FROM nvidia/cuda:12.1.1-devel-ubuntu20.04
+
+# Set the working directory
+WORKDIR /app
+
+# Install Python 3.9
+RUN apt-get update && \
+    DEBIAN_FRONTEND=noninteractive apt-get install -y software-properties-common && \
+    add-apt-repository ppa:deadsnakes/ppa && \
+    apt-get update && \
+    DEBIAN_FRONTEND=noninteractive apt-get install -y python3.9 python3.9-dev python3.9-distutils curl && \
+    # Install pip for python3.9
+    curl -sS https://bootstrap.pypa.io/get-pip.py | python3.9 && \
+    # Make python3 point to python3.9
+    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1 && \
+    # Clean up
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Copy the requirements file into the container
+COPY requirements.txt .
+
+# Install the required packages
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the application code
+COPY . ./gpu_container/
+
+# Command to run the application
+CMD ["uvicorn", "gpu_container.app:app", "--host", "0.0.0.0", "--port", "8000"]
diff --git a/gpu_container/__init__.py b/gpu_container/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/gpu_container/app.py b/gpu_container/app.py
new file mode 100644
index 000000000..226c1df3f
--- /dev/null
+++ b/gpu_container/app.py
@@ -0,0 +1,25 @@
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
+from gpu_container.embeddings.lifespan import lifespan as embeddings_lifespan
+from gpu_container.embeddings.router import router as embeddings_router
+from gpu_container.vllm.lifespan import lifespan as vllm_lifespan
+from gpu_container.vllm.router import router as vllm_router
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """
+    A top-level lifespan handler that calls the lifespan handlers
+    for different parts of the application.
+    """
+    async with embeddings_lifespan(app):
+        async with vllm_lifespan(app):
+            yield
+
+
+app = FastAPI(lifespan=lifespan)
+
+app.include_router(embeddings_router)
+app.include_router(vllm_router)
diff --git a/gpu_container/docker-compose.yml b/gpu_container/docker-compose.yml
new file mode 100644
index 000000000..c29496ac3
--- /dev/null
+++ b/gpu_container/docker-compose.yml
@@ -0,0 +1,16 @@
+services:
+  gpu-app:
+    build: .
+    ports:
+      - "8000:8000"
+    environment:
+      - MODEL_ID=WhereIsAI/UAE-Large-V1
+      - VLLM_MODEL_ID=mrfakename/mistral-small-3.1-24b-instruct-2503-hf
+      - VLLM_GPU_UTILIZATION=0.8
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
diff --git a/gpu_container/embeddings/__init__.py b/gpu_container/embeddings/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/gpu_container/embeddings/lifespan.py b/gpu_container/embeddings/lifespan.py
new file mode 100644
index 000000000..c6d31780a
--- /dev/null
+++ b/gpu_container/embeddings/lifespan.py
@@ -0,0 +1,42 @@
+import os
+from contextlib import asynccontextmanager
+
+import torch
+from angle_emb import AnglE
+from fastapi import FastAPI
+
+
+def load_config_from_env():
+    """Loads configuration from environment variables."""
+    model_id = os.getenv("MODEL_ID", "WhereIsAI/UAE-Large-V1")
+    device = os.getenv("DEVICE", "cpu")
+
+    return {"model_id": model_id, "device": device}
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Handle embedding model startup and shutdown."""
+    print("Loading embeddings model...")
+    config = load_config_from_env()
+    print(f"Loading model: {config['model_id']} on device: {config['device']}")
+
+    model = AnglE.from_pretrained(config["model_id"], pooling_strategy="cls")
+
+    if config["device"] == "cuda" and torch.cuda.is_available():
+        model.to(torch.device("cuda"))
+        print("Embeddings model moved to CUDA.")
+    else:
+        model.to(torch.device("cpu"))
+        print("Embeddings model moved to CPU.")
+
+    app.state.embeddings_model = model
+    app.state.embeddings_model_id = config["model_id"]
+    print("Embeddings model loaded.")
+
+    yield
+
+    print("Shutting down embeddings model...")
+    app.state.model = None
+    app.state.model_id = None
+    print("Embeddings model shut down.")
diff --git a/gpu_container/embeddings/router.py b/gpu_container/embeddings/router.py
new file mode 100644
index 000000000..476a49df0
--- /dev/null
+++ b/gpu_container/embeddings/router.py
@@ -0,0 +1,44 @@
+from typing import List
+
+import numpy as np
+from fastapi import APIRouter, Request
+from pydantic import BaseModel
+
+router = APIRouter()
+
+
+class EmbeddingRequest(BaseModel):
+    input: List[str]
+
+
+class Embedding(BaseModel):
+    object: str = "embedding"
+    index: int
+    embedding: List[float]
+
+
+class EmbeddingResponse(BaseModel):
+    object: str = "list"
+    data: List[Embedding]
+    model: str
+
+
+@router.post("/v1/embeddings", response_model=EmbeddingResponse)
+async def get_embeddings(request: Request, body: EmbeddingRequest):
+    """Generate embeddings for a list of texts."""
+    model = request.app.state.embeddings_model
+    model_id = request.app.state.embeddings_model_id
+
+    if model is None:
+        return {"error": "Model not loaded"}, 503
+
+    # Generate embeddings
+    embeddings = model.encode(body.input, to_numpy=True)
+
+    # Ensure embeddings are a list of lists of floats
+    if isinstance(embeddings, np.ndarray):
+        embeddings = embeddings.tolist()
+
+    response_data = [Embedding(index=i, embedding=embedding) for i, embedding in enumerate(embeddings)]
+
+    return EmbeddingResponse(data=response_data, model=model_id)
diff --git a/gpu_container/requirements.txt b/gpu_container/requirements.txt
new file mode 100644
index 000000000..c61d03ef6
--- /dev/null
+++ b/gpu_container/requirements.txt
@@ -0,0 +1,6 @@
+angle-emb
+torch
+fastapi
+uvicorn
+pydantic
+vllm==0.8.5
diff --git a/gpu_container/vllm/__init__.py b/gpu_container/vllm/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/gpu_container/vllm/lifespan.py b/gpu_container/vllm/lifespan.py
new file mode 100644
index 000000000..a71875ddf
--- /dev/null
+++ b/gpu_container/vllm/lifespan.py
@@ -0,0 +1,34 @@
+import os
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
+from gpu_container.vllm.reproducible_vllm import ReproducibleVLLM
+
+
+def load_config_from_env():
+    """Loads vLLM configuration from environment variables."""
+    vllm_model_id = os.getenv("VLLM_MODEL_ID", "default_model_id")
+    device = os.getenv("DEVICE", "cuda")
+    # Add any other vLLM-specific environment variables here
+    return {"vllm_model_id": vllm_model_id, "device": device}
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Handle vLLM engine startup and shutdown."""
+    print("Loading vLLM engine...")
+    config = load_config_from_env()
+
+    engine = ReproducibleVLLM(model_id=config["vllm_model_id"], device=config["device"])
+
+    app.state.vllm_engine = engine
+    app.state.vllm_model_id = config["vllm_model_id"]
+    print("vLLM engine loaded.")
+
+    yield
+
+    print("Shutting down vLLM engine...")
+    app.state.vllm_engine = None
+    app.state.vllm_model_id = None
+    print("vLLM engine shut down.")
diff --git a/gpu_container/vllm/reproducible_vllm.py b/gpu_container/vllm/reproducible_vllm.py
new file mode 100644
index 000000000..dd2f532e9
--- /dev/null
+++ b/gpu_container/vllm/reproducible_vllm.py
@@ -0,0 +1,169 @@
+import random
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from vllm import LLM, SamplingParams
+
+
+class ReproducibleVLLM:
+    def __init__(
+        self,
+        model_id: str = "mrfakename/mistral-small-3.1-24b-instruct-2503-hf",
+        device: str = "cuda:0",
+        sampling_params: Optional[Dict[str, Union[str, float, int, bool]]] = None,
+    ):
+        """Deterministic VLLM model."""
+        self._device = device
+        self.model_id = model_id
+        self.sampling_params = {} if sampling_params is None else sampling_params
+
+        self.model = LLM(
+            model=model_id,
+            trust_remote_code=True,
+            gpu_memory_utilization=0.9,
+        )
+
+        # Store tokenizer from VLLM for consistency
+        self.tokenizer = self.model.get_tokenizer()
+
+    @classmethod
+    async def get_max_tokens(
+        cls,
+        sampling_params: Dict[str, Union[str, float, int, bool]],
+        default_value: int = 512,
+    ) -> int:
+        # Process max tokens with backward compatibility.
+        max_tokens = sampling_params.get("max_tokens")
+        if max_tokens is None:
+            max_tokens = sampling_params.get("max_new_tokens")
+        if max_tokens is None:
+            max_tokens = sampling_params.get("max_completion_tokens", default_value)
+        return max_tokens
+
+    @classmethod
+    async def prepare_sampling_params(
+        cls, sampling_params: Optional[Dict[str, Union[str, float, int, bool]]] = None
+    ) -> SamplingParams:
+        sampling_params = sampling_params or {}
+        max_tokens = await cls.get_max_tokens(sampling_params)
+
+        params = SamplingParams(
+            temperature=float(sampling_params.get("temperature", 1.0)),
+            top_p=float(sampling_params.get("top_p", 1.0)),
+            max_tokens=int(max_tokens),
+            presence_penalty=float(sampling_params.get("presence_penalty", 0.0)),
+            frequency_penalty=float(sampling_params.get("frequency_penalty", 0.0)),
+            top_k=int(sampling_params.get("top_k", -1)),
+            logprobs=sampling_params.get("logprobs", None),
+        )
+        return params
+
+    async def generate(
+        self,
+        messages: Union[List[str], List[Dict[str, str]]],
+        sampling_params: Optional[Dict[str, Union[str, float, int, bool]]] = None,
+        seed: Optional[int] = None,
+        continue_last_message: bool = False,
+    ) -> str:
+        """Generate text with optimized performance using VLLM."""
+        self.set_random_seeds(seed)
+
+        # Convert chat messages to prompt string using tokenizer's chat template
+        if isinstance(messages, list) and isinstance(messages[0], dict):
+            try:
+                # Extract any trailing whitespace before applying template
+                trailing_space = ""
+                if continue_last_message and messages[-1]["content"]:
+                    content = messages[-1]["content"]
+                    stripped = content.rstrip()
+                    if len(content) > len(stripped):
+                        trailing_space = content[len(stripped) :]
+
+                # Try using the tokenizer's chat template
+                prompt = self.tokenizer.apply_chat_template(
+                    conversation=messages,
+                    tokenize=False,
+                    add_generation_prompt=not continue_last_message,
+                    continue_final_message=continue_last_message,
+                )
+
+                # Append back just the trailing whitespace if it was stripped
+                if trailing_space:
+                    prompt += trailing_space
+            except (AttributeError, NotImplementedError):
+                raise ValueError(f"Chat template not supported for model {self.model_id}")
+        else:
+            prompt = messages[0] if isinstance(messages, list) else messages
+
+        # Convert sampling parameters to vLLM format.
+        params = sampling_params if sampling_params is not None else self.sampling_params
+        vllm_params = await self.prepare_sampling_params(params)
+        outputs = self.model.generate(prompt, vllm_params)
+
+        if not outputs:
+            return ""
+
+        result = outputs[0].outputs[0].text
+        return {"choices": [{"message": {"content": result}}]}
+
+    async def generate_logits(
+        self,
+        messages: Union[List[str], List[Dict[str, str]]],
+        top_logprobs: int = 10,
+        sampling_params: Optional[Dict[str, Union[str, float, int, bool]]] = None,
+        seed: Optional[int] = None,
+        continue_last_message: bool = False,
+    ) -> dict[str, float]:
+        """Generate logits for the next token prediction.
+
+        Args:
+            messages: Input messages or text.
+            top_logprobs: Number of top logits to return (default: 10).
+            sampling_params: Generation parameters.
+            seed: Random seed for reproducibility.
+            continue_last_message: Whether to continue the last message in chat format.
+
+        Returns:
+            Dictionary mapping tokens to their log probabilities.
+        """
+        self.set_random_seeds(seed)
+        params = sampling_params if sampling_params is not None else self.sampling_params
+        params = params.copy()
+        params["max_tokens"] = 1
+        params["logprobs"] = top_logprobs
+        vllm_params = await self.prepare_sampling_params(params)
+
+        prompt = self.tokenizer.apply_chat_template(
+            conversation=messages,
+            tokenize=False,
+            add_generation_prompt=not continue_last_message,
+            continue_final_message=continue_last_message,
+        )
+
+        outputs = self.model.generate(prompt, vllm_params)
+
+        if not outputs or not outputs[0].outputs[0].logprobs:
+            return {}
+
+        logprobs = outputs[0].outputs[0].logprobs[0]
+        token_logprobs = {self.tokenizer.decode([token]): logprob.logprob for token, logprob in logprobs.items()}
+        sorted_token_logprobs = dict(sorted(token_logprobs.items(), key=lambda item: item[1], reverse=True))
+        return sorted_token_logprobs, prompt
+
+    def set_random_seeds(self, seed: Optional[int] = 42):
+        """Set random seeds for reproducibility across all relevant libraries."""
+        if seed is not None:
+            random.seed(seed)
+            np.random.seed(seed)
+            torch.manual_seed(seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(seed)
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+
+    @staticmethod
+    def format_messages(
+        messages: Union[List[str], List[Dict[str, str]]],
+    ) -> List[Dict[str, Union[str, List[Dict[str, str]]]]]:
+        return messages
diff --git a/gpu_container/vllm/router.py b/gpu_container/vllm/router.py
new file mode 100644
index 000000000..821c5fb76
--- /dev/null
+++ b/gpu_container/vllm/router.py
@@ -0,0 +1,26 @@
+from fastapi import APIRouter, Request
+
+router = APIRouter()
+
+
+@router.post("/v1/chat/generate_logits")
+async def generate_logits(request: Request):
+    json_request = await request.json()
+    return await request.app.state.vllm_engine.generate_logits(
+        messages=json_request["messages"],
+        sampling_params=json_request["sampling_params"],
+        seed=json_request["seed"],
+        continue_last_message=json_request["continue_last_message"],
+        top_logprobs=json_request["top_logprobs"],
+    )
+
+
+@router.post("/v1/chat/generate")
+async def generate(request: Request):
+    json_request = await request.json()
+    return await request.app.state.vllm_engine.generate(
+        messages=json_request["messages"],
+        sampling_params=json_request["sampling_params"],
+        seed=json_request["seed"],
+        continue_last_message=json_request.get("continue_last_message", False),
+    )
diff --git a/neurons/validator.py b/neurons/validator.py
index dffe6a1c4..b8f5c4600 100644
--- a/neurons/validator.py
+++ b/neurons/validator.py
@@ -21,7 +21,6 @@
 
 settings.shared_settings = settings.SharedSettings.load(mode="validator")
 
-from prompting.llms.model_manager import AsyncModelScheduler, ModelManager
 from prompting.rewards.scoring import task_scorer
 from shared.logging import init_wandb
 
@@ -57,7 +56,6 @@ def init_process_logging(name: str):
 
 
 async def create_loop_process(
-    model_scheduler: AsyncModelScheduler,
     task_queue: list,
     scoring_queue: list,
     reward_events: list,
@@ -71,10 +69,9 @@ async def create_loop_process(
 
     all_tasks: list[asyncio.Task] = []
 
-    async def cleanup(model_scheduler):
+    async def cleanup():
         logger.info("Cleaning up resources...")
         torch.distributed.destroy_process_group()
-        await model_scheduler.llm_model.cleanup()
         for t in all_tasks:
             t.cancel()
         await asyncio.gather(*all_tasks, return_exceptions=True)
@@ -92,10 +89,8 @@ async def spawn_loops(task_queue: list, scoring_queue: list, reward_events: list
         task_loop_task = asyncio.create_task(
             task_loop.start(task_queue, scoring_queue, miners_dict, simultaneous_loops=1), name="TaskLoop"
         )
-        model_scheduler_task = asyncio.create_task(model_scheduler.start(scoring_queue), name="ModelScheduler")
         task_scorer_task = asyncio.create_task(
             task_scorer.start(
-                model_scheduler,
                 scoring_queue,
                 reward_events,
                 mp_lock=mp_lock,
@@ -104,7 +99,7 @@ async def spawn_loops(task_queue: list, scoring_queue: list, reward_events: list
             ),
             name="TaskScorer",
         )
-        all_tasks.extend([profile, task_loop_task, model_scheduler_task, task_scorer_task])
+        all_tasks.extend([profile, task_loop_task, task_scorer_task])
 
         try:
             while True:
@@ -112,8 +107,6 @@ async def spawn_loops(task_queue: list, scoring_queue: list, reward_events: list
                 logger.debug(
                     f"Task Queue {len(task_queue)}. Scoring Queue {len(scoring_queue)}. Reward Events {len(reward_events)}"
                 )
-                if model_scheduler.memory_error is not None:
-                    raise model_scheduler.memory_error
         except asyncio.CancelledError:
             logger.info("spawn_loops received cancellation signal.")
             raise
@@ -122,12 +115,12 @@ async def spawn_loops(task_queue: list, scoring_queue: list, reward_events: list
         await spawn_loops(task_queue, scoring_queue, reward_events, miners_dict)
     except MemoryError as e:
         logger.error(f"MemoryError encountered. Terminating program: {e}")
-        await cleanup(model_scheduler)
+        await cleanup()
         sys.exit(1)
     except Exception as e:
         logger.exception(f"Terminating loop process: {e}")
     finally:
-        await cleanup(model_scheduler)
+        await cleanup()
 
 
 def start_api(
@@ -289,8 +282,6 @@ async def main(
         tasks: list[asyncio.Task] = []
         event_stop = mp.Event()
 
-        model_scheduler = AsyncModelScheduler(llm_model_manager=ModelManager(), mp_lock=mp_lock, sync=True)
-
         try:
             # Start checking the availability of miners at regular intervals
             if not settings.shared_settings.NEURON_DISABLE_SET_WEIGHTS:
@@ -315,7 +306,6 @@ async def main(
 
             loop_task = asyncio.create_task(
                 create_loop_process(
-                    model_scheduler=model_scheduler,
                     task_queue=task_queue,
                     scoring_queue=scoring_queue,
                     reward_events=reward_events,
diff --git a/prompting/api/weight_syncing/api.py b/prompting/api/weight_syncing/api.py
index 8bf6ff2ca..70f18e52d 100644
--- a/prompting/api/weight_syncing/api.py
+++ b/prompting/api/weight_syncing/api.py
@@ -14,6 +14,10 @@ def get_weight_dict(request: Request):
     return request.app.state.weight_dict
 
 
+def get_uid_from_hotkey(hotkey: str):
+    return shared_settings.METAGRAPH.hotkeys.index(hotkey)
+
+
 async def verify_weight_signature(request: Request):
     signed_by = request.headers.get("Epistula-Signed-By")
     signed_for = request.headers.get("Epistula-Signed-For")
@@ -26,6 +30,9 @@ async def verify_weight_signature(request: Request):
         raise HTTPException(status_code=401, detail="Signer not the expected ss58 address")
     now = time.time()
     body = await request.body()
+    if body["uid"] != get_uid_from_hotkey(signed_by):
+        logger.error("Invalid uid")
+        raise HTTPException(status_code=400, detail="Invalid uid in body")
     err = verify_signature(
         request.headers.get("Epistula-Request-Signature"),
         body,
diff --git a/prompting/llms/model_manager.py b/prompting/llms/model_manager.py
index aa5388273..effc35bd9 100644
--- a/prompting/llms/model_manager.py
+++ b/prompting/llms/model_manager.py
@@ -1,17 +1,4 @@
 import asyncio
-import gc
-from multiprocessing.managers import AcquirerProxy
-from typing import ClassVar
-
-import torch
-from loguru import logger
-from pydantic import BaseModel, ConfigDict
-
-from prompting.llms.model_zoo import ModelConfig, ModelZoo
-from prompting.llms.utils import GPUInfo, model_factory
-from prompting.llms.vllm_llm import ReproducibleVLLM
-from shared import settings
-from shared.loop_runner import AsyncLoopRunner
 
 
 class AsyncRLock:
@@ -45,224 +32,3 @@ async def __aenter__(self):
 
     async def __aexit__(self, exc_type, exc, tb):
         self.release()
-
-
-class ModelManager(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-    total_ram: float = settings.shared_settings.LLM_MODEL_RAM
-    active_models: dict[ModelConfig, ReproducibleVLLM] = {}
-    loading_tasks: dict[ModelConfig, asyncio.Future] = {}
-    used_ram: float = 0.0
-    lock: ClassVar[AsyncRLock] = AsyncRLock()
-    # lock: ClassVar[AsyncRLock] = asyncio.Lock()
-
-    async def load_model(self, model_config: ModelConfig, force: bool = True) -> ReproducibleVLLM:
-        """Load model into GPU.
-
-        Warning: This operation will block execution until the model is successfully loaded into VRAM.
-
-        Args:
-            model_config: Model config to load.
-            force: If enabled, will unload all other models.
-        """
-        async with self.lock:
-            # Copy active models, since they will be modified in the loop.
-            active_models = set(self.active_models.keys())
-
-            if model_config in active_models:
-                logger.debug(f"Model {model_config.llm_model_id} is already loaded.")
-                return self.active_models[model_config]
-
-            if force:
-                logger.debug(f"Forcing model {model_config.llm_model_id} to load.")
-                for active_model in active_models:
-                    logger.debug(f"Unloading {active_model.llm_model_id} to make room for {model_config.llm_model_id}")
-                    await self._unload_model(active_model)
-                await self.cleanup()
-
-            try:
-                GPUInfo.log_gpu_info()
-                model_class = model_factory(model_config.llm_model_id)
-                model = model_class(
-                    model_id=model_config.llm_model_id,
-                    device=settings.shared_settings.NEURON_DEVICE,
-                    sampling_params=settings.shared_settings.SAMPLING_PARAMS,
-                )
-                self.active_models[model_config] = model
-                self.used_ram += model_config.min_ram
-                logger.info(
-                    f"Model {model_config.llm_model_id} has been successfully loaded. "
-                    f"Approx. used VRAM: {self.used_ram:.0f}GB"
-                )
-                await asyncio.sleep(1.0)
-                return model
-            except BaseException as e:
-                await self.cleanup()
-                # In case of VRAM leak, raise an exception to terminate the process.
-                raise MemoryError(f"Failed to load model {model_config.llm_model_id}: {e}")
-
-    async def _cleanup_model(self, model_instance: ReproducibleVLLM, cpu_offload: bool = False):
-        """Free VRAM from given model."""
-        if cpu_offload:
-            try:
-                model_instance.model = model_instance.model.to("cpu")
-            except NotImplementedError as e:
-                logger.exception(f"Standard move to CPU failed: {str(e)}")
-                try:
-                    # Fallback for meta tensors.
-                    model_instance.model = model_instance.model.to_empty("cpu")
-                except Exception as fallback_e:
-                    logger.exception(f"Could not move meta model to CPU, proceeding with generic GC: {str(fallback_e)}")
-            except Exception as e:
-                logger.exception(f"Unexpected error when moving model to CPU: {str(e)}")
-
-        model_instance.unload_model()
-        del model_instance
-
-    async def _unload_model(self, model_config: ModelConfig):
-        if model_config not in self.active_models:
-            logger.warning(f"Couldn't find given model to unload: {model_config}")
-            return
-
-        try:
-            initial_free_memory = GPUInfo.free_memory
-            logger.debug(f"Initial free GPU memory before unloading: {initial_free_memory} GB")
-            # async with self.rlock:
-            model_instance = self.active_models.pop(model_config)
-            await self._cleanup_model(model_instance, cpu_offload=False)
-            await self.cleanup()
-
-            memory_freed = GPUInfo.free_memory - initial_free_memory
-            logger.info(f"Successfully unloaded model {model_config.llm_model_id}. Memory freed: {memory_freed:.2f} GB")
-
-        except Exception as ex:
-            logger.error(f"Failed to unload model {model_config.llm_model_id}. Error: {str(ex)}")
-
-        # Update used RAM tracking
-        self.used_ram -= model_config.min_ram
-
-        GPUInfo.log_gpu_info()
-
-    async def get_model(self, llm_model: ModelConfig | str) -> ReproducibleVLLM:
-        async with self.lock:
-            if not llm_model:
-                llm_model = next(iter(self.active_models.keys())) if self.active_models else ModelZoo.get_random()
-            if isinstance(llm_model, str):
-                llm_model = ModelZoo.get_model_by_id(llm_model)
-            if llm_model in self.active_models:
-                return self.active_models[llm_model]
-
-        return await self.load_model(llm_model)
-
-    async def generate(
-        self,
-        messages: list[str] | list[dict],
-        roles: list[str] | None = None,
-        model: ModelConfig | str | None = None,
-        seed: int = None,
-        sampling_params: dict[str, float] = None,
-    ) -> str:
-        if messages and isinstance(messages[0], dict):
-            dict_messages = messages
-        else:
-            dict_messages = [{"content": message, "role": role} for message, role in zip(messages, roles)]
-
-        async with self.lock:
-            if isinstance(model, str):
-                model = ModelZoo.get_model_by_id(model)
-            if not model:
-                model = ModelZoo.get_random(max_ram=self.total_ram)
-
-        model_instance: ReproducibleVLLM = await self.get_model(model)
-
-        async with self.lock:
-            if model_instance is None:
-                raise ValueError("Model is None, which may indicate the model is still loading.")
-            responses = await model_instance.generate(
-                messages=dict_messages, sampling_params=sampling_params, seed=seed
-            )
-            return responses
-
-    async def generate_logits(
-        self,
-        messages: list[str],
-        model: ModelConfig | str | None = None,
-        sampling_params: dict[str, float] = None,
-        seed: int = None,
-        continue_last_message: bool = False,
-        top_logprobs: int = 10,
-    ):
-        model_instance: ReproducibleVLLM = await self.get_model(model)
-        return await model_instance.generate_logits(
-            messages=messages,
-            sampling_params=sampling_params,
-            seed=seed,
-            continue_last_message=continue_last_message,
-            top_logprobs=top_logprobs,
-        )
-
-    async def cleanup(self):
-        """Perform VRAM clean-up."""
-        for _, model in self.active_models.items():
-            del model.model
-            del model
-
-        self.active_models = {}
-        self.used_ram = 0.0
-
-        if torch.cuda.is_available():
-            # Reset all CUDA cached memory.
-            try:
-                torch.cuda.synchronize()
-                gc.collect()
-                torch.cuda.empty_cache()
-                torch.cuda.reset_peak_memory_stats()
-                torch.cuda.reset_accumulated_memory_stats()
-                await asyncio.sleep(1.0)
-            except BaseException as e:
-                logger.warning(f"Error during CUDA empty cache: {e}")
-        else:
-            logger.warning("CUDA is not available")
-
-        gc.collect()
-        gc.collect(generation=2)
-        await asyncio.sleep(1.0)
-
-        logger.info(f"VRAM clean-up completed. Current GPU usage: {GPUInfo.gpu_utilization * 100:.2f}%")
-        GPUInfo.log_gpu_info()
-
-
-class AsyncModelScheduler(AsyncLoopRunner):
-    llm_model_manager: ModelManager
-    mp_lock: AcquirerProxy
-    interval: int = 3600
-    scoring_queue: list | None = None
-    memory_error: MemoryError | None = None
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
-    async def start(self, scoring_queue: list, name: str | None = None, **kwargs):
-        self.scoring_queue = scoring_queue
-        await super().start(name=name, **kwargs)
-        # Load the model immediately.
-        await self.run_step()
-
-    async def run_step(self):
-        """This method is called periodically according to the interval."""
-        # try to load the model belonging to the oldest task in the queue
-        with self.mp_lock:
-            selected_model = self.scoring_queue[0].task.llm_model if self.scoring_queue else None
-        if not selected_model:
-            selected_model = ModelZoo.get_random(max_ram=self.llm_model_manager.total_ram)
-        logger.info(f"Loading model {selected_model.llm_model_id} for {self.interval} seconds.")
-
-        if selected_model in self.llm_model_manager.active_models:
-            logger.info(f"Model {selected_model.llm_model_id} is already loaded.")
-            return
-
-        try:
-            await self.llm_model_manager.load_model(selected_model)
-        except MemoryError as e:
-            self.memory_error = e
-        logger.debug(f"Active models: {list(self.llm_model_manager.active_models.keys())}")
-        await asyncio.sleep(0.01)
diff --git a/prompting/rewards/exact_match.py b/prompting/rewards/exact_match.py
index 19000b23d..250df786e 100644
--- a/prompting/rewards/exact_match.py
+++ b/prompting/rewards/exact_match.py
@@ -7,11 +7,11 @@
 from loguru import logger
 from openai.types.chat import ChatCompletionChunk
 
-from prompting.llms.model_manager import ModelManager
 from prompting.rewards.reward import BaseRewardModel, BatchRewardOutput
 from prompting.tasks.base_task import BaseTextTask
-from shared import settings
+from shared import constants, settings
 from shared.dendrite import DendriteResponseEvent
+from shared.docker_utils import get_logits
 
 shared_settings = settings.shared_settings
 
@@ -36,13 +36,9 @@ async def reward(  # noqa: C901
         reference: str,
         response_event: DendriteResponseEvent,
         task: BaseTextTask,
-        model_manager: ModelManager,
         **kwargs,
     ) -> BatchRewardOutput:
         """Calculate rewards based on the logits of the response and verifies them."""
-        if model_manager is None:
-            raise ValueError("Model manager must be set")
-
         all_chunks: list[list[str]] = response_event.stream_results_all_chunks
         all_chunk_dicts_raw: list[list[ChatCompletionChunk | dict]] = response_event.stream_results_all_chunk_dicts_raw
         uids: np.ndarray | list[float] = response_event.uids
@@ -64,10 +60,10 @@ async def reward(  # noqa: C901
             raise ValueError("Timeout must be greater than 0.")
 
         # If max_tokens are not provided, always check for eos.
-        model = await model_manager.get_model(task.llm_model_id)
-        max_tokens = await model.get_max_tokens(sampling_parameters, default_value=2048)
-        eos_token = model.tokenizer.eos_token
-        bos_token = model.tokenizer.bos_token
+        model = task.llm_model_id
+        max_tokens = sampling_parameters.get("max_tokens", 2048)
+        eos_token = constants.SPECIAL_TOKENS.get(model, {}).get("eos_token")
+        bos_token = constants.SPECIAL_TOKENS.get(model, {}).get("bos_token")
         special_tokens = set([bos_token, eos_token])
         timing_verified: list[list[float]] = []
         rewards: list[float] = []
@@ -110,14 +106,14 @@ async def reward(  # noqa: C901
                     to_complete = "".join(chunks[:check_idx])
                     if to_complete:
                         messages.extend([{"role": "assistant", "content": to_complete}])
-
-                    verification_logits, _ = await model_manager.generate_logits(
+                    response = await get_logits(
                         model=task.llm_model_id,
                         messages=messages,
                         top_logprobs=TOP_LOGPROBS,
                         sampling_params=sampling_parameters,
                         continue_last_message=len(to_complete) > 0,
                     )
+                    verification_logits = response[0]
                     if check_idx < eos_idx:
                         if chunks[check_idx] in special_tokens:
                             raise ValueError("Special tokens mid-completion")
diff --git a/prompting/rewards/inference_reward_model.py b/prompting/rewards/inference_reward_model.py
index a70db9d4a..999c1203e 100644
--- a/prompting/rewards/inference_reward_model.py
+++ b/prompting/rewards/inference_reward_model.py
@@ -14,17 +14,15 @@ async def reward(
         response_event: DendriteResponseEvent,
         model_id: str | None = None,
         task: BaseTextTask | None = None,
-        model_manager=None,
         **kwargs,
     ) -> BatchRewardOutput:
         """Gives an exact reward of 1 if the response matches the reference, 0 otherwise"""
-        if model_manager is None:
-            raise ValueError("Model manager must be set")
+        logger.info(f"model_id: {model_id}")
 
         if model_id or task.organic:
             logger.info("Using logits reward model")
             logits_reward_model = LogitsRewardModel()
-            return await logits_reward_model.reward(reference, response_event, task, model_manager=model_manager)
+            return await logits_reward_model.reward(reference, response_event, task)
 
         relevance_reward_model = RelevanceRewardModel()
-        return await relevance_reward_model.reward(reference, response_event, model_manager=model_manager)
+        return await relevance_reward_model.reward(reference, response_event)
diff --git a/prompting/rewards/relevance.py b/prompting/rewards/relevance.py
index 98a67d790..65691a5f3 100644
--- a/prompting/rewards/relevance.py
+++ b/prompting/rewards/relevance.py
@@ -1,24 +1,46 @@
 import time
-from typing import ClassVar, Optional
+from typing import Optional
 
 import numpy as np
-from angle_emb import AnglE
+import requests
 from pydantic import ConfigDict
 from scipy import spatial
 
 from prompting.rewards.reward import BaseRewardModel, BatchRewardOutput
-from shared import settings
+from shared import constants, settings
 from shared.dendrite import DendriteResponseEvent
 
 shared_settings = settings.shared_settings
 
 
+def get_embeddings(inputs):
+    """
+    Sends a POST request to the local embeddings endpoint and returns the response.
+
+    Args:
+        inputs (str or list of str): A single input string or a list of input strings to embed.
+
+    Returns:
+        dict: JSON response from the embeddings server.
+    """
+    if isinstance(inputs, str):
+        inputs = [inputs]  # convert single string to list
+
+    url = f"{constants.DOCKER_BASE_URL}/v1/embeddings"
+    headers = {"Content-Type": "application/json"}
+    payload = {"input": inputs}
+
+    try:
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        return response.json()
+    except requests.RequestException as e:
+        return {"error": str(e)}
+
+
 class RelevanceRewardModel(BaseRewardModel):
     threshold: Optional[float] = None
     model_config = ConfigDict(arbitrary_types_allowed=True)
-    embedding_model: ClassVar[AnglE] = AnglE.from_pretrained(
-        "WhereIsAI/UAE-Large-V1", pooling_strategy="cls", device=shared_settings.NEURON_DEVICE
-    ).to(shared_settings.NEURON_DEVICE)
 
     async def reward(
         self, reference: str, response_event: DendriteResponseEvent, model_manager=None, **kwargs
@@ -31,15 +53,13 @@ async def reward(
         """
         if not reference:
             raise Exception("Reference is empty - something went wrong during the reference generation")
-        reference_embedding = self.embedding_model.encode(reference, to_numpy=True)
-        reference_emb_flatten = reference_embedding.flatten()
+        reference_embedding = np.array(get_embeddings(reference)["data"][0]["embedding"])
         rewards: list[float] = []
         timings: list[float] = []
         completions: list[str] = response_event.completions
         # baseline is the cosine similarity between the reference and an empty string
-        baseline = 1 - float(
-            spatial.distance.cosine(reference_emb_flatten, self.embedding_model.encode("", to_numpy=True).flatten())
-        )
+        baseline_embedding = np.array(get_embeddings("")["data"][0]["embedding"])
+        baseline = 1 - float(spatial.distance.cosine(reference_embedding, baseline_embedding))
 
         for comp in completions:
             if len(comp) == 0:
@@ -47,9 +67,9 @@ async def reward(
                 timings.append(0)
                 continue
             t0 = time.time()
-            emb = self.embedding_model.encode(comp, to_numpy=True)
+            emb = np.array(get_embeddings(comp)["data"][0]["embedding"])
             # Calculate cosine similarity between reference and completion embeddings, and subtract baseline
-            score = 1 - float(spatial.distance.cosine(reference_emb_flatten, emb.flatten() - baseline))
+            score = 1 - float(spatial.distance.cosine(reference_embedding, emb) - baseline)
 
             rewards.append(score)
             timings.append(time.time() - t0)
diff --git a/prompting/rewards/reward.py b/prompting/rewards/reward.py
index f7d3044b0..16d6e0aac 100644
--- a/prompting/rewards/reward.py
+++ b/prompting/rewards/reward.py
@@ -5,7 +5,6 @@
 import numpy as np
 from pydantic import BaseModel, ConfigDict, model_validator
 
-from prompting.llms.model_manager import ModelManager
 from prompting.tasks.base_task import BaseTextTask
 from shared.dendrite import DendriteResponseEvent
 
@@ -75,7 +74,6 @@ async def reward(
         self,
         reference: str,
         response_event: DendriteResponseEvent,
-        model_manager: ModelManager = None,
         task_queue: list[BaseTextTask] | None = None,
         **kwargs,
     ) -> BatchRewardOutput:
@@ -88,7 +86,6 @@ async def apply(
         challenge: str | None = None,
         reward_type: Literal["reward", "penalty"] = "reward",
         task: BaseTextTask | None = None,
-        model_manager: ModelManager | None = None,
         task_queue: list[BaseTextTask] | None = None,
         **kwargs,
     ) -> WeightedRewardEvent:
@@ -97,7 +94,7 @@ async def apply(
         t0 = time.time()
         comparator = reference if reward_type == "reward" else challenge
         batch_rewards_output: BatchRewardOutput = await self.reward(
-            comparator, response_event, task=task, model_manager=model_manager, task_queue=task_queue, **kwargs
+            comparator, response_event, task=task, task_queue=task_queue, **kwargs
         )
         batch_rewards_time = time.time() - t0
         uids = batch_rewards_output.uids if batch_rewards_output.uids is not None else response_event.uids
@@ -159,7 +156,6 @@ async def apply(
         challenge: str | None = None,
         model_id: str | None = None,
         task: BaseTextTask | None = None,
-        model_manager: ModelManager | None = None,
         task_queue: list[BaseTextTask] | None = None,
     ) -> list[WeightedRewardEvent]:
         if task_queue is None:
@@ -174,7 +170,6 @@ async def apply(
                     reward_type="reward",
                     model_id=model_id,
                     task=task,
-                    model_manager=model_manager,
                     task_queue=task_queue,
                 ),
             )
diff --git a/prompting/rewards/scoring.py b/prompting/rewards/scoring.py
index 468eaefb4..ee7bcef97 100644
--- a/prompting/rewards/scoring.py
+++ b/prompting/rewards/scoring.py
@@ -6,7 +6,6 @@
 from loguru import logger
 from pydantic import ConfigDict
 
-from prompting.llms.model_manager import AsyncModelScheduler
 from prompting.rewards.scoring_config import ScoringConfig
 from prompting.tasks.base_task import BaseTextTask
 from prompting.tasks.MSRv2_task import MSRv2Task
@@ -26,7 +25,6 @@ class TaskScorer(AsyncLoopRunner):
 
     mp_lock: AcquirerProxy | None = None
     is_running: bool = False
-    model_scheduler: AsyncModelScheduler | None = None
     thread: threading.Thread = None
     interval: int = 1
     scoring_queue: list | None = None
@@ -36,7 +34,6 @@ class TaskScorer(AsyncLoopRunner):
 
     async def start(
         self,
-        model_scheduler: AsyncModelScheduler,
         scoring_queue,
         reward_events,
         mp_lock: AcquirerProxy,
@@ -46,7 +43,6 @@ async def start(
     ):
         self.scoring_queue = scoring_queue
         self.reward_events = reward_events
-        self.model_scheduler = model_scheduler
         self.mp_lock = mp_lock
         self.task_queue = task_queue
         return await super().start(name=name, **kwargs)
@@ -73,26 +69,19 @@ def add_to_queue(
 
     async def run_step(self) -> RewardLoggingEvent:
         await asyncio.sleep(0.1)
-        # Only score responses for which the model is loaded
-        await self.model_scheduler.llm_model_manager.lock.acquire()
-        with self.mp_lock:
-            scorable = [
-                scoring_config
-                for scoring_config in self.scoring_queue
-                if (scoring_config.task.llm_model in self.model_scheduler.llm_model_manager.active_models.keys())
-                or (scoring_config.task.llm_model is None)
-            ]
-            if len(scorable) == 0:
-                return
-            self.scoring_queue.remove(scorable[0])
-        scoring_config: ScoringConfig = scorable.pop(0)
+
+        if not self.scoring_queue:
+            return
+
+        # TODO: Filter based on active models before selecting an item to score.
+        scoring_config: ScoringConfig = self.scoring_queue.pop(0)
 
         # here we generate the actual reference
         with Timer(label=f"Generating reference for {scoring_config.task.__class__.__name__}"):
             await scoring_config.task.make_reference(
                 dataset_entry=scoring_config.dataset_entry,
-                model_manager=self.model_scheduler.llm_model_manager,
             )
+            logger.info(f"Reference: {scoring_config.task.reference}")
 
         # and there we then calculate the reward
         reward_pipeline = TaskRegistry.get_task_reward(scoring_config.task)
@@ -103,9 +92,8 @@ async def run_step(self) -> RewardLoggingEvent:
                 response_event=scoring_config.response,
                 challenge=scoring_config.task.query,
                 reference=scoring_config.task.reference,
-                model_id=scoring_config.task.llm_model,
+                model_id=scoring_config.task.llm_model_id,
                 task=scoring_config.task,
-                model_manager=self.model_scheduler.llm_model_manager,
                 task_queue=self.task_queue,
             )
 
@@ -164,7 +152,6 @@ async def run_step(self) -> RewardLoggingEvent:
                 source=source,
             )
         )
-        self.model_scheduler.llm_model_manager.lock.release()
         await asyncio.sleep(0.01)
 
 
diff --git a/prompting/rewards/web_retrieval.py b/prompting/rewards/web_retrieval.py
index 2d33f9fb8..7d67800d3 100644
--- a/prompting/rewards/web_retrieval.py
+++ b/prompting/rewards/web_retrieval.py
@@ -8,6 +8,7 @@
 
 import numpy as np
 import pandas as pd
+import requests
 import whois
 from loguru import logger
 from pydantic import BaseModel
@@ -18,9 +19,36 @@
 from prompting.rewards.relevance import RelevanceRewardModel
 from prompting.rewards.reward import BatchRewardOutput
 from prompting.tasks.base_task import BaseTextTask
+from shared import constants
 from shared.dendrite import DendriteResponseEvent
 from shared.misc import async_lru_cache
 
+
+def get_embeddings(inputs):
+    """
+    Sends a POST request to the local embeddings endpoint and returns the response.
+
+    Args:
+        inputs (str or list of str): A single input string or a list of input strings to embed.
+
+    Returns:
+        dict: JSON response from the embeddings server.
+    """
+    if isinstance(inputs, str):
+        inputs = [inputs]  # convert single string to list
+
+    url = f"{constants.DOCKER_BASE_URL}/v1/embeddings"
+    headers = {"Content-Type": "application/json"}
+    payload = {"input": inputs}
+
+    try:
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        return response.json()
+    except requests.RequestException as e:
+        return {"error": str(e)}
+
+
 MIN_RELEVANT_CHARS = 300
 MIN_MATCH_THRESHOLD = 98
 
@@ -140,8 +168,8 @@ async def domain_age_days(domain: str, fallback_age: int = 1_000_000) -> int:
     @async_lru_cache(maxsize=1000)
     async def _cosine_similarity(self, content1: str, content2: str) -> float:
         """Calculate the cosine similarity between sentence embeddings of the reference and completions."""
-        reference_emb_flatten = self.embedding_model.encode(content1, to_numpy=True).flatten()
-        response_emb_flatten = self.embedding_model.encode(content2, to_numpy=True).flatten()
+        reference_emb_flatten = get_embeddings(content1)["data"][0]["embedding"]
+        response_emb_flatten = get_embeddings(content2)["data"][0]["embedding"]
         return 1.0 - float(spatial.distance.cosine(reference_emb_flatten, response_emb_flatten))
 
     async def score_website_result(
@@ -152,7 +180,6 @@ async def score_website_result(
 
         # Extract domain from URL.
         netloc = extract_main_domain(response_url)
-        logger.debug(f"Scoring url: {response_url}")
 
         if any(term in response_url for term in BLACKLISTED_TERMS):
             logger.debug(f"Domain {response_url} contains blacklisted term, scoring 0")
@@ -192,18 +219,18 @@ async def score_website_result(
             # Content scraped from the URL provided in the completion.
             reference_website_content = DDGDataset.extract_website_content(response_url)
             if not reference_website_content or len(reference_website_content) == 0:
-                logger.debug(f"Failed to extract miner {uid} content from website: {response_url}")
+                # logger.debug(f"Failed to extract miner {uid} content from website: {response_url}")
                 return 0
 
             if fuzz.ratio(response_content, reference_website_content) < MIN_MATCH_THRESHOLD:
-                logger.debug(f"Miner {uid} returned text that doesn't match the website, scoring 0")
+                # logger.debug(f"Miner {uid} returned text that doesn't match the website, scoring 0")
                 return 0
 
             if len(response_relevant) > len(response_content) or len(response_relevant) < MIN_RELEVANT_CHARS:
-                logger.debug(
-                    f"Miner {uid} relevant section is too short (<{MIN_RELEVANT_CHARS} chars) or longer than the whole "
-                    f"website content {len(response_relevant)} > {len(response_content)}"
-                )
+                # logger.debug(
+                #     f"Miner {uid} relevant section is too short (<{MIN_RELEVANT_CHARS} chars) or longer than the whole "
+                #     f"website content {len(response_relevant)} > {len(response_content)}"
+                # )
                 return 0
 
             if response_relevant not in response_content:
@@ -212,10 +239,10 @@ async def score_website_result(
             similarity = await self._cosine_similarity(content1=dataset_entry.query, content2=response_relevant)
             if similarity < PENALIZE_SIM_THRESHOLD:
                 # Penalise if similarity is too low.
-                logger.debug(f"Miner {uid} returned text that doesn't match the query")
+                # logger.debug(f"Miner {uid} returned text that doesn't match the query")
                 return PENALTY
             elif similarity < MIN_SIM_THRESHOLD:
-                logger.debug(f"Miner {uid} returned text has low similarity")
+                # logger.debug(f"Miner {uid} returned text has low similarity")
                 return 0
             return similarity * discount_factor
 
diff --git a/prompting/tasks/MSRv2_task.py b/prompting/tasks/MSRv2_task.py
index 2d3429c1f..2fcb753d8 100644
--- a/prompting/tasks/MSRv2_task.py
+++ b/prompting/tasks/MSRv2_task.py
@@ -4,7 +4,6 @@
 from loguru import logger
 
 from prompting.datasets.random_website import DDGDatasetEntry
-from prompting.llms.model_manager import ModelManager
 from prompting.rewards.MSRv2_reward import MSRv2RewardModel
 from prompting.rewards.reward import BaseRewardConfig, BaseRewardModel
 from prompting.tasks.multi_step_reasoning import MultiStepReasoningTask
@@ -56,11 +55,11 @@ def make_query(self, dataset_entry: DDGDatasetEntry):
         else:
             return self.reference or self.generative_miner_answer
 
-    async def make_reference(self, dataset_entry: Context, model_manager: ModelManager):
+    async def make_reference(self, dataset_entry: Context):
         if self.stage == "generative":
             if random.random() < self.REAL_REFERENCE_PROBABILITY:
                 # Validator's turn to generate the reference
-                reference_attempt = await super().make_reference(dataset_entry, model_manager=model_manager)
+                reference_attempt = await super().make_reference(dataset_entry)
                 self.reference = reference_attempt if isinstance(reference_attempt, str) else None
                 self.validator_generated_reference = self.reference  # Store the validator's generated reference
                 return self.reference
diff --git a/prompting/tasks/base_task.py b/prompting/tasks/base_task.py
index ae6469979..eae68aadd 100644
--- a/prompting/tasks/base_task.py
+++ b/prompting/tasks/base_task.py
@@ -8,10 +8,9 @@
 
 from prompting.llms.apis.gpt_wrapper import LLMMessage, LLMMessages
 from prompting.llms.apis.llm_wrapper import LLMWrapper
-from prompting.llms.model_manager import ModelManager
-from prompting.llms.model_zoo import ModelConfig
 from shared import settings
 from shared.base import DatasetEntry
+from shared.docker_utils import get_generation
 
 
 def CHATTENSOR_SYSTEM_PROMPT():
@@ -59,7 +58,7 @@ class BaseTextTask(BaseTask):
     roles: list[str] | None = None
     messages: list[str] | list[dict] | None = None
     reference: str | None = None
-    llm_model: ModelConfig = None
+    llm_model: str | None = None
     llm_model_id: str = None
     seed: int = Field(default_factory=lambda: random.randint(0, 1000000), allow_mutation=False)
     query_system_prompt: ClassVar[str | None] = None
@@ -83,14 +82,13 @@ def get_model_id_and_seed(self) -> "BaseTextTask":
     async def make_query(self, dataset_entry: DatasetEntry, **kwargs) -> str:
         return self.query
 
-    async def make_reference(self, dataset_entry: DatasetEntry, model_manager: ModelManager | None = None) -> str:
+    async def make_reference(self, dataset_entry: DatasetEntry) -> str:
         return self.reference
 
-    async def generate_reference(self, messages: list[str], model_manager: ModelManager | None = None) -> str:
+    async def generate_reference(self, messages: list[str]) -> str:
         """Generate reference answer to be used for scoring miner completions"""
-        model = await model_manager.get_model(settings.shared_settings.LLM_MODEL[0])
-        self.reference = await model.generate(messages=messages)
-        if self.reference is None:
+        self.reference = await get_generation(messages=messages, model=settings.shared_settings.LLM_MODEL[0])
+        if not self.reference:
             raise Exception("Reference generation failed")
 
         return self.reference
diff --git a/prompting/tasks/inference.py b/prompting/tasks/inference.py
index b8a8826a8..9373f6360 100644
--- a/prompting/tasks/inference.py
+++ b/prompting/tasks/inference.py
@@ -1,15 +1,15 @@
 import random
 from typing import ClassVar
 
+from loguru import logger
 from pydantic import Field, model_validator
 
 from prompting.datasets.sn13 import ChatEntry
-from prompting.llms.model_manager import ModelManager
-from prompting.llms.model_zoo import ModelConfig, ModelZoo
 from prompting.rewards.inference_reward_model import InferenceRewardModel
 from prompting.rewards.reward import BaseRewardConfig, BaseRewardModel
 from prompting.tasks.base_task import BaseTextTask
 from shared import settings
+from shared.docker_utils import get_generation
 
 shared_settings = settings.shared_settings
 
@@ -43,8 +43,7 @@ class InferenceTask(BaseTextTask):
     query: str | list = []
     reference: str | None = None
     system_prompt: str | None = None
-    llm_model: ModelConfig | None = None
-    llm_model_id: str | None = Field(default_factory=lambda: random.choice(ModelZoo.models_configs).llm_model_id)
+    llm_model_id: str | None = Field(default_factory=lambda: random.choice(settings.shared_settings.LLM_MODEL))
     seed: int = Field(default_factory=lambda: random.randint(0, 1_000_000), allow_mutation=False)
     sampling_params: dict[str, float] = shared_settings.SAMPLING_PARAMS.copy()
     messages: list[dict] | None = None
@@ -56,7 +55,6 @@ def random_llm_model_id(self):
             return self
         # self.sampling_params["temperature"] = random.randint(1, 10) / 10
         # self.sampling_params["max_new_tokens"] = random.choice([256, 512, 1024, 2048])
-        self.llm_model = ModelZoo.get_model_by_id(self.llm_model_id)
         return self
 
     async def make_query(self, dataset_entry: ChatEntry) -> str:
@@ -69,17 +67,18 @@ async def make_query(self, dataset_entry: ChatEntry) -> str:
 
         return self.query
 
-    async def make_reference(self, dataset_entry: ChatEntry, model_manager: ModelManager | None = None) -> str:
-        assert model_manager is not None, f"Model manager must be provided for {self.__class__.__name__}"
+    async def make_reference(self, dataset_entry: ChatEntry) -> str:
         # With logits scoring there is no reference, and instead we need to generate the logits based
         # on the miner's completions.
-        if not self.organic and (self.llm_model or self.llm_model_id):
+        logger.info(f"self.llm_model: {self.llm_model}")
+        logger.info(f"self.llm_model_id: {self.llm_model_id}")
+        if self.organic or self.llm_model_id:
             self.reference = ""
             return self.reference
 
-        self.reference = await model_manager.generate(
+        self.reference = await get_generation(
             messages=self.messages,
-            model=self.llm_model,
+            model=self.llm_model_id,
             seed=self.seed,
             sampling_params=self.sampling_params,
         )
diff --git a/prompting/tasks/multi_step_reasoning.py b/prompting/tasks/multi_step_reasoning.py
index e3ebcc99d..3ab18fe96 100644
--- a/prompting/tasks/multi_step_reasoning.py
+++ b/prompting/tasks/multi_step_reasoning.py
@@ -4,7 +4,6 @@
 from loguru import logger
 
 from prompting.datasets.random_website import DDGDatasetEntry
-from prompting.llms.model_manager import ModelManager
 from prompting.rewards.relevance import RelevanceRewardModel
 from prompting.rewards.reward import BaseRewardConfig, BaseRewardModel
 from prompting.tasks.qa import WebQuestionAnsweringTask
@@ -99,7 +98,7 @@ async def _async_generate_reference(self):
                 logger.debug(f"**Total thinking time: {total_thinking_time:.2f} seconds**")
         return steps[-1][1]
 
-    async def make_reference(self, dataset_entry: Context, model_manager: ModelManager | None = None):
+    async def make_reference(self, dataset_entry: Context):
         try:
             logger.debug(f"Generating reference for MSR: {self.messages}")
             # Run the async function in a new event loop
diff --git a/prompting/tasks/qa.py b/prompting/tasks/qa.py
index 1c9a6fa45..fd6143e41 100644
--- a/prompting/tasks/qa.py
+++ b/prompting/tasks/qa.py
@@ -1,7 +1,6 @@
 from typing import ClassVar
 
 from prompting.datasets.random_website import DDGDatasetEntry
-from prompting.llms.model_manager import ModelManager
 from prompting.rewards.relevance import RelevanceRewardModel
 from prompting.rewards.reward import BaseRewardConfig, BaseRewardModel
 from prompting.rewards.rouge import RougeRewardModel
@@ -65,11 +64,7 @@ async def make_query(self, dataset_entry: DDGDatasetEntry):
         self.query = await self.generate_query(messages=[query_prompt])
         return self.query
 
-    async def make_reference(self, dataset_entry: DDGDatasetEntry, model_manager: ModelManager | None = None):
-        assert model_manager is not None, f"Model manager must be provided for {self.__class__.__name__}"
+    async def make_reference(self, dataset_entry: DDGDatasetEntry):
         reference_prompt = REFERENCE_PROMPT_TEMPLATE.format(context=dataset_entry.website_content, question=self.query)
-        self.reference = await self.generate_reference(
-            messages=[{"role": "user", "content": reference_prompt}],
-            model_manager=model_manager,
-        )
+        self.reference = await self.generate_reference(messages=[{"role": "user", "content": reference_prompt}])
         return self.reference
diff --git a/prompting/tasks/task_registry.py b/prompting/tasks/task_registry.py
index 8504afecd..15b4dcd95 100644
--- a/prompting/tasks/task_registry.py
+++ b/prompting/tasks/task_registry.py
@@ -29,16 +29,16 @@ def __hash__(self):
 
 class TaskRegistry(BaseModel):
     task_configs: ClassVar[list[TaskConfig]] = [
-        TaskConfig(task=MSRv2Task, probability=0.10, datasets=[DDGDataset], reward_model=MSRv2RewardConfig),
+        TaskConfig(task=MSRv2Task, probability=0.20, datasets=[DDGDataset], reward_model=MSRv2RewardConfig),
         TaskConfig(
             task=InferenceTask,
-            probability=0.55,
+            probability=0.50,
             datasets=[SN13Dataset],
             reward_model=InferenceRewardConfig,
         ),
         TaskConfig(
             task=WebRetrievalTask,
-            probability=0.35,
+            probability=0.30,
             datasets=[DDGDataset],
             reward_model=WebRetrievalRewardConfig,
         ),
diff --git a/prompting/tasks/web_retrieval.py b/prompting/tasks/web_retrieval.py
index aed37b945..ae4537ab3 100644
--- a/prompting/tasks/web_retrieval.py
+++ b/prompting/tasks/web_retrieval.py
@@ -6,7 +6,6 @@
 from pydantic import Field
 
 from prompting.datasets.random_website import DDGDatasetEntry
-from prompting.llms.model_manager import ModelManager
 from prompting.rewards.reward import BaseRewardConfig, BaseRewardModel
 from prompting.rewards.web_retrieval import WebRetrievalRewardModel
 from prompting.tasks.base_task import BaseTextTask
@@ -48,7 +47,7 @@ async def make_query(self, dataset_entry: DDGDatasetEntry) -> str:
         )
         return self.query
 
-    async def make_reference(self, dataset_entry: DDGDatasetEntry, model_manager: ModelManager | None = None) -> str:
+    async def make_reference(self, dataset_entry: DDGDatasetEntry) -> str:
         dataset_entry.query = self.query
         ref_dict = dataset_entry.model_dump_json()
         self.reference = json.dumps(ref_dict)
diff --git a/run.sh b/run.sh
index caada4ce6..fc7c6e191 100755
--- a/run.sh
+++ b/run.sh
@@ -13,6 +13,8 @@ old_args=$@
 
 bash scripts/install.sh
 
+bash scripts/manage_container.sh
+
 # Loop through all command line arguments
 while [[ $# -gt 0 ]]; do
   arg="$1"
diff --git a/scripts/manage_container.sh b/scripts/manage_container.sh
new file mode 100755
index 000000000..fe34dfaf3
--- /dev/null
+++ b/scripts/manage_container.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+# This script restarts the gpu-app docker container if it's running,
+# or starts it if it's not.
+
+# The directory where docker-compose.yml is located
+COMPOSE_DIR="gpu_container"
+
+# The name of the service to manage
+SERVICE_NAME="gpu-app"
+
+# Change to the compose directory and exit if it fails
+cd "$COMPOSE_DIR" || { echo "Directory $COMPOSE_DIR not found." >&2; exit 1; }
+
+# Check if the service is running.
+# 'docker compose ps -q' will output container IDs if they are up.
+# We check if the output is non-empty.
+if [ -n "$(docker compose ps -q "$SERVICE_NAME")" ]; then
+    echo "Service '$SERVICE_NAME' is running. Restarting..."
+    docker compose restart "$SERVICE_NAME"
+else
+    # This will handle both 'stopped' and 'not-created' states.
+    # The --build flag ensures the image is up-to-date.
+    echo "Service '$SERVICE_NAME' is not running. Starting..."
+    docker compose up -d --build "$SERVICE_NAME"
+fi
+
+# Go back to the original directory
+cd - >/dev/null
+
+echo "Script finished."
diff --git a/shared/constants.py b/shared/constants.py
index a2f8c11e6..e49891e1f 100644
--- a/shared/constants.py
+++ b/shared/constants.py
@@ -1 +1,5 @@
 WHITELISTED_VALIDATORS_UIDS = [5, 518, 674, 966, 502, 520, 993, 24]  # OTF  # WildSageLabs  # Rizzo  # Macrocosmos
+
+DOCKER_BASE_URL = "http://localhost:8000"
+
+SPECIAL_TOKENS = {"mrfakename/mistral-small-3.1-24b-instruct-2503-hf": {"bos_token": "<s>", "eos_token": "</s>"}}
diff --git a/shared/docker_utils.py b/shared/docker_utils.py
new file mode 100644
index 000000000..6a189f34b
--- /dev/null
+++ b/shared/docker_utils.py
@@ -0,0 +1,57 @@
+import requests
+from loguru import logger
+
+from shared import constants
+
+
+async def get_generation(
+    messages: list[str] | list[dict],
+    roles: list[str] | None = None,
+    model: str | None = None,
+    seed: int = None,
+    sampling_params: dict[str, float] = None,
+) -> str:
+    if messages and isinstance(messages[0], dict):
+        dict_messages = messages
+    else:
+        dict_messages = [
+            {"content": message, "role": role} for message, role in zip(messages, roles or ["user"] * len(messages))
+        ]
+    url = f"{constants.DOCKER_BASE_URL}/v1/chat/generate"
+    headers = {"Content-Type": "application/json"}
+    payload = {"messages": dict_messages, "seed": seed, "sampling_params": sampling_params}
+    response = requests.post(url, headers=headers, json=payload)
+    try:
+        json_response = response.json()
+        logger.info(f"Response: {json_response}")
+        return json_response["choices"][0]["message"]["content"]
+    except requests.exceptions.JSONDecodeError:
+        logger.error(f"Error generating response. Status: {response.status_code}, Body: {response.text}")
+        return ""
+
+
+# @async_lru_cache(maxsize=1000)
+async def get_logits(
+    messages: list[str],
+    model: None = None,
+    sampling_params: dict[str, float] = None,
+    seed: int = None,
+    continue_last_message: bool = False,
+    top_logprobs: int = 10,
+):
+    url = f"{constants.DOCKER_BASE_URL}/v1/chat/generate_logits"
+    headers = {"Content-Type": "application/json"}
+    payload = {
+        "messages": messages,
+        "seed": seed,
+        "sampling_params": sampling_params,
+        "top_logprobs": top_logprobs,
+        "continue_last_message": continue_last_message,
+    }
+    response = requests.post(url, headers=headers, json=payload)
+    try:
+        json_response = response.json()
+        return json_response
+    except requests.exceptions.JSONDecodeError:
+        logger.error(f"Error generating logits. Status: {response.status_code}, Body: {response.text}")
+        return ""
diff --git a/shared/epistula.py b/shared/epistula.py
index 8b06b5dd9..02e73ca36 100644
--- a/shared/epistula.py
+++ b/shared/epistula.py
@@ -169,7 +169,7 @@ async def query_miners(
                 )
             else:
                 responses_error += 1
-                logger.error(f"Unknown response type: {response}")
+                # logger.error(f"Unknown response type: {response}")
                 results.append(SynapseStreamResult(uid=uid, exception=f"Unknown response type: {response}"))
 
         logger.info(
@@ -262,8 +262,8 @@ async def make_openai_query(
             stream=True,
             extra_body=extra_body,
         )
-    except BaseException as e:
-        logger.warning(f"Error while querying UID {uid}: {e}")
+    except BaseException:
+        # logger.warning(f"Error while querying UID {uid}: {e}")
         return
 
     if stream:
diff --git a/shared/settings.py b/shared/settings.py
index 206cd1dc8..73497231f 100644
--- a/shared/settings.py
+++ b/shared/settings.py
@@ -133,7 +133,7 @@ class SharedSettings(BaseSettings):
         "temperature": 0.7,
         "top_p": 0.95,
         "top_k": 50,
-        "max_tokens": 512,
+        "max_tokens": 2048,
     }
     LLM_MODEL_RAM: float = Field(70, env="LLM_MODEL_RAM")
     OPENAI_API_KEY: str | None = Field(None, env="OPENAI_API_KEY")
diff --git a/tests/prompting/rewards/test_exact_match.py b/tests/prompting/rewards/test_exact_match.py
index 7d8325f92..e2f3489b8 100644
--- a/tests/prompting/rewards/test_exact_match.py
+++ b/tests/prompting/rewards/test_exact_match.py
@@ -1,342 +1,342 @@
-from typing import Any
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import numpy as np
-import pytest
-
-from prompting.llms.model_manager import ModelManager
-from prompting.rewards.exact_match import (
-    INCORRECT_PENALTY,
-    MAX_VERIFY_TOKENS,
-    MIN_SMOOTH_PENALTY_SCALE,
-    MIN_VERIFY_TOKENS,
-    PARTIAL_PENALTY,
-    TOP_LOGPROBS,
-    VERIFICATION_THRESH_SIM,
-    LogitsRewardModel,
-)
-from prompting.rewards.reward import BatchRewardOutput
-from prompting.tasks.base_task import BaseTextTask
-from shared.dendrite import DendriteResponseEvent
-
-
-@pytest.fixture
-def model_manager():
-    """Mock ModelManager for testing."""
-    manager = MagicMock(spec=ModelManager)
-    model = MagicMock()
-    tokenizer = MagicMock()
-    tokenizer.eos_token = "<|endoftext|>"
-
-    model.tokenizer = tokenizer
-    model.get_max_tokens = AsyncMock(return_value=2048)
-
-    manager.get_model.return_value = model
-
-    async def mock_generate_logits(*args, **kwargs):
-        return {"token1": -0.1, "token2": -0.5, "<|endoftext|>": -1.0}, "prompt"
-
-    manager.generate_logits = AsyncMock(side_effect=mock_generate_logits)
-    return manager
-
-
-@pytest.fixture
-def task():
-    """Mock Task for testing."""
-    task = MagicMock(spec=BaseTextTask)
-    task.llm_model_id = "mockmodel"
-    task.task_messages = [
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": "Tell me a joke."},
-    ]
-    task.sampling_params = {"temperature": 0.7, "max_tokens": 100}
-    return task
-
-
-def create_chat_completion_chunk(
-    content: str = "",
-    logprobs: dict[str, float] | None = None,
-    top_logprobs: int = 5,
-) -> dict[str, Any]:
-    """Return a dict that looks like an OpenAI `ChatCompletionChunk`."""
-
-    # Default log-probabilities if none provided.
-    if logprobs is None:
-        logprobs = {
-            content: -0.1,
-            "token2": -0.5,
-            "token3": -0.6,
-            "token4": -0.7,
-            "<|endoftext|>": -1.0,
-        }
-
-    choice_dict: dict[str, Any] = {
-        "index": 0,
-        "delta": {"role": "assistant", "content": content},
-    }
-
-    # Only include the `logprobs` block when tokens were supplied.
-    if logprobs:
-        choice_dict["logprobs"] = {
-            "content": [
-                {"top_logprobs": [{"token": tok, "logprob": lp} for tok, lp in list(logprobs.items())[:top_logprobs]]}
-            ]
-        }
-    else:
-        choice_dict["logprobs"] = None
-
-    # Assemble the full chunk.
-    chunk_dict: dict[str, Any] = {
-        "id": "chunk_id",
-        "object": "chat.completion.chunk",
-        "created": 1234567890,
-        "model": "VeryStronkModel",
-        "choices": [choice_dict],
-        "usage": None,
-    }
-
-    return chunk_dict
-
-
-async def create_response_event_mock(chunks_all, timings_all, timeout: float = 10) -> MagicMock:
-    completions = ["".join(chunks) for chunks in chunks_all]
-    chunk_dicts_raw = []
-    for chunks in chunks_all:
-        chunk_dicts_raw.append([create_chat_completion_chunk(chunk) for chunk in chunks])
-
-    response_event = MagicMock(spec=DendriteResponseEvent)
-    response_event.stream_results_all_chunks = chunks_all
-    response_event.stream_results_all_chunk_dicts_raw = chunk_dicts_raw
-    response_event.uids = list(range(len(chunks_all)))
-    response_event.stream_results_all_chunks_timings = timings_all
-    response_event.completions = completions
-    response_event.timeout = timeout
-    return response_event
-
-
-@pytest.mark.asyncio
-async def test_correct_completion(model_manager, task):
-    """Test case 1: Correct completion with reward >0.5 and ≤1."""
-    chunks_all = [["Hello", ", ", "world", "!"]]
-    chunks_timings_all = [[0.1, 0.1, 0.1, 0.1]]
-    response_event = await create_response_event_mock(chunks_all, chunks_timings_all)
-    chunk_dicts_raw = []
-    for chunks in chunks_all:
-        chunk_dicts_raw.append([create_chat_completion_chunk(chunk) for chunk in chunks])
-
-    with (
-        patch("prompting.rewards.exact_match.MIN_VERIFY_TOKENS", 2),
-        patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_similarity", return_value=1),
-        patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_contains", return_value=1),
-    ):
-        reward_model = LogitsRewardModel()
-        result = await reward_model.reward(
-            reference="", response_event=response_event, task=task, model_manager=model_manager
-        )
-        assert isinstance(result, BatchRewardOutput)
-        assert len(result.rewards) == 1
-        assert result.rewards[0] == pytest.approx(1.0)
-
-
-@pytest.mark.asyncio
-async def test_mixed_completions(model_manager, task):
-    """Test case 2: One ideal completion, one with missing logprobs penalized."""
-    top_logprobs = 5
-    chunks_timings_all = [[0.1, 0.2, 0.3, 0.4] for _ in range(3)]
-    chunks_all = [["Hello", ", ", "world", "!"], ["Fail", "ed", " ", "completion"], ["Wro", "ng", " ", "completion"]]
-    chunk_dicts_raw: list[list[dict[str, float]]] = []
-
-    correct_logprobs: list[dict[str, float]] = []
-    for part in chunks_all[0]:
-        correct_logprobs.append(create_chat_completion_chunk(part, top_logprobs=top_logprobs))
-    chunk_dicts_raw.append(correct_logprobs)
-
-    incorrect_logprobs: list[dict[str, float]] = []
-    wrong_logprobs: dict[str, float] = {
-        "wrong": -0.1,
-        "log": -5.43,
-        "prob": -8.54,
-        "defined": -11,
-        "<|endoftext|>": -3000000,
-    }
-    for part in chunks_all[1]:
-        incorrect_logprobs.append(create_chat_completion_chunk(part, logprobs=wrong_logprobs))
-    chunk_dicts_raw.append(incorrect_logprobs)
-
-    empty_logprobs: list[dict[str, float]] = []
-    for part in chunks_all[2]:
-        empty_logprobs.append(create_chat_completion_chunk(part, logprobs={}))
-    chunk_dicts_raw.append(empty_logprobs)
-
-    response_event = await create_response_event_mock(chunks_all, chunks_timings_all)
-    response_event.stream_results_all_chunk_dicts_raw = chunk_dicts_raw
-
-    def mock_verify_sim(original_logits, verification_logits):
-        return 1.0 if original_logits and "wrong" not in original_logits else VERIFICATION_THRESH_SIM * 0.9
-
-    with (
-        patch("prompting.rewards.exact_match.MIN_VERIFY_TOKENS", 2),
-        patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_similarity", side_effect=mock_verify_sim),
-        patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_contains", return_value=1),
-    ):
-        reward_model = LogitsRewardModel()
-        result = await reward_model.reward(
-            reference="", response_event=response_event, task=task, model_manager=model_manager
-        )
-
-        assert isinstance(result, BatchRewardOutput)
-        assert len(result.rewards) == len(chunk_dicts_raw)
-        assert 0.2 < result.rewards[0] <= 1.0
-        assert result.rewards[1] == INCORRECT_PENALTY
-        assert result.rewards[2] == INCORRECT_PENALTY
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize(
-    "eos_in_logits, expected_penalty",
-    [
-        (True, None),
-        (False, PARTIAL_PENALTY),
-    ],
-    ids=["eos_present", "eos_missing"],
-)
-async def test_eos_handling(eos_in_logits, expected_penalty, model_manager, task):
-    emitted = ["Hello", ", ", "world", "!"]
-    timings = [[0.1] * len(emitted)]
-    response_event = await create_response_event_mock([emitted], timings)
-    verify_logits = {"tokA": -0.1, "tokB": -0.5}
-    if eos_in_logits:
-        verify_logits["<|endoftext|>"] = -1.0
-    model_manager.generate_logits = AsyncMock(return_value=(verify_logits, "prompt"))
-
-    with (
-        patch("prompting.rewards.exact_match.MIN_VERIFY_TOKENS", 2),
-        patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_similarity", return_value=1),
-        patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_contains", return_value=1),
-    ):
-        reward_model = LogitsRewardModel()
-        result: BatchRewardOutput = await reward_model.reward(
-            reference="",
-            response_event=response_event,
-            task=task,
-            model_manager=model_manager,
-        )
-
-    assert isinstance(result, BatchRewardOutput)
-    assert len(result.rewards) == 1
-    if expected_penalty is None:
-        # eos present.
-        assert result.rewards[0] != PARTIAL_PENALTY
-    else:
-        # eos missing.
-        assert result.rewards[0] == pytest.approx(expected_penalty)
-
-
-def test_verify_logit_similarity():
-    """Test the verify_logit_similarity similarity metric."""
-    original = {f"token{idx}": -0.01 for idx in range(TOP_LOGPROBS)}
-    # Identical distributions -> 1.0.
-    assert LogitsRewardModel.verify_logit_similarity(original, original) == pytest.approx(1.0)
-
-    with patch("prompting.rewards.exact_match.TOP_LOGPROBS", 5):
-        # Disjoint tokens -> near zero.
-        disjoint = {"foo": -0.1, "bar": -0.5, "foo1": -1.0, "bar1": -1.5, "foo2": -2.0}
-        sim = LogitsRewardModel.verify_logit_similarity(original, disjoint)
-        assert sim == pytest.approx(0.0)
-
-        # Partial overlap -> between 0 and 1.
-        original = {"token1": -0.1, "token2": -0.5, "token3": -1.0, "foo1": -1.5, "bar1": -2.0}
-        partial = {"token1": -0.1, "token2": -0.5, "token3": -1.0, "token4": -1.5, "token5": -2.0}
-        sim2 = LogitsRewardModel.verify_logit_similarity(original, partial)
-        assert sim2 == pytest.approx(0.6)
-
-
-def test_smooth_reward_scale():
-    """Test the smooth_reward_scale function under various conditions."""
-    # Test empty timings list.
-    assert LogitsRewardModel.smooth_timings_reward([]) == 0.0
-
-    # Test uniform timings (should give maximum reward).
-    uniform_timings = [0.1, 0.1, 0.1, 0.1, 0.1]
-    assert LogitsRewardModel.smooth_timings_reward(uniform_timings) == pytest.approx(1.0)
-
-    # Test high variance timings (should give minimum reward).
-    high_var_timings = [0.1, 0.1, 15.0, 0.1, 0.1]
-    assert LogitsRewardModel.smooth_timings_reward(high_var_timings) == MIN_SMOOTH_PENALTY_SCALE
-
-    # Test moderate variance timings.
-    moderate_var_timings = [0.3, 0.2, 0.4, 0.1, 0.1]
-    assert LogitsRewardModel.smooth_timings_reward(moderate_var_timings) == pytest.approx(1.0)
-
-    # Test with custom minimum reward.
-    custom_min = 0.8
-    assert LogitsRewardModel.smooth_timings_reward(high_var_timings, min_reward=custom_min) == custom_min
-
-    # Test with single timing value.
-    single_timing = [1.5]
-    assert LogitsRewardModel.smooth_timings_reward(single_timing) == 1.0
-
-
-@pytest.mark.parametrize(
-    "value, min_value, expected",
-    [
-        # Linear mapping.
-        (0.6, 0.2, (0.6 - 0.2) / (1.0 - 0.2)),
-        # Below min clips to 0.0.
-        (0.1, 0.3, 0.0),
-        # Above max clips to 1.0.
-        (1.2, 0.0, 1.0),
-        # At min boundary.
-        (0.3, 0.3, 0.0),
-        # At max boundary.
-        (1.0, 0.3, 1.0),
-    ],
-)
-def test_rescale_various_cases(value, min_value, expected):
-    assert LogitsRewardModel.rescale(value, min_value=min_value) == pytest.approx(expected)
-
-
-@pytest.mark.parametrize(
-    "values, expected",
-    [
-        # All valid.
-        ([[0.1, 1.0], [5.0, 0.1], [6.5]], 0.55),
-        # Mixed values.
-        ([[-1.0, 0.5], [2.0, 0.1]], 1.05),
-        # All negative.
-        ([[-3.0, -0.1], [-2.5]], 1e-6),
-        # Empty lists.
-        ([[], []], 1e-6),
-        # Zeros included.
-        ([[0.0, -1.0], [0.0]], 0.0),
-    ],
-)
-def test_fastest_timing_various_cases(values, expected):
-    assert LogitsRewardModel.fastest_timing(values) == pytest.approx(expected)
-
-
-@pytest.mark.parametrize(
-    "completion_length",
-    [
-        5,
-        (MIN_VERIFY_TOKENS + MAX_VERIFY_TOKENS) // 2,
-        MAX_VERIFY_TOKENS,
-        MAX_VERIFY_TOKENS + 5,
-    ],
-)
-def test_sample_verification_indices_properties(completion_length):
-    indices = LogitsRewardModel.sample_verification_indices(completion_length)
-
-    # Compute expected number of sampled tokens with first and eos indices.
-    expected_k = int(np.clip(completion_length, 1, MAX_VERIFY_TOKENS))
-
-    # The result should have expected_k samples plus one EOS index.
-    assert isinstance(indices, list)
-    assert len(indices) == expected_k
-    assert indices == sorted(indices)
-    assert indices[-1] == completion_length
-    # All other indices should be in the range [0, completion_length).
-    sample_indices = indices[:-1]
-    assert all(0 <= idx < completion_length for idx in sample_indices)
-    # No duplicates overall.
-    assert len(set(indices)) == len(indices)
+# from typing import Any
+# from unittest.mock import AsyncMock, MagicMock, patch
+
+# import numpy as np
+# import pytest
+
+# from prompting.llms.model_manager import ModelManager
+# from prompting.rewards.exact_match import (
+#     INCORRECT_PENALTY,
+#     MAX_VERIFY_TOKENS,
+#     MIN_SMOOTH_PENALTY_SCALE,
+#     MIN_VERIFY_TOKENS,
+#     PARTIAL_PENALTY,
+#     TOP_LOGPROBS,
+#     VERIFICATION_THRESH_SIM,
+#     LogitsRewardModel,
+# )
+# from prompting.rewards.reward import BatchRewardOutput
+# from prompting.tasks.base_task import BaseTextTask
+# from shared.dendrite import DendriteResponseEvent
+
+
+# @pytest.fixture
+# def model_manager():
+#     """Mock ModelManager for testing."""
+#     manager = MagicMock(spec=ModelManager)
+#     model = MagicMock()
+#     tokenizer = MagicMock()
+#     tokenizer.eos_token = "<|endoftext|>"
+
+#     model.tokenizer = tokenizer
+#     model.get_max_tokens = AsyncMock(return_value=2048)
+
+#     manager.get_model.return_value = model
+
+#     async def mock_generate_logits(*args, **kwargs):
+#         return {"token1": -0.1, "token2": -0.5, "<|endoftext|>": -1.0}, "prompt"
+
+#     manager.generate_logits = AsyncMock(side_effect=mock_generate_logits)
+#     return manager
+
+
+# @pytest.fixture
+# def task():
+#     """Mock Task for testing."""
+#     task = MagicMock(spec=BaseTextTask)
+#     task.llm_model_id = "mockmodel"
+#     task.task_messages = [
+#         {"role": "system", "content": "You are a helpful assistant."},
+#         {"role": "user", "content": "Tell me a joke."},
+#     ]
+#     task.sampling_params = {"temperature": 0.7, "max_tokens": 100}
+#     return task
+
+
+# def create_chat_completion_chunk(
+#     content: str = "",
+#     logprobs: dict[str, float] | None = None,
+#     top_logprobs: int = 5,
+# ) -> dict[str, Any]:
+#     """Return a dict that looks like an OpenAI `ChatCompletionChunk`."""
+
+#     # Default log-probabilities if none provided.
+#     if logprobs is None:
+#         logprobs = {
+#             content: -0.1,
+#             "token2": -0.5,
+#             "token3": -0.6,
+#             "token4": -0.7,
+#             "<|endoftext|>": -1.0,
+#         }
+
+#     choice_dict: dict[str, Any] = {
+#         "index": 0,
+#         "delta": {"role": "assistant", "content": content},
+#     }
+
+#     # Only include the `logprobs` block when tokens were supplied.
+#     if logprobs:
+#         choice_dict["logprobs"] = {
+#             "content": [
+#                 {"top_logprobs": [{"token": tok, "logprob": lp} for tok, lp in list(logprobs.items())[:top_logprobs]]}
+#             ]
+#         }
+#     else:
+#         choice_dict["logprobs"] = None
+
+#     # Assemble the full chunk.
+#     chunk_dict: dict[str, Any] = {
+#         "id": "chunk_id",
+#         "object": "chat.completion.chunk",
+#         "created": 1234567890,
+#         "model": "VeryStronkModel",
+#         "choices": [choice_dict],
+#         "usage": None,
+#     }
+
+#     return chunk_dict
+
+
+# async def create_response_event_mock(chunks_all, timings_all, timeout: float = 10) -> MagicMock:
+#     completions = ["".join(chunks) for chunks in chunks_all]
+#     chunk_dicts_raw = []
+#     for chunks in chunks_all:
+#         chunk_dicts_raw.append([create_chat_completion_chunk(chunk) for chunk in chunks])
+
+#     response_event = MagicMock(spec=DendriteResponseEvent)
+#     response_event.stream_results_all_chunks = chunks_all
+#     response_event.stream_results_all_chunk_dicts_raw = chunk_dicts_raw
+#     response_event.uids = list(range(len(chunks_all)))
+#     response_event.stream_results_all_chunks_timings = timings_all
+#     response_event.completions = completions
+#     response_event.timeout = timeout
+#     return response_event
+
+
+# @pytest.mark.asyncio
+# async def test_correct_completion(model_manager, task):
+#     """Test case 1: Correct completion with reward >0.5 and ≤1."""
+#     chunks_all = [["Hello", ", ", "world", "!"]]
+#     chunks_timings_all = [[0.1, 0.1, 0.1, 0.1]]
+#     response_event = await create_response_event_mock(chunks_all, chunks_timings_all)
+#     chunk_dicts_raw = []
+#     for chunks in chunks_all:
+#         chunk_dicts_raw.append([create_chat_completion_chunk(chunk) for chunk in chunks])
+
+#     with (
+#         patch("prompting.rewards.exact_match.MIN_VERIFY_TOKENS", 2),
+#         patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_similarity", return_value=1),
+#         patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_contains", return_value=1),
+#     ):
+#         reward_model = LogitsRewardModel()
+#         result = await reward_model.reward(
+#             reference="", response_event=response_event, task=task, model_manager=model_manager
+#         )
+#         assert isinstance(result, BatchRewardOutput)
+#         assert len(result.rewards) == 1
+#         assert result.rewards[0] == pytest.approx(1.0)
+
+
+# @pytest.mark.asyncio
+# async def test_mixed_completions(model_manager, task):
+#     """Test case 2: One ideal completion, one with missing logprobs penalized."""
+#     top_logprobs = 5
+#     chunks_timings_all = [[0.1, 0.2, 0.3, 0.4] for _ in range(3)]
+#     chunks_all = [["Hello", ", ", "world", "!"], ["Fail", "ed", " ", "completion"], ["Wro", "ng", " ", "completion"]]
+#     chunk_dicts_raw: list[list[dict[str, float]]] = []
+
+#     correct_logprobs: list[dict[str, float]] = []
+#     for part in chunks_all[0]:
+#         correct_logprobs.append(create_chat_completion_chunk(part, top_logprobs=top_logprobs))
+#     chunk_dicts_raw.append(correct_logprobs)
+
+#     incorrect_logprobs: list[dict[str, float]] = []
+#     wrong_logprobs: dict[str, float] = {
+#         "wrong": -0.1,
+#         "log": -5.43,
+#         "prob": -8.54,
+#         "defined": -11,
+#         "<|endoftext|>": -3000000,
+#     }
+#     for part in chunks_all[1]:
+#         incorrect_logprobs.append(create_chat_completion_chunk(part, logprobs=wrong_logprobs))
+#     chunk_dicts_raw.append(incorrect_logprobs)
+
+#     empty_logprobs: list[dict[str, float]] = []
+#     for part in chunks_all[2]:
+#         empty_logprobs.append(create_chat_completion_chunk(part, logprobs={}))
+#     chunk_dicts_raw.append(empty_logprobs)
+
+#     response_event = await create_response_event_mock(chunks_all, chunks_timings_all)
+#     response_event.stream_results_all_chunk_dicts_raw = chunk_dicts_raw
+
+#     def mock_verify_sim(original_logits, verification_logits):
+#         return 1.0 if original_logits and "wrong" not in original_logits else VERIFICATION_THRESH_SIM * 0.9
+
+#     with (
+#         patch("prompting.rewards.exact_match.MIN_VERIFY_TOKENS", 2),
+#         patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_similarity", side_effect=mock_verify_sim),
+#         patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_contains", return_value=1),
+#     ):
+#         reward_model = LogitsRewardModel()
+#         result = await reward_model.reward(
+#             reference="", response_event=response_event, task=task, model_manager=model_manager
+#         )
+
+#         assert isinstance(result, BatchRewardOutput)
+#         assert len(result.rewards) == len(chunk_dicts_raw)
+#         assert 0.2 < result.rewards[0] <= 1.0
+#         assert result.rewards[1] == INCORRECT_PENALTY
+#         assert result.rewards[2] == INCORRECT_PENALTY
+
+
+# @pytest.mark.asyncio
+# @pytest.mark.parametrize(
+#     "eos_in_logits, expected_penalty",
+#     [
+#         (True, None),
+#         (False, PARTIAL_PENALTY),
+#     ],
+#     ids=["eos_present", "eos_missing"],
+# )
+# async def test_eos_handling(eos_in_logits, expected_penalty, model_manager, task):
+#     emitted = ["Hello", ", ", "world", "!"]
+#     timings = [[0.1] * len(emitted)]
+#     response_event = await create_response_event_mock([emitted], timings)
+#     verify_logits = {"tokA": -0.1, "tokB": -0.5}
+#     if eos_in_logits:
+#         verify_logits["<|endoftext|>"] = -1.0
+#     model_manager.generate_logits = AsyncMock(return_value=(verify_logits, "prompt"))
+
+#     with (
+#         patch("prompting.rewards.exact_match.MIN_VERIFY_TOKENS", 2),
+#         patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_similarity", return_value=1),
+#         patch("prompting.rewards.exact_match.LogitsRewardModel.verify_logit_contains", return_value=1),
+#     ):
+#         reward_model = LogitsRewardModel()
+#         result: BatchRewardOutput = await reward_model.reward(
+#             reference="",
+#             response_event=response_event,
+#             task=task,
+#             model_manager=model_manager,
+#         )
+
+#     assert isinstance(result, BatchRewardOutput)
+#     assert len(result.rewards) == 1
+#     if expected_penalty is None:
+#         # eos present.
+#         assert result.rewards[0] != PARTIAL_PENALTY
+#     else:
+#         # eos missing.
+#         assert result.rewards[0] == pytest.approx(expected_penalty)
+
+
+# def test_verify_logit_similarity():
+#     """Test the verify_logit_similarity similarity metric."""
+#     original = {f"token{idx}": -0.01 for idx in range(TOP_LOGPROBS)}
+#     # Identical distributions -> 1.0.
+#     assert LogitsRewardModel.verify_logit_similarity(original, original) == pytest.approx(1.0)
+
+#     with patch("prompting.rewards.exact_match.TOP_LOGPROBS", 5):
+#         # Disjoint tokens -> near zero.
+#         disjoint = {"foo": -0.1, "bar": -0.5, "foo1": -1.0, "bar1": -1.5, "foo2": -2.0}
+#         sim = LogitsRewardModel.verify_logit_similarity(original, disjoint)
+#         assert sim == pytest.approx(0.0)
+
+#         # Partial overlap -> between 0 and 1.
+#         original = {"token1": -0.1, "token2": -0.5, "token3": -1.0, "foo1": -1.5, "bar1": -2.0}
+#         partial = {"token1": -0.1, "token2": -0.5, "token3": -1.0, "token4": -1.5, "token5": -2.0}
+#         sim2 = LogitsRewardModel.verify_logit_similarity(original, partial)
+#         assert sim2 == pytest.approx(0.6)
+
+
+# def test_smooth_reward_scale():
+#     """Test the smooth_reward_scale function under various conditions."""
+#     # Test empty timings list.
+#     assert LogitsRewardModel.smooth_timings_reward([]) == 0.0
+
+#     # Test uniform timings (should give maximum reward).
+#     uniform_timings = [0.1, 0.1, 0.1, 0.1, 0.1]
+#     assert LogitsRewardModel.smooth_timings_reward(uniform_timings) == pytest.approx(1.0)
+
+#     # Test high variance timings (should give minimum reward).
+#     high_var_timings = [0.1, 0.1, 15.0, 0.1, 0.1]
+#     assert LogitsRewardModel.smooth_timings_reward(high_var_timings) == MIN_SMOOTH_PENALTY_SCALE
+
+#     # Test moderate variance timings.
+#     moderate_var_timings = [0.3, 0.2, 0.4, 0.1, 0.1]
+#     assert LogitsRewardModel.smooth_timings_reward(moderate_var_timings) == pytest.approx(1.0)
+
+#     # Test with custom minimum reward.
+#     custom_min = 0.8
+#     assert LogitsRewardModel.smooth_timings_reward(high_var_timings, min_reward=custom_min) == custom_min
+
+#     # Test with single timing value.
+#     single_timing = [1.5]
+#     assert LogitsRewardModel.smooth_timings_reward(single_timing) == 1.0
+
+
+# @pytest.mark.parametrize(
+#     "value, min_value, expected",
+#     [
+#         # Linear mapping.
+#         (0.6, 0.2, (0.6 - 0.2) / (1.0 - 0.2)),
+#         # Below min clips to 0.0.
+#         (0.1, 0.3, 0.0),
+#         # Above max clips to 1.0.
+#         (1.2, 0.0, 1.0),
+#         # At min boundary.
+#         (0.3, 0.3, 0.0),
+#         # At max boundary.
+#         (1.0, 0.3, 1.0),
+#     ],
+# )
+# def test_rescale_various_cases(value, min_value, expected):
+#     assert LogitsRewardModel.rescale(value, min_value=min_value) == pytest.approx(expected)
+
+
+# @pytest.mark.parametrize(
+#     "values, expected",
+#     [
+#         # All valid.
+#         ([[0.1, 1.0], [5.0, 0.1], [6.5]], 0.55),
+#         # Mixed values.
+#         ([[-1.0, 0.5], [2.0, 0.1]], 1.05),
+#         # All negative.
+#         ([[-3.0, -0.1], [-2.5]], 1e-6),
+#         # Empty lists.
+#         ([[], []], 1e-6),
+#         # Zeros included.
+#         ([[0.0, -1.0], [0.0]], 0.0),
+#     ],
+# )
+# def test_fastest_timing_various_cases(values, expected):
+#     assert LogitsRewardModel.fastest_timing(values) == pytest.approx(expected)
+
+
+# @pytest.mark.parametrize(
+#     "completion_length",
+#     [
+#         5,
+#         (MIN_VERIFY_TOKENS + MAX_VERIFY_TOKENS) // 2,
+#         MAX_VERIFY_TOKENS,
+#         MAX_VERIFY_TOKENS + 5,
+#     ],
+# )
+# def test_sample_verification_indices_properties(completion_length):
+#     indices = LogitsRewardModel.sample_verification_indices(completion_length)
+
+#     # Compute expected number of sampled tokens with first and eos indices.
+#     expected_k = int(np.clip(completion_length, 1, MAX_VERIFY_TOKENS))
+
+#     # The result should have expected_k samples plus one EOS index.
+#     assert isinstance(indices, list)
+#     assert len(indices) == expected_k
+#     assert indices == sorted(indices)
+#     assert indices[-1] == completion_length
+#     # All other indices should be in the range [0, completion_length).
+#     sample_indices = indices[:-1]
+#     assert all(0 <= idx < completion_length for idx in sample_indices)
+#     # No duplicates overall.
+#     assert len(set(indices)) == len(indices)