astrai/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.8"
+__version__ = "1.3.7"
 __author__ = "ViperEkura"
 
 from astrai.config import (

astrai/inference/__init__.py

@@ -30,14 +30,10 @@ from astrai.inference.api.openai import OpenAIResponseBuilder
 from astrai.inference.core import (
     STOP,
     Allocator,
-    CacheView,
-    ContiguousCache,
-    ContiguousCacheView,
     Executor,
     InferenceScheduler,
     KVCache,
-    PageCache,
+    KvcacheView,
-    PageCacheView,
     PagePool,
     PrefixCache,
     Storage,
@@ -67,12 +63,8 @@ __all__ = [
     "TaskManager",
     "TaskStatus",
     "Allocator",
-    "CacheView",
     "KVCache",
-    "ContiguousCache",
+    "KvcacheView",
-    "ContiguousCacheView",
-    "PageCache",
-    "PageCacheView",
     "PagePool",
     "PrefixCache",
     "Storage",

astrai/inference/core/__init__.py

@@ -2,12 +2,8 @@
 
 from astrai.inference.core.cache import (
     Allocator,
-    CacheView,
-    ContiguousCache,
-    ContiguousCacheView,
     KVCache,
-    PageCache,
+    KvcacheView,
-    PageCacheView,
     PagePool,
     PrefixCache,
     Storage,
@@ -20,12 +16,8 @@ from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
 
 __all__ = [
     "Allocator",
-    "CacheView",
     "KVCache",
-    "ContiguousCache",
+    "KvcacheView",
-    "ContiguousCacheView",
-    "PageCache",
-    "PageCacheView",
     "PagePool",
     "PrefixCache",
     "Storage",

astrai/inference/core/cache.py

@@ -1,5 +1,4 @@
 import threading
-from abc import ABC, abstractmethod
 from collections import OrderedDict
 from typing import Callable, Dict, List, Optional, Tuple
 
@@ -63,7 +62,6 @@ class Allocator:
 
     def touch(self, idx: int):
         with self._lock:
-            if idx in self._lru:
             self._lru.move_to_end(idx)
 
 
@@ -276,42 +274,7 @@ class Storage:
         return k, v
 
 
-class CacheView(ABC):
+class KvcacheView:
-    """Abstract view passed to attention layers for KV-cache I/O."""
-
-    @abstractmethod
-    def write(self, layer_id: int, k: Tensor, v: Tensor): ...
-
-    @abstractmethod
-    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: ...
-
-
-class KVCache(ABC):
-    """Abstract KV-cache facade for scheduler/executor."""
-
-    @abstractmethod
-    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: ...
-
-    @abstractmethod
-    def task_free(self, task_id: str): ...
-
-    @abstractmethod
-    def task_extend(self, task_id: str, pos: int) -> bool: ...
-
-    @abstractmethod
-    def bind_tasks(
-        self, task_ids: List[str], total_len: int, device: torch.device
-    ) -> CacheView: ...
-
-    def task_cached(self, task_id: str) -> int:
-        return 0
-
-    def task_record_hashes(
-        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
-    ): ...
-
-
-class PageCacheView(CacheView):
     """Bundles Storage + page_table + total_len for attention layers."""
 
     def __init__(self, storage: Storage, page_table: Tensor, total_len: int = 0):
@@ -327,8 +290,8 @@ class PageCacheView(CacheView):
         return self._storage.gather(layer_id, self._page_table, self._total_len)
 
 
-class PageCache(KVCache):
+class KVCache:
-    """Paged KV-cache with prefix sharing."""
+    """Facade: page management + KV-cache I/O for continuous batching."""
 
     def __init__(
         self,
@@ -398,102 +361,8 @@ class PageCache(KVCache):
         for i in range(start_logical_page, full_pages):
             self._pool.record(page_table[i], prompt_ids, i)
 
-    def bind_tasks(
+    def make_table_tensor(self, task_ids: List[str], device: torch.device) -> Tensor:
-        self, task_ids: List[str], total_len: int, device: torch.device
+        return self._table.table_tensor(task_ids, device)
-    ) -> PageCacheView:
-        page_table = self._table.table_tensor(task_ids, device)
-        return PageCacheView(self._storage, page_table, total_len)
 
-
+    def bind(self, page_table: Tensor, total_len: int = 0) -> KvcacheView:
-class ContiguousCacheView(CacheView):
+        return KvcacheView(self._storage, page_table, total_len)
-    """Contiguous KV-cache view for attention layers."""
-
-    def __init__(
-        self, cache: "ContiguousCache", batch_indices: Tensor, total_len: int = 0
-    ):
-        self._cache = cache
-        self._batch_indices = batch_indices
-        self._total_len = total_len
-
-    def write(self, layer_id: int, k: Tensor, v: Tensor):
-        seq_len = k.size(1)
-        start_pos = self._total_len - seq_len
-        indices = self._batch_indices
-        self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k
-        self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v
-        new_len = start_pos + seq_len
-        for s in indices.tolist():
-            cur = self._cache._slot_len.get(s, 0)
-            if new_len > cur:
-                self._cache._slot_len[s] = new_len
-
-    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
-        max_len = max(
-            self._cache._slot_len.get(int(s), 0) for s in self._batch_indices.tolist()
-        )
-        indices = self._batch_indices
-        k = self._cache.k[layer_id, indices, :max_len]
-        v = self._cache.v[layer_id, indices, :max_len]
-        return k, v
-
-
-class ContiguousCache(KVCache):
-    """Contiguous per-slot KV cache (default implementation)."""
-
-    def __init__(
-        self,
-        n_layers: int,
-        max_batch_size: int,
-        max_seq_len: int,
-        n_kv_heads: int,
-        head_dim: int,
-        device: torch.device,
-        dtype: torch.dtype,
-    ):
-        self.max_seq_len = max_seq_len
-        self.k = torch.zeros(
-            n_layers,
-            max_batch_size,
-            max_seq_len,
-            n_kv_heads,
-            head_dim,
-            device=device,
-            dtype=dtype,
-        )
-        self.v = torch.zeros(
-            n_layers,
-            max_batch_size,
-            max_seq_len,
-            n_kv_heads,
-            head_dim,
-            device=device,
-            dtype=dtype,
-        )
-        self._slot_len: Dict[int, int] = {}
-        self._task_slot: Dict[str, int] = {}
-        self._free_slots = list(range(max_batch_size))
-        self._device = device
-
-    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
-        if not self._free_slots:
-            return False
-        slot = self._free_slots.pop(0)
-        self._task_slot[task_id] = slot
-        self._slot_len[slot] = 0
-        return True
-
-    def task_free(self, task_id: str):
-        slot = self._task_slot.pop(task_id, None)
-        if slot is not None:
-            self._slot_len.pop(slot, None)
-            self._free_slots.append(slot)
-
-    def task_extend(self, task_id: str, pos: int) -> bool:
-        return pos < self.max_seq_len
-
-    def bind_tasks(
-        self, task_ids: List[str], total_len: int, device: torch.device
-    ) -> ContiguousCacheView:
-        slots = [self._task_slot[tid] for tid in task_ids]
-        batch_indices = torch.tensor(slots, dtype=torch.long, device=device)
-        return ContiguousCacheView(self, batch_indices, total_len)

astrai/inference/core/executor.py

@@ -19,13 +19,13 @@ class Executor:
         self,
         model: AutoModel,
         tokenizer: AutoTokenizer,
-        kv_cache: KVCache,
+        page_cache: KVCache,
         device: Optional[str] = None,
         dtype: Optional[torch.dtype] = None,
     ):
         self.model = model
         self.tokenizer = tokenizer
-        self.kv_cache = kv_cache
+        self.page_cache = page_cache
         self.device = device or next(model.parameters()).device
         self.dtype = dtype or next(model.parameters()).dtype
 
@@ -43,6 +43,7 @@ class Executor:
         )
 
         task_ids = [t.task_id for t in tasks]
+        page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
 
         with torch.inference_mode():
             self.model(
@@ -52,7 +53,7 @@ class Executor:
                 )
                 .unsqueeze(0)
                 .expand(batch_sz, -1),
-                paged_cache=self.kv_cache.bind_tasks(task_ids, prompt_len, self.device),
+                paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
             )
 
     def execute_decode(self, tasks: List[Task]) -> List[int]:
@@ -71,6 +72,7 @@ class Executor:
         total_len = position_ids.max().item() + 1
 
         task_ids = [t.task_id for t in tasks]
+        page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
 
         temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
         top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
@@ -79,7 +81,7 @@ class Executor:
         with torch.inference_mode():
             outputs = self.model(
                 input_ids.unsqueeze(1),
-                paged_cache=self.kv_cache.bind_tasks(task_ids, total_len, self.device),
+                paged_cache=self.page_cache.bind(page_tables, total_len=total_len),
                 position_ids=position_ids.unsqueeze(1),
             )
             logits = outputs["logits"][:, -1, :]

astrai/inference/core/scheduler.py

@@ -4,7 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 
-from astrai.inference.core.cache import ContiguousCache, KVCache
+from astrai.inference.core.cache import KVCache
 from astrai.inference.core.executor import Executor
 from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
 from astrai.model.automodel import AutoModel
@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
 
 
 class InferenceScheduler:
-    """Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups)."""
+    """Four-phase continuous batching loop: cleanup -> refill -> prefill -> decode."""
 
     def __init__(
         self,
@@ -23,9 +23,9 @@ class InferenceScheduler:
         max_batch_size: int = 16,
         max_seq_len: Optional[int] = None,
         max_prompt_len: int = 2048,
+        page_size: int = 64,
         device: Optional[str] = None,
         dtype: Optional[torch.dtype] = None,
-        cache: Optional[KVCache] = None,
     ):
         config = model.config
 
@@ -41,17 +41,16 @@ class InferenceScheduler:
         self.device = device or next(model.parameters()).device
         self.dtype = dtype or next(model.parameters()).dtype
 
-        head_dim = config.dim // config.n_heads
+        n_pages = (
+            max_batch_size * (self.max_seq_len + page_size) + page_size - 1
+        ) // page_size
 
-        if cache is not None:
+        self._page_cache = KVCache(
-            self._cache = cache
-        else:
-            self._cache = ContiguousCache(
             config.n_layers,
-                max_batch_size,
+            n_pages,
-                self.max_seq_len,
+            page_size,
             config.n_kv_heads,
-                head_dim,
+            config.dim // config.n_heads,
             self.device,
             self.dtype,
         )
@@ -66,7 +65,7 @@ class InferenceScheduler:
         self._executor = Executor(
             model=model,
             tokenizer=tokenizer,
-            kv_cache=self._cache,
+            page_cache=self._page_cache,
             device=self.device,
             dtype=self.dtype,
         )
@@ -79,19 +78,18 @@ class InferenceScheduler:
 
     def remove_task(self, task_id: str):
         for task in self._task_mgr.remove_task(task_id):
-            self._cache.task_free(task.task_id)
+            self._page_cache.task_free(task.task_id)
 
     def get_stats(self) -> Dict[str, Any]:
         return self._task_mgr.get_stats()
 
     def _run_generation_loop(self):
         stop_ids = self._task_mgr.tokenizer.stop_ids
-        cache = self._cache
         try:
             while not self._stop_event.is_set():
                 finished = self._task_mgr.remove_finished_tasks(stop_ids)
                 for task in finished:
-                    cache.task_free(task.task_id)
+                    self._page_cache.task_free(task.task_id)
 
                 active = self._task_mgr.get_active_tasks()
                 available = self._task_mgr.max_batch_size - len(active)
@@ -99,7 +97,7 @@ class InferenceScheduler:
                     candidates = self._task_mgr.pull_candidates(available)
                     failed = []
                     for task in candidates:
-                        if cache.task_alloc(task.task_id, task.prompt_ids):
+                        if self._page_cache.task_alloc(task.task_id, task.prompt_ids):
                             self._task_mgr.activate(task)
                         else:
                             failed.append(task)
@@ -114,7 +112,7 @@ class InferenceScheduler:
                     t
                     for t in self._task_mgr.get_active_tasks()
                     if t.output_tokens == 0
-                    and cache.task_cached(t.task_id) < len(t.prompt_ids)
+                    and self._page_cache.task_cached(t.task_id) < len(t.prompt_ids)
                 ]
                 if to_prefill:
                     for t in to_prefill:
@@ -124,34 +122,36 @@ class InferenceScheduler:
                     for t in to_prefill:
                         key = (
                             len(t.prompt_ids),
-                            cache.task_cached(t.task_id),
+                            self._page_cache.task_cached(t.task_id),
                         )
                         groups.setdefault(key, []).append(t)
 
                     for (prompt_len, start_pos), group in groups.items():
                         self._executor.execute_prefill(group, prompt_len, start_pos)
-                        start_logical_page = start_pos // getattr(
+                        start_logical_page = start_pos // self._page_cache.page_size
-                            cache, "page_size", 64
-                        )
                         for t in group:
-                            cache.task_record_hashes(
+                            self._page_cache.task_record_hashes(
-                                t.task_id, t.prompt_ids, start_logical_page
+                                t.task_id,
+                                t.prompt_ids,
+                                start_logical_page=start_logical_page,
                             )
 
                 pos_groups: Dict[int, List[Task]] = {}
                 for t in self._task_mgr.get_active_tasks():
                     pos_groups.setdefault(t.next_pos, []).append(t)
 
-                for next_pos in sorted(pos_groups.keys()):
+                if pos_groups:
-                    group = sorted(pos_groups[next_pos], key=lambda t: t.task_id)
+                    best_key = max(pos_groups, key=lambda k: len(pos_groups[k]))
+                    group = sorted(pos_groups[best_key], key=lambda t: t.task_id)
 
                     valid: List[Task] = []
                     for t in group:
-                        if cache.task_extend(t.task_id, t.next_pos):
+                        if self._page_cache.task_extend(t.task_id, t.next_pos):
                             valid.append(t)
                         else:
                             t.status = TaskStatus.ABORTED
-                            self._task_mgr.invoke_callback(t.task_id, STOP)
+                            if t.stream_callback:
+                                t.stream_callback(STOP)
 
                     if valid:
                         next_tokens = self._executor.execute_decode(valid)
@@ -159,23 +159,32 @@ class InferenceScheduler:
                         for t, ntok in zip(valid, next_tokens):
                             t.output_ids.append(ntok)
                             t.output_tokens += 1
-                            self._task_mgr.invoke_callback(
+                            pos = t.input_tokens + t.output_tokens
-                                t.task_id,
+                            extend_ok = self._page_cache.task_extend(t.task_id, pos)
-                                self._task_mgr.tokenizer.decode([ntok]),
+                            if t.stream_callback:
+                                t.stream_callback(
+                                    self._task_mgr.tokenizer.decode([ntok])
                                 )
+                            if not extend_ok:
+                                t.status = TaskStatus.ABORTED
+                                if t.stream_callback:
+                                    t.stream_callback(STOP)
 
                         for t in valid:
                             if t.is_finished(stop_ids):
-                                self._task_mgr.invoke_callback(t.task_id, STOP)
+                                if t.stream_callback:
+                                    t.stream_callback(STOP)
 
         except Exception as e:
             self._stop_event.set()
             logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
             for task in self._task_mgr.get_active_tasks():
-                self._task_mgr.invoke_callback(task.task_id, STOP)
+                if task.stream_callback:
-                cache.task_free(task.task_id)
+                    task.stream_callback(STOP)
+                self._page_cache.task_free(task.task_id)
             for task in self._task_mgr.get_waiting_tasks():
-                self._task_mgr.invoke_callback(task.task_id, STOP)
+                if task.stream_callback:
+                    task.stream_callback(STOP)
             self._task_mgr.clear_queues()
 
     def start(self):
@@ -193,10 +202,12 @@ class InferenceScheduler:
             self._loop_thread.join(timeout=2.0)
         self._loop_thread = None
         for task in self._task_mgr.get_active_tasks():
-            self._task_mgr.invoke_callback(task.task_id, STOP)
+            if task.stream_callback:
-            self._cache.task_free(task.task_id)
+                task.stream_callback(STOP)
+            self._page_cache.task_free(task.task_id)
         for task in self._task_mgr.get_waiting_tasks():
-            self._task_mgr.invoke_callback(task.task_id, STOP)
+            if task.stream_callback:
+                task.stream_callback(STOP)
         self._task_mgr.clear_queues()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()

astrai/inference/core/task.py

@@ -33,6 +33,7 @@ class Task:
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = 50,
+        stream_callback: Optional[Callable[[str], None]] = None,
     ):
         self.task_id = task_id
         self.prompt_ids = prompt_ids
@@ -47,6 +48,7 @@ class Task:
         self.output_tokens: int = 0
         self.arrival_time = time.time()
         self.finish_time: Optional[float] = None
+        self.stream_callback = stream_callback
 
     @property
     def next_pos(self) -> int:
@@ -77,7 +79,6 @@ class TaskManager:
 
         self.waiting_queue: Deque[Task] = deque()
         self.active_tasks: List[Task] = []
-        self._callbacks: Dict[str, Callable[[str], None]] = {}
 
         self._task_event = threading.Event()
         self._lock = threading.Lock()
@@ -116,13 +117,12 @@ class TaskManager:
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
+            stream_callback=stream_callback,
         )
 
         with self._lock:
             self.waiting_queue.append(task)
             self._total_tasks += 1
-            if stream_callback:
-                self._callbacks[task_id] = stream_callback
 
         self._task_event.set()
         return task_id
@@ -134,14 +134,8 @@ class TaskManager:
                 t for t in self.waiting_queue if t.task_id != task_id
             )
             self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
-            self._callbacks.pop(task_id, None)
         return removed_active
 
-    def invoke_callback(self, task_id: str, token: str):
-        cb = self._callbacks.get(task_id)
-        if cb:
-            cb(token)
-
     def get_stats(self) -> Dict[str, Any]:
         return {
             "total_tasks": self._total_tasks,
@@ -210,7 +204,6 @@ class TaskManager:
         with self._lock:
             self.waiting_queue.clear()
             self.active_tasks.clear()
-            self._callbacks.clear()
 
     def wake(self):
         self._task_event.set()

astrai/inference/engine.py

@@ -8,7 +8,6 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple,
 import torch
 import torch.nn as nn
 
-from astrai.inference.core.cache import KVCache
 from astrai.inference.core.scheduler import InferenceScheduler
 from astrai.inference.core.task import STOP
 from astrai.tokenize import AutoTokenizer
@@ -102,7 +101,6 @@ class InferenceEngine:
         max_seq_len: Optional[int] = None,
         max_prompt_len: int = 2048,
         page_size: int = 128,
-        cache: Optional[KVCache] = None,
     ):
         self.model = model
         self.tokenizer = tokenizer
@@ -112,7 +110,7 @@ class InferenceEngine:
             max_batch_size=max_batch_size,
             max_seq_len=max_seq_len,
             max_prompt_len=max_prompt_len,
-            cache=cache,
+            page_size=page_size,
         )
 
         self.scheduler.start()

astrai/model/components/attention.py

@@ -6,7 +6,7 @@ import torch.nn.functional as F
 from torch import Tensor
 
 from astrai.factory import BaseFactory
-from astrai.inference.core.cache import CacheView
+from astrai.inference.core.cache import KvcacheView
 from astrai.model.components.linear import Linear
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import apply_rotary_emb
@@ -75,7 +75,7 @@ class GQA(nn.Module):
         x: Tensor,
         rotary_emb: Tensor,
         attn_mask: Tensor = None,
-        paged_cache: Optional[CacheView] = None,
+        paged_cache: Optional[KvcacheView] = None,
     ) -> Tensor:
         is_causal = attn_mask is None
 
@@ -162,7 +162,7 @@ class MLA(nn.Module):
         x: Tensor,
         rotary_emb: Tensor,
         attn_mask: Tensor = None,
-        paged_cache: Optional[CacheView] = None,
+        paged_cache: Optional[KvcacheView] = None,
     ) -> Tensor:
         bsz, seq_len, _ = x.size()
         is_causal = attn_mask is None

astrai/model/components/decoder_block.py

@@ -4,7 +4,7 @@ from typing import Optional
 import torch.nn as nn
 from torch import Tensor
 
-from astrai.inference.core.cache import CacheView
+from astrai.inference.core.cache import KvcacheView
 from astrai.model.components.attention import AttnFactory
 from astrai.model.components.mlp import FFNFactory
 from astrai.model.components.norm import RMSNorm
@@ -25,7 +25,7 @@ class DecoderBlock(nn.Module):
         x: Tensor,
         rotary_emb: Tensor,
         attention_mask: Optional[Tensor] = None,
-        paged_cache: Optional[CacheView] = None,
+        paged_cache: Optional[KvcacheView] = None,
     ) -> Tensor:
         attn_output = self.attention(
             self.input_norm(x),

astrai/model/transformer.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 from torch import Tensor
 
 from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.inference.core.cache import CacheView
+from astrai.inference.core.cache import KvcacheView
 from astrai.model.automodel import AutoModel
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
@@ -112,7 +112,7 @@ class AutoRegressiveLM(AutoModel):
         self,
         input_ids: Tensor,
         input_mask: Optional[Tensor] = None,
-        paged_cache: Optional[CacheView] = None,
+        paged_cache: Optional[KvcacheView] = None,
         position_ids: Optional[Tensor] = None,
     ) -> Dict[str, Tensor]:
         assert input_ids.ndim == 2

scripts/eval/evaluate_humaneval.py

@@ -1,22 +1,22 @@
-"""HumanEval benchmark — functional pipeline design.
+"""HumanEval code generation benchmark.
 
-Pipeline:
+Generates n completions per problem, extracts function bodies, executes
-    load -> generate -> extract -> test -> score -> report
+against hidden tests, and computes pass@k.
 
-Each stage is a pure function (except GPU/CPU-bound I/O stages).
+Usage::
-Config is a single dataclass; side effects are isolated at pipeline boundaries.
+
+    python scripts/tools/evaluate_humaneval.py --param_path ./params \
+        --data_path HumanEval.jsonl.gz --output results.json \
+        --num_samples 200 --temperature 0.8 --max_tokens 512
 """
 
 import argparse
-import itertools
 import json
 import os
 import re
-import subprocess
-import sys
-from dataclasses import dataclass
 from math import prod
-from typing import Dict, Iterator, List, Optional, Sequence, Tuple
+from multiprocessing import Process, Queue
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -26,15 +26,11 @@ from astrai.inference import InferenceEngine
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 
-# ---------------------------------------------------------------------------
-# Config
-# ---------------------------------------------------------------------------
-
 HUMANEVAL_URL = (
     "https://github.com/openai/human-eval/raw/master/data/HumanEval.jsonl.gz"
 )
 
-STOP_SEQUENCES = [
+_STOP_SEQUENCES = [
     "\nclass ",
     "\ndef ",
     "\n# ",
@@ -44,85 +40,43 @@ STOP_SEQUENCES = [
 ]
 
 
-@dataclass
+def _download_humaneval(data_path: str):
-class EvalConfig:
+    if os.path.exists(data_path):
-    param_path: str = "./params"
-    data_path: str = "./humaneval/HumanEval.jsonl"
-    output: Optional[str] = None
-
-    test_only: Optional[str] = None
-    generate_only: bool = False
-
-    num_samples: int = 200
-    max_tokens: int = 512
-    temperature: float = 0.8
-    top_p: float = 0.95
-    top_k: int = 50
-    batch_size: int = 32
-    test_timeout: float = 3.0
-    test_workers: int = 8
-    k_values: Tuple[int, ...] = (1, 10, 100)
-    problem_indices: Optional[List[int]] = None
-
-
-def download(url: str, path: str):
-    if os.path.exists(path):
         return
     import gzip
     import urllib.request
 
-    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    os.makedirs(os.path.dirname(data_path) or ".", exist_ok=True)
-    print(f"Downloading {url} ...")
+    print(f"Downloading HumanEval from {HUMANEVAL_URL} ...")
-    tmp = path + ".tmp"
+    tmp = data_path + ".tmp"
-    urllib.request.urlretrieve(url, tmp)
+    urllib.request.urlretrieve(HUMANEVAL_URL, tmp)
     with gzip.open(tmp, "rb") as f_in:
-        with open(path, "wb") as f_out:
+        with open(data_path, "wb") as f_out:
             f_out.write(f_in.read())
     os.remove(tmp)
-    print(f"  saved to {path}")
+    print(f"  saved to {data_path}")
 
 
-def load_jsonl(path: str) -> List[dict]:
+def _load_problems(data_path: str) -> List[dict]:
-    rows = []
+    problems = []
-    with open(path, encoding="utf-8") as f:
+    with open(data_path, "r", encoding="utf-8") as f:
         for line in f:
             line = line.strip()
             if line:
-                rows.append(json.loads(line))
+                problems.append(json.loads(line))
-    return rows
+    return problems
 
 
-def save_json(path: str, data):
+def _extract_function_body(code: str, entry_point: str) -> Optional[str]:
-    with open(path, "w", encoding="utf-8") as f:
+    """Extract the function body from a completion."""
-        json.dump(data, f, indent=2, ensure_ascii=False)
-
-
-def create_engine(param_path: str, batch_size: int) -> InferenceEngine:
-    model = AutoModel.from_pretrained(param_path)
-    tokenizer = AutoTokenizer.from_pretrained(param_path)
-    model.to(device="cuda", dtype=torch.bfloat16)
-    return InferenceEngine(
-        model=model,
-        tokenizer=tokenizer,
-        max_batch_size=batch_size,
-    )
-
-
-def trim_stop(text: str) -> str:
-    for stop in STOP_SEQUENCES:
-        idx = text.find(stop)
-        if idx != -1:
-            text = text[:idx]
-    return text
-
-
-def extract_body(code: str, entry_point: str) -> Optional[str]:
     pattern = rf"def\s+{re.escape(entry_point)}\b[^:]*:"
     match = re.search(pattern, code)
     if not match:
+        # Use the full code as-is if we can't find the function
         return code
 
-    lines = code[match.end() :].split("\n")
+    body_start = match.end()
+    lines = code[body_start:].split("\n")
     body_lines = []
     started = False
 
@@ -140,253 +94,240 @@ def extract_body(code: str, entry_point: str) -> Optional[str]:
         body_lines.append(stripped)
 
     body = "\n".join(body_lines)
-    return body if body.strip() else None
+    if not body.strip():
+        return None
+    return body
 
 
-def deduplicate(seq: Sequence[str]) -> List[str]:
+def _trim_stop_sequences(text: str) -> str:
+    for stop in _STOP_SEQUENCES:
+        idx = text.find(stop)
+        if idx != -1:
+            text = text[:idx]
+    return text
+
+
+def _execute_code(problem: dict, completion: str, timeout: float = 3.0) -> bool:
+    """Run the completion against hidden tests in a subprocess."""
+
+    def _worker(queue, full_code):
+        try:
+            namespace = {}
+            exec(full_code, namespace)
+            check = namespace.get("check")
+            if check is None:
+                queue.put(False)
+                return
+            check(namespace.get(problem["entry_point"]))
+            queue.put(True)
+        except Exception:
+            queue.put(False)
+
+    full_code = problem["prompt"] + completion + "\n" + problem["test"]
+
+    queue: Queue = Queue()
+    proc = Process(target=_worker, args=(queue, full_code))
+    proc.start()
+    proc.join(timeout)
+
+    if proc.is_alive():
+        proc.terminate()
+        proc.join()
+        return False
+
+    try:
+        return queue.get_nowait()
+    except Exception:
+        return False
+
+
+def _pass_at_k(n: int, c: int, k: int) -> float:
+    """Unbiased estimator of pass@k."""
+    if n - c < k:
+        return 1.0
+    return 1.0 - float(prod(1.0 - k / np.arange(n - c + 1, n + 1)))
+
+
+def _deduplicate(completions: List[str]) -> List[str]:
     seen = set()
-    return [x for x in seq if not (x in seen or seen.add(x))]
+    unique = []
+    for c in completions:
+        if c not in seen:
+            seen.add(c)
+            unique.append(c)
+    return unique
 
 
-def generate_batch(
+def _generate(
     engine: InferenceEngine,
     prompt: str,
-    n: int,
+    num_samples: int,
-    batch_size: int,
     max_tokens: int,
     temperature: float,
     top_p: float,
     top_k: int,
+    batch_size: int,
 ) -> List[str]:
+    batches = [prompt] * min(batch_size, num_samples)
     completions = []
-    remaining = n
+    remaining = num_samples
+
     while remaining > 0:
         current = min(batch_size, remaining)
+        batch_prompts = batches[:current]
         outputs = engine.generate(
-            prompt=[prompt] * current,
+            prompt=batch_prompts,
             stream=False,
             max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
         )
-        completions.extend(outputs if isinstance(outputs, list) else [outputs])
+        if isinstance(outputs, str):
+            outputs = [outputs]
+        completions.extend(outputs)
         remaining -= current
-    return deduplicate(completions)
+
+    return _deduplicate(completions)
 
 
-def extract_completions(
+def evaluate(
-    raw: Sequence[str],
-    entry_point: str,
-) -> List[str]:
-    bodies = []
-    for r in raw:
-        t = trim_stop(r)
-        body = extract_body(t, entry_point)
-        if body:
-            bodies.append(body)
-    return bodies
-
-
-def generate_all(
     engine: InferenceEngine,
-    problems: Sequence[dict],
+    problems: List[dict],
-    cfg: EvalConfig,
+    num_samples: int,
-) -> List[dict]:
+    max_tokens: int,
-    results = []
+    temperature: float,
-    for problem in tqdm.tqdm(problems, desc="Generating", unit="problem"):
+    top_p: float,
-        raw = generate_batch(
+    top_k: int,
-            engine,
+    batch_size: int,
-            problem["prompt"],
+    k_values: Tuple[int, ...] = (1, 10, 100),
-            cfg.num_samples,
-            cfg.batch_size,
-            cfg.max_tokens,
-            cfg.temperature,
-            cfg.top_p,
-            cfg.top_k,
-        )
-        bodies = extract_completions(raw, problem["entry_point"])
-        results.append(
-            dict(
-                task_id=problem["task_id"],
-                entry_point=problem["entry_point"],
-                prompt=problem["prompt"],
-                test=problem["test"],
-                completions=bodies,
-            )
-        )
-    return results
-
-
-def execute_one(args: tuple) -> bool:
-    full_code, entry_point, timeout = args
-    try:
-        r = subprocess.run(
-            [sys.executable, "-c", full_code],
-            capture_output=True,
-            timeout=timeout,
-        )
-        return r.returncode == 0
-    except subprocess.TimeoutExpired:
-        return False
-    except Exception:
-        return False
-
-
-def test_one(item: dict, cfg: EvalConfig) -> Tuple[str, int, int]:
-    from concurrent.futures import ProcessPoolExecutor
-
-    task_id = item["task_id"]
-    completions = item["completions"]
-    codes = [
-        (
-            item["prompt"] + c + "\n" + item["test"],
-            item["entry_point"],
-            cfg.test_timeout,
-        )
-        for c in completions
-    ]
-    n = len(codes)
-    passed = 0
-    with ProcessPoolExecutor(max_workers=cfg.test_workers) as pool:
-        for ok in pool.map(execute_one, codes):
-            if ok:
-                passed += 1
-    return task_id, n, passed
-
-
-def test_all(
-    items: Sequence[dict],
-    cfg: EvalConfig,
-) -> Iterator[Tuple[str, int, int]]:
-    for item in tqdm.tqdm(items, desc="Testing", unit="problem"):
-        yield test_one(item, cfg)
-
-
-def pass_at_k(n: int, c: int, k: int) -> float:
-    if n - c < k:
-        return 1.0
-    return 1.0 - float(prod(1.0 - k / np.arange(n - c + 1, n + 1)))
-
-
-def score_results(
-    results: Iterator[Tuple[str, int, int]],
-    k_values: Tuple[int, ...],
 ) -> Dict:
-    # filter to k <= n (peek first result to get n)
+    results = {}
-    first = next(results)
+    all_pass_at_k = {k: [] for k in k_values}
-    results = itertools.chain([first], results)
-    n = first[1]
-    k_values = tuple(k for k in k_values if k <= n)
 
-    scores = {k: [] for k in k_values}
+    for problem in tqdm.tqdm(problems, desc="HumanEval", unit="problem"):
-    output = {}
+        task_id = problem["task_id"]
-    for task_id, n, passed in results:
+        prompt = problem["prompt"]
-        entry = {"task_id": task_id, "n": n, "passed": passed}
+        entry_point = problem["entry_point"]
+
+        raw_completions = _generate(
+            engine,
+            prompt,
+            num_samples,
+            max_tokens,
+            temperature,
+            top_p,
+            top_k,
+            batch_size,
+        )
+
+        completions = []
+        for raw in raw_completions:
+            trimmed = _trim_stop_sequences(raw)
+            body = _extract_function_body(trimmed, entry_point)
+            if body:
+                completions.append(body)
+
+        passed = 0
+        for comp in completions:
+            if _execute_code(problem, comp):
+                passed += 1
+
+        n = len(completions)
+        c = passed
+        result = {"task_id": task_id, "n": n, "passed": c}
         for k in k_values:
-            pk = round(pass_at_k(n, passed, k), 4)
+            result[f"pass@{k}"] = round(_pass_at_k(n, c, k), 4)
-            entry[f"pass@{k}"] = pk
+            all_pass_at_k[k].append(_pass_at_k(n, c, k))
-            scores[k].append(pk)
+        results[task_id] = result
-        output[task_id] = entry
 
     summary = {}
     for k in k_values:
-        vals = scores[k]
+        vals = all_pass_at_k[k]
         summary[f"pass@{k}"] = round(float(np.mean(vals)), 4)
-    output["_summary"] = summary
+    results["_summary"] = summary
-    return output
+
+    return results
 
 
-def run_pipeline(cfg: EvalConfig) -> Dict:
+def main():
-    if cfg.test_only:
+    parser = argparse.ArgumentParser(description="HumanEval benchmark")
-        with open(cfg.test_only, encoding="utf-8") as f:
+    parser.add_argument(
-            generated = json.load(f)
+        "--param_path", type=str, default="./params", help="Model directory"
-    else:
+    )
-        download(HUMANEVAL_URL, cfg.data_path)
+    parser.add_argument(
-
+        "--data_path",
-        problems = load_jsonl(cfg.data_path)
-        if cfg.problem_indices:
-            problems = [problems[i] for i in cfg.problem_indices if i < len(problems)]
-
-        engine = create_engine(cfg.param_path, cfg.batch_size)
-
-        try:
-            generated = generate_all(engine, problems, cfg)
-        finally:
-            engine.shutdown()
-
-        if cfg.output:
-            mid = cfg.output.replace(".json", "_completions.json")
-            save_json(mid, generated)
-            print(f"Completions saved to {mid}")
-
-        if cfg.generate_only:
-            return {}
-
-    results = test_all(generated, cfg)
-    scored = score_results(results, cfg.k_values)
-    return scored
-
-
-def parse_args(argv: Optional[List[str]] = None) -> EvalConfig:
-    p = argparse.ArgumentParser(description="HumanEval benchmark")
-    p.add_argument("--param_path", type=str, default="./params")
-    p.add_argument("--data_path", type=str, default="./humaneval/HumanEval.jsonl")
-    p.add_argument("--output", type=str, default=None)
-    p.add_argument(
-        "--test_only",
         type=str,
+        default="./humaneval/HumanEval.jsonl",
+        help="HumanEval JSONL file (auto-download if missing)",
+    )
+    parser.add_argument("--output", type=str, default=None, help="Output JSON path")
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=200,
+        help="Completions per problem",
+    )
+    parser.add_argument(
+        "--max_tokens", type=int, default=512, help="Max generation tokens"
+    )
+    parser.add_argument(
+        "--temperature", type=float, default=0.8, help="Sampling temperature"
+    )
+    parser.add_argument("--top_p", type=float, default=0.95, help="Top-p sampling")
+    parser.add_argument("--top_k", type=int, default=50, help="Top-k sampling")
+    parser.add_argument(
+        "--batch_size", type=int, default=1, help="Inference batch size"
+    )
+    parser.add_argument(
+        "--problems",
+        type=int,
+        nargs="+",
         default=None,
-        help="Skip generation, test existing completions JSON",
+        help="Specific problem indices (0-based)",
     )
-    p.add_argument(
+    args = parser.parse_args()
-        "--generate_only", action="store_true", help="Only generate, skip testing"
-    )
-    p.add_argument("--num_samples", type=int, default=200)
-    p.add_argument("--max_tokens", type=int, default=512)
-    p.add_argument("--temperature", type=float, default=0.8)
-    p.add_argument("--top_p", type=float, default=0.95)
-    p.add_argument("--top_k", type=int, default=50)
-    p.add_argument("--batch_size", type=int, default=32)
-    p.add_argument("--test_workers", type=int, default=8)
-    p.add_argument("--test_timeout", type=float, default=3.0)
-    p.add_argument("--problems", type=int, nargs="+", default=None)
-    args = p.parse_args(argv)
 
-    return EvalConfig(
+    _download_humaneval(args.data_path)
-        param_path=args.param_path,
+    problems = _load_problems(args.data_path)
-        data_path=args.data_path,
+    if args.problems:
-        output=args.output,
+        problems = [problems[i] for i in args.problems if i < len(problems)]
-        test_only=args.test_only,
+
-        generate_only=args.generate_only,
+    model = AutoModel.from_pretrained(args.param_path)
+    tokenizer = AutoTokenizer.from_pretrained(args.param_path)
+    model.to(device="cuda", dtype=torch.bfloat16)
+
+    engine = InferenceEngine(
+        model=model,
+        tokenizer=tokenizer,
+        max_batch_size=args.batch_size,
+    )
+
+    results = evaluate(
+        engine=engine,
+        problems=problems,
         num_samples=args.num_samples,
         max_tokens=args.max_tokens,
         temperature=args.temperature,
         top_p=args.top_p,
         top_k=args.top_k,
         batch_size=args.batch_size,
-        test_workers=args.test_workers,
+        k_values=(1, 10, 100),
-        test_timeout=args.test_timeout,
-        problem_indices=args.problems,
     )
 
-
+    summary = results.pop("_summary")
-def report(scored: Dict):
-    summary = scored.pop("_summary", {})
     print(f"\n{'=' * 60}")
     for k, v in summary.items():
         print(f"  {k}: {v:.2%}")
     print(f"{'=' * 60}")
-    scored["_summary"] = summary
 
+    if args.output:
+        results["_summary"] = summary
+        with open(args.output, "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2, ensure_ascii=False)
+        print(f"Results saved to {args.output}")
 
-def main():
+    engine.shutdown()
-    cfg = parse_args()
-    scored = run_pipeline(cfg)
-    report(scored)
-    if cfg.output:
-        save_json(cfg.output, scored)
-        print(f"Results saved to {cfg.output}")
 
 
 if __name__ == "__main__":

scripts/eval/evaluate_ifd.py

@@ -4,15 +4,6 @@ IFD = conditional_NLL / unconditional_NLL
 
 - Messages format: plain text concatenation (no chat template)
 - Plain format: raw instr_key + resp_key fields
-
-v2 changelog:
-  - Same token set: unconditional pass prefixes resp with a plain-text sentinel
-    (default ``\\n``; use ``--sentinel_text ""`` for bos/pad fallback).
-    Both branches predict the identical N resp tokens.
-    Single-token answers (rl=1) are now supported.
-  - ctx_len tracked in output
-  - skip_reason for None samples (no more silent None)
-  - --per_token for per-token IFD breakdown
 """
 
 import argparse
@@ -30,7 +21,7 @@ from astrai.tokenize import AutoTokenizer
 def _pack_bins(pairs, max_len):
     """BFD bin packing: pack (c+r) into bins of max total length."""
     indexed = sorted(enumerate(pairs), key=lambda x: -(len(x[1][0]) + len(x[1][1])))
-    bins = []
+    bins = []  # each bin: list of (orig_idx, ctx_ids, resp_ids)
     lengths = []
     for orig_idx, (c, r) in indexed:
         size = len(c) + len(r)
@@ -48,51 +39,19 @@ def _pack_bins(pairs, max_len):
     return bins
 
 
-def _resolve_sentinel_ids(tokenizer, sentinel_text):
-    """Tokenize the sentinel text for the unconditional pass prefix.
-
-    Falls back to bos/pad_token_id when sentinel_text is empty or
-    cannot be encoded.
-    """
-    if sentinel_text:
-        ids = tokenizer.encode(sentinel_text, add_special_tokens=False)
-        if ids:
-            return ids
-    for attr in ("bos_token_id", "pad_token_id", "eos_token_id"):
-        tid = getattr(tokenizer, attr, None)
-        if tid is not None:
-            return [tid]
-    return [0]
-
-
 @torch.inference_mode()
-def _score_batch(
+def _score_batch(pairs, model, device, max_len=2048):
-    pairs, model, device, max_len=2048, sentinel_ids=None, per_token=False
+    """BFD-packed IFD: pack items into bins, one forward pass per bin."""
-):
-    """BFD-packed IFD with text-sentinel-anchored unconditional pass.
-
-    Conditional:   (ctx + resp[0..i-1]) → resp[i],  i = 0..N-1
-    Unconditional: (<sentinel> + resp[0..i-1]) → resp[i],  i = 0..N-1
-
-    Both branches predict the identical N response tokens.  A short
-    plain-text sentinel gives the unconditional pass a prefix so that
-    every response token can be predicted.  Single-token answers (rl=1)
-    are supported.
-    """
     if not pairs:
         return []
-
-    if sentinel_ids is None:
-        sentinel_ids = [0]
-
     bins = _pack_bins(pairs, max_len)
+
     result = [None] * len(pairs)
 
-    # ---- conditional pass (packed, per-document position IDs) ----
     for bin_items in bins:
         seq_ids = []
-        global_pos = []
+        global_pos = []  # doc-reset position IDs for RoPE
-        doc_ids = []
+        doc_ids = []  # document index for attention mask
         doc_offsets = []
 
         for di, (orig_idx, c, r) in enumerate(bin_items):
@@ -108,10 +67,8 @@ def _score_batch(
 
         full_ids = torch.tensor([seq_ids], device=device, dtype=torch.long)
         pos_ids = torch.tensor([global_pos], device=device, dtype=torch.long)
-        seq_len = len(seq_ids)
+        T = len(seq_ids)
-        causal = torch.tril(
+        causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=device))
-            torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
-        )
         doc_t = torch.tensor([doc_ids], device=device)
         doc_mask = doc_t.unsqueeze(-1) == doc_t.unsqueeze(-2)
         attn_mask = (causal & doc_mask[0]).unsqueeze(0).unsqueeze(0)
@@ -121,72 +78,46 @@ def _score_batch(
 
         for start, end, orig_idx, ctx_len in doc_offsets:
             rl = end - start - ctx_len
+            if rl < 2:
+                continue
             resp_start = start + ctx_len - 1
             resp_logits = logits_full[resp_start : end - 1]
             resp_targets = torch.tensor(
                 seq_ids[start + ctx_len : end], device=device, dtype=torch.long
             )
-            cond_losses = F.cross_entropy(
+            L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()
-                resp_logits, resp_targets, reduction="none"
+            result[orig_idx] = (L_cond, rl)
-            ).cpu()
-            result[orig_idx] = {
-                "_cond_losses": cond_losses,
-                "_rl": rl,
-                "_ctx_len": ctx_len,
-            }
 
-    # ---- unconditional pass (sentinel-prefixed, batched 2D) ----
+    # unconditional pass: batch all responses separately (sorted by length)
-    valid_items = [
+    resp_seqs = [
-        (
+        (i, result[i][1], pairs[i][1])
-            i,
-            result[i]["_rl"],
-            result[i]["_ctx_len"],
-            result[i]["_cond_losses"],
-            pairs[i][1],
-        )
         for i in range(len(pairs))
-        if result[i] is not None and "_cond_losses" in result[i]
+        if result[i] is not None
     ]
-    if not valid_items:
+    if resp_seqs:
-        return result
+        resp_seqs.sort(key=lambda x: -x[1])
-
+        r_batch = torch.zeros(
-    valid_items.sort(key=lambda x: -x[1])
+            len(resp_seqs),
-    prefix_len = len(sentinel_ids)
+            max(len(r) for _, _, r in resp_seqs),
-    max_rl = prefix_len + max(rl for _, rl, _, _, _ in valid_items)
+            dtype=torch.long,
-    bsz = len(valid_items)
+            device=device,
-
-    u_batch = torch.zeros(bsz, max_rl, dtype=torch.long, device=device)
-    for ri, (_, rl, _, _, r_ids) in enumerate(valid_items):
-        u_batch[ri, :prefix_len] = torch.tensor(sentinel_ids, dtype=torch.long)
-        u_batch[ri, prefix_len : prefix_len + rl] = torch.tensor(
-            r_ids, dtype=torch.long
         )
+        for ri, (_, rl, r_ids) in enumerate(resp_seqs):
+            r_batch[ri, :rl] = torch.tensor(r_ids, dtype=torch.long)
+        logits_resp = model(r_batch)["logits"]
 
-    logits_resp = model(u_batch)["logits"]
+        for ri, (orig_idx, rl, _) in enumerate(resp_seqs):
-
+            L_cond = result[orig_idx][0]
-    for ri, (orig_idx, rl, ctx_len, cond_losses, _) in enumerate(valid_items):
+            unp_logits = logits_resp[ri, : rl - 1]
-        unp_logits = logits_resp[ri, prefix_len - 1 : prefix_len - 1 + rl]
+            unp_targets = r_batch[ri, 1:rl]
-        unp_targets = u_batch[ri, prefix_len : prefix_len + rl]
+            L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
-        uncond_losses = F.cross_entropy(unp_logits, unp_targets, reduction="none").cpu()
-
-        L_cond = cond_losses.mean().item()
-        L_uncond = uncond_losses.mean().item()
             ifd = L_cond / L_uncond if L_uncond > 0 else None
-
+            result[orig_idx] = {
-        out = {
                 "L_cond": round(L_cond, 6),
                 "L_uncond": round(L_uncond, 6),
                 "ifd": round(ifd, 6) if ifd is not None else None,
-            "ctx_len": ctx_len,
                 "resp_len": rl,
             }
-        if per_token:
-            per = [
-                (round(c.item() / u.item(), 6) if u.item() > 0 else None)
-                for c, u in zip(cond_losses, uncond_losses)
-            ]
-            out["ifd_per_token"] = per
-        result[orig_idx] = out
 
     return result
 
@@ -204,40 +135,17 @@ def _trim(context_ids, resp_ids, max_len):
     return context_ids[overflow:], resp_ids
 
 
-def score_plain(
+def score_plain(model, tokenizer, instruction, response, device, max_len=2048):
-    model,
-    tokenizer,
-    instruction,
-    response,
-    device,
-    max_len=2048,
-    sentinel_ids=None,
-    per_token=False,
-):
     """Compute IFD for a single instruction-response pair (plain format)."""
     ctx_ids = tokenizer.encode(instruction, add_special_tokens=False)
     resp_ids = tokenizer.encode(response, add_special_tokens=False)
     ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
     if not ctx_ids or not resp_ids:
-        return {
+        return {"L_cond": None, "L_uncond": None, "ifd": None, "error": "empty"}
-            "L_cond": None,
+    return _score_batch([(ctx_ids, resp_ids)], model, device, max_len)[0]
-            "L_uncond": None,
-            "ifd": None,
-            "skip_reason": "empty ctx or resp",
-        }
-    return _score_batch(
-        [(ctx_ids, resp_ids)],
-        model,
-        device,
-        max_len,
-        sentinel_ids=sentinel_ids,
-        per_token=per_token,
-    )[0]
 
 
-def score_messages(
+def score_messages(model, tokenizer, messages, device, max_len=2048):
-    model, tokenizer, messages, device, max_len=2048, sentinel_ids=None, per_token=False
-):
     """Compute IFD for each assistant turn in a messages array."""
     turns = []
     for i, msg in enumerate(messages):
@@ -251,10 +159,8 @@ def score_messages(
             turns.append((ctx_ids, resp_ids))
     if not turns:
         return None
-    raw_scores = _score_batch(
+    raw_scores = _score_batch(turns, model, device, max_len)
-        turns, model, device, max_len, sentinel_ids=sentinel_ids, per_token=per_token
+    valid = [s for s in raw_scores if s is not None and s["ifd"] is not None]
-    )
-    valid = [s for s in raw_scores if s is not None and s.get("ifd") is not None]
     if not valid:
         return {"ifd": None, "ifd_turns": raw_scores}
     avg = sum(s["ifd"] for s in valid) / len(valid)
@@ -275,8 +181,6 @@ def process_file(
     data_format="plain",
     batch_size=1,
     device=None,
-    sentinel_text="\n",
-    per_token=False,
 ):
     if device is None:
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -287,8 +191,6 @@ def process_file(
     model.to(device=device, dtype=dtype)
     model.eval()
 
-    sentinel_ids = _resolve_sentinel_ids(tokenizer, sentinel_text)
-
     with open(input_file, encoding="utf-8") as f:
         data = [json.loads(line) for line in f if line.strip()]
 
@@ -309,14 +211,7 @@ def process_file(
                 if ctx_ids and resp_ids:
                     turns.append((ctx_ids, resp_ids))
             if not turns:
-                results.append(
+                results.append({**item, "ifd": None, "ifd_turns": []})
-                    {
-                        **item,
-                        "ifd": None,
-                        "skip_reason": "no valid assistant turns",
-                        "ifd_turns": [],
-                    }
-                )
                 continue
             buffer.append((item, turns, "messages"))
         else:
@@ -324,32 +219,15 @@ def process_file(
             resp_ids = tokenizer.encode(item[resp_key], add_special_tokens=False)
             ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
             if not ctx_ids or not resp_ids:
-                results.append(
+                results.append({**item, "ifd": None, "ifd_detail": {"error": "empty"}})
-                    {
-                        **item,
-                        "ifd": None,
-                        "ifd_detail": {"skip_reason": "empty ctx or resp"},
-                    }
-                )
                 continue
             buffer.append((item, [(ctx_ids, resp_ids)], "plain"))
 
         if len(buffer) >= batch_size:
-            _flush_buffer(
+            _flush_buffer(buffer, results, all_ifds, model, device, max_len)
-                buffer,
-                results,
-                all_ifds,
-                model,
-                device,
-                max_len,
-                sentinel_ids,
-                per_token,
-            )
 
     if buffer:
-        _flush_buffer(
+        _flush_buffer(buffer, results, all_ifds, model, device, max_len)
-            buffer, results, all_ifds, model, device, max_len, sentinel_ids, per_token
-        )
 
     with open(output_file, "w", encoding="utf-8") as f:
         for item in results:
@@ -360,20 +238,17 @@ def process_file(
         print(f"\n{'=' * 50}")
         print(f"  Samples:    {len(data)}")
         print(f"  Valid IFD:  {len(valid_ifd)}")
-        print(f"  Skipped:       {len(data) - len(valid_ifd)}")
         print(f"  Mean IFD:   {statistics.mean(valid_ifd):.4f}")
         print(f"  Median IFD: {statistics.median(valid_ifd):.4f}")
-        if len(valid_ifd) > 1:
         print(f"  Stdev IFD:  {statistics.stdev(valid_ifd):.4f}")
         print(f"  Min IFD:    {min(valid_ifd):.4f}")
         print(f"  Max IFD:    {max(valid_ifd):.4f}")
         print(f"{'=' * 50}")
+
     print(f"Results saved to {output_file}")
 
 
-def _flush_buffer(
+def _flush_buffer(buffer, results, all_ifds, model, device, max_len=2048):
-    buffer, results, all_ifds, model, device, max_len, sentinel_ids, per_token
-):
     all_pairs = []
     indices = []
     for item, turns, fmt in buffer:
@@ -381,21 +256,12 @@ def _flush_buffer(
         all_pairs.extend(turns)
         indices.append((item, turns, fmt, start, len(all_pairs)))
 
-    raw = _score_batch(
+    raw = _score_batch(all_pairs, model, device, max_len)
-        all_pairs,
-        model,
-        device,
-        max_len,
-        sentinel_ids=sentinel_ids,
-        per_token=per_token,
-    )
 
     for item, turns, fmt, start, end in indices:
         turn_scores = raw[start:end]
         if fmt == "messages":
-            valid = [
+            valid = [s for s in turn_scores if s is not None and s["ifd"] is not None]
-                s for s in turn_scores if s is not None and s.get("ifd") is not None
-            ]
             if not valid:
                 results.append({**item, "ifd": None, "ifd_turns": turn_scores})
             else:
@@ -411,8 +277,8 @@ def _flush_buffer(
                 )
         else:
             score = turn_scores[0]
-            all_ifds.append(score.get("ifd"))
+            all_ifds.append(score["ifd"])
-            results.append({**item, "ifd": score.get("ifd"), "ifd_detail": score})
+            results.append({**item, "ifd": score["ifd"], "ifd_detail": score})
 
     buffer.clear()
 
@@ -442,17 +308,6 @@ def main():
         "--batch_size", type=int, default=8, help="Batch size for model forward passes"
     )
     parser.add_argument("--device", type=str, default=None, help="Device (e.g. cuda:0)")
-    parser.add_argument(
-        "--sentinel_text",
-        type=str,
-        default="\n",
-        help='Plain-text prefix for unconditional pass (default: "\\n"). Use "" for bos/pad fallback.',
-    )
-    parser.add_argument(
-        "--per_token",
-        action="store_true",
-        help="Include per-token IFD breakdown in output",
-    )
     args = parser.parse_args()
 
     process_file(
@@ -465,8 +320,6 @@ def main():
         data_format=args.format,
         batch_size=args.batch_size,
         device=args.device,
-        sentinel_text=args.sentinel_text,
-        per_token=args.per_token,
     )
 
 

scripts/eval/evaluate_rouge.py

@@ -1,153 +0,0 @@
-"""ROUGE evaluation (manual implementation, no external deps).
-
-Computes ROUGE-1, ROUGE-2, ROUGE-L precision, recall, and F1.
-
-Usage::
-
-    # Batch evaluation from JSONL (each line: {"reference": ..., "candidate": ...})
-    python scripts/eval/evaluate_rouge.py --data_path preds.jsonl --output results.json
-
-    # As a library
-    from scripts.eval.evaluate_rouge import compute_rouge
-    scores = compute_rouge("the cat sat on the mat", "the cat sat")
-"""
-
-import argparse
-import json
-from collections import Counter
-from typing import Dict, List, Tuple
-
-
-def _tokenize(text: str) -> List[str]:
-    return text.split()
-
-
-def _ngrams(tokens: List[str], n: int) -> Counter:
-    return Counter(zip(*[tokens[i:] for i in range(n)]))
-
-
-def _lcs(x: List[str], y: List[str]) -> int:
-    m, n = len(x), len(y)
-    dp = [[0] * (n + 1) for _ in range(m + 1)]
-    for i in range(1, m + 1):
-        xi = x[i - 1]
-        dpi = dp[i]
-        dpi_1 = dp[i - 1]
-        for j in range(1, n + 1):
-            if xi == y[j - 1]:
-                dpi[j] = dpi_1[j - 1] + 1
-            else:
-                dpi[j] = dpi_1[j] if dpi_1[j] > dpi[j - 1] else dpi[j - 1]
-    return dp[m][n]
-
-
-def _f1(precision: float, recall: float) -> float:
-    if precision + recall == 0:
-        return 0.0
-    return 2 * precision * recall / (precision + recall)
-
-
-def _rouge_n(ref_tokens: List[str], cand_tokens: List[str], n: int) -> Dict[str, float]:
-    ref_ngrams = _ngrams(ref_tokens, n)
-    cand_ngrams = _ngrams(cand_tokens, n)
-
-    overlap = sum((cand_ngrams & ref_ngrams).values())
-    cand_total = sum(cand_ngrams.values())
-    ref_total = sum(ref_ngrams.values())
-
-    precision = overlap / cand_total if cand_total > 0 else 0.0
-    recall = overlap / ref_total if ref_total > 0 else 0.0
-    f1 = _f1(precision, recall)
-
-    return {"precision": precision, "recall": recall, "f1": f1}
-
-
-def _rouge_l(ref_tokens: List[str], cand_tokens: List[str]) -> Dict[str, float]:
-    lcs_len = _lcs(ref_tokens, cand_tokens)
-    ref_len = len(ref_tokens)
-    cand_len = len(cand_tokens)
-
-    recall = lcs_len / ref_len if ref_len > 0 else 0.0
-    precision = lcs_len / cand_len if cand_len > 0 else 0.0
-    f1 = _f1(precision, recall)
-
-    return {"precision": precision, "recall": recall, "f1": f1}
-
-
-def compute_rouge(
-    reference: str, candidate: str, n: int = 2
-) -> Dict[str, Dict[str, float]]:
-    """Compute ROUGE-N (1..n) and ROUGE-L scores.
-
-    Returns::
-
-        {
-            "rouge-1": {"precision": ..., "recall": ..., "f1": ...},
-            "rouge-2": {"precision": ..., "recall": ..., "f1": ...},
-            "rouge-l": {"precision": ..., "recall": ..., "f1": ...},
-        }
-    """
-    ref_tokens = _tokenize(reference)
-    cand_tokens = _tokenize(candidate)
-
-    results = {}
-    for i in range(1, n + 1):
-        results[f"rouge-{i}"] = _rouge_n(ref_tokens, cand_tokens, i)
-    results["rouge-l"] = _rouge_l(ref_tokens, cand_tokens)
-    return results
-
-
-def evaluate_file(data_path: str) -> Dict:
-    with open(data_path, "r", encoding="utf-8") as f:
-        pairs = [json.loads(line) for line in f if line.strip()]
-
-    agg = {
-        k: {"precision": 0.0, "recall": 0.0, "f1": 0.0}
-        for k in ("rouge-1", "rouge-2", "rouge-l")
-    }
-    per_item = []
-
-    for item in pairs:
-        ref = item["reference"]
-        cand = item["candidate"]
-        scores = compute_rouge(ref, cand)
-        per_item.append({**item, "scores": scores})
-        for k, v in scores.items():
-            agg[k]["precision"] += v["precision"]
-            agg[k]["recall"] += v["recall"]
-            agg[k]["f1"] += v["f1"]
-
-    n = len(pairs)
-    for k in agg:
-        agg[k] = {m: v / n for m, v in agg[k].items()}
-
-    return {"num_samples": n, "aggregate": agg, "per_item": per_item}
-
-
-def main():
-    parser = argparse.ArgumentParser(description="ROUGE evaluation")
-    parser.add_argument(
-        "--data_path", required=True, help="JSONL with reference/candidate per line"
-    )
-    parser.add_argument("--output", type=str, default=None, help="Output JSON path")
-    args = parser.parse_args()
-
-    results = evaluate_file(args.data_path)
-    agg = results["aggregate"]
-
-    print(f"Samples: {results['num_samples']}")
-    print()
-    for metric in ("rouge-1", "rouge-2", "rouge-l"):
-        s = agg[metric]
-        print(
-            f"  {metric:8s}  P={s['precision']:.4f}  R={s['recall']:.4f}  F1={s['f1']:.4f}"
-        )
-
-    if args.output:
-        with open(args.output, "w", encoding="utf-8") as f:
-            json.dump(results, f, indent=2, ensure_ascii=False)
-        print(f"\nSaved to {args.output}")
-
-
-if __name__ == "__main__":
-    main()

tests/inference/test_cache.py

@@ -4,7 +4,7 @@ import torch
 
 from astrai.inference import (
     Allocator,
-    PageCache,
+    KVCache,
     PagePool,
     PrefixCache,
     Storage,
@@ -161,7 +161,7 @@ def test_task_table_pop():
 
 
 def test_kv_cache_task_extend_allocates():
-    cache = PageCache(
+    cache = KVCache(
         n_layers=1,
         n_pages=8,
         page_size=64,
@@ -177,7 +177,7 @@ def test_kv_cache_task_extend_allocates():
 
 
 def test_kv_cache_task_extend_fails_when_pool_full():
-    cache = PageCache(
+    cache = KVCache(
         n_layers=1,
         n_pages=2,
         page_size=64,