Compare commits
No commits in common. "bbe6ff2d8f7245ef964e733ce95e128f4d3ef43e" and "c7158418dd53fb1cdb556a2714dbdf45b713adef" have entirely different histories.
bbe6ff2d8f
...
c7158418dd
|
|
@ -1,4 +1,4 @@
|
|||
__version__ = "1.3.8"
|
||||
__version__ = "1.3.7"
|
||||
__author__ = "ViperEkura"
|
||||
|
||||
from astrai.config import (
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
PageCacheView,
|
||||
KvcacheView,
|
||||
PagePool,
|
||||
PrefixCache,
|
||||
Storage,
|
||||
|
|
@ -67,12 +63,8 @@ __all__ = [
|
|||
"TaskManager",
|
||||
"TaskStatus",
|
||||
"Allocator",
|
||||
"CacheView",
|
||||
"KVCache",
|
||||
"ContiguousCache",
|
||||
"ContiguousCacheView",
|
||||
"PageCache",
|
||||
"PageCacheView",
|
||||
"KvcacheView",
|
||||
"PagePool",
|
||||
"PrefixCache",
|
||||
"Storage",
|
||||
|
|
|
|||
|
|
@ -2,12 +2,8 @@
|
|||
|
||||
from astrai.inference.core.cache import (
|
||||
Allocator,
|
||||
CacheView,
|
||||
ContiguousCache,
|
||||
ContiguousCacheView,
|
||||
KVCache,
|
||||
PageCache,
|
||||
PageCacheView,
|
||||
KvcacheView,
|
||||
PagePool,
|
||||
PrefixCache,
|
||||
Storage,
|
||||
|
|
@ -20,12 +16,8 @@ from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
|
|||
|
||||
__all__ = [
|
||||
"Allocator",
|
||||
"CacheView",
|
||||
"KVCache",
|
||||
"ContiguousCache",
|
||||
"ContiguousCacheView",
|
||||
"PageCache",
|
||||
"PageCacheView",
|
||||
"KvcacheView",
|
||||
"PagePool",
|
||||
"PrefixCache",
|
||||
"Storage",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import threading
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import OrderedDict
|
||||
from typing import Callable, Dict, List, Optional, Tuple
|
||||
|
||||
|
|
@ -63,8 +62,7 @@ class Allocator:
|
|||
|
||||
def touch(self, idx: int):
|
||||
with self._lock:
|
||||
if idx in self._lru:
|
||||
self._lru.move_to_end(idx)
|
||||
self._lru.move_to_end(idx)
|
||||
|
||||
|
||||
class PrefixCache:
|
||||
|
|
@ -276,42 +274,7 @@ class Storage:
|
|||
return k, v
|
||||
|
||||
|
||||
class CacheView(ABC):
|
||||
"""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):
|
||||
class KvcacheView:
|
||||
"""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):
|
||||
"""Paged KV-cache with prefix sharing."""
|
||||
class KVCache:
|
||||
"""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(
|
||||
self, task_ids: List[str], total_len: int, device: torch.device
|
||||
) -> PageCacheView:
|
||||
page_table = self._table.table_tensor(task_ids, device)
|
||||
return PageCacheView(self._storage, page_table, total_len)
|
||||
def make_table_tensor(self, task_ids: List[str], device: torch.device) -> Tensor:
|
||||
return self._table.table_tensor(task_ids, device)
|
||||
|
||||
|
||||
class ContiguousCacheView(CacheView):
|
||||
"""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)
|
||||
def bind(self, page_table: Tensor, total_len: int = 0) -> KvcacheView:
|
||||
return KvcacheView(self._storage, page_table, total_len)
|
||||
|
|
|
|||
|
|
@ -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, :]
|
||||
|
|
|
|||
|
|
@ -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,20 +41,19 @@ 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._cache = cache
|
||||
else:
|
||||
self._cache = ContiguousCache(
|
||||
config.n_layers,
|
||||
max_batch_size,
|
||||
self.max_seq_len,
|
||||
config.n_kv_heads,
|
||||
head_dim,
|
||||
self.device,
|
||||
self.dtype,
|
||||
)
|
||||
self._page_cache = KVCache(
|
||||
config.n_layers,
|
||||
n_pages,
|
||||
page_size,
|
||||
config.n_kv_heads,
|
||||
config.dim // config.n_heads,
|
||||
self.device,
|
||||
self.dtype,
|
||||
)
|
||||
|
||||
self._task_mgr = TaskManager(
|
||||
tokenizer=tokenizer,
|
||||
|
|
@ -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(
|
||||
cache, "page_size", 64
|
||||
)
|
||||
start_logical_page = start_pos // self._page_cache.page_size
|
||||
for t in group:
|
||||
cache.task_record_hashes(
|
||||
t.task_id, t.prompt_ids, start_logical_page
|
||||
self._page_cache.task_record_hashes(
|
||||
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()):
|
||||
group = sorted(pos_groups[next_pos], key=lambda t: t.task_id)
|
||||
if pos_groups:
|
||||
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(
|
||||
t.task_id,
|
||||
self._task_mgr.tokenizer.decode([ntok]),
|
||||
)
|
||||
pos = t.input_tokens + t.output_tokens
|
||||
extend_ok = self._page_cache.task_extend(t.task_id, pos)
|
||||
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)
|
||||
cache.task_free(task.task_id)
|
||||
if task.stream_callback:
|
||||
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)
|
||||
self._cache.task_free(task.task_id)
|
||||
if task.stream_callback:
|
||||
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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,22 +1,22 @@
|
|||
"""HumanEval benchmark — functional pipeline design.
|
||||
"""HumanEval code generation benchmark.
|
||||
|
||||
Pipeline:
|
||||
load -> generate -> extract -> test -> score -> report
|
||||
Generates n completions per problem, extracts function bodies, executes
|
||||
against hidden tests, and computes pass@k.
|
||||
|
||||
Each stage is a pure function (except GPU/CPU-bound I/O stages).
|
||||
Config is a single dataclass; side effects are isolated at pipeline boundaries.
|
||||
Usage::
|
||||
|
||||
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
|
||||
class EvalConfig:
|
||||
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):
|
||||
def _download_humaneval(data_path: str):
|
||||
if os.path.exists(data_path):
|
||||
return
|
||||
import gzip
|
||||
import urllib.request
|
||||
|
||||
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
||||
print(f"Downloading {url} ...")
|
||||
tmp = path + ".tmp"
|
||||
urllib.request.urlretrieve(url, tmp)
|
||||
os.makedirs(os.path.dirname(data_path) or ".", exist_ok=True)
|
||||
print(f"Downloading HumanEval from {HUMANEVAL_URL} ...")
|
||||
tmp = data_path + ".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]:
|
||||
rows = []
|
||||
with open(path, encoding="utf-8") as f:
|
||||
def _load_problems(data_path: str) -> List[dict]:
|
||||
problems = []
|
||||
with open(data_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line:
|
||||
rows.append(json.loads(line))
|
||||
return rows
|
||||
problems.append(json.loads(line))
|
||||
return problems
|
||||
|
||||
|
||||
def save_json(path: str, data):
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
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]:
|
||||
def _extract_function_body(code: str, entry_point: str) -> Optional[str]:
|
||||
"""Extract the function body from a completion."""
|
||||
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,
|
||||
batch_size: int,
|
||||
num_samples: 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(
|
||||
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(
|
||||
def evaluate(
|
||||
engine: InferenceEngine,
|
||||
problems: Sequence[dict],
|
||||
cfg: EvalConfig,
|
||||
) -> List[dict]:
|
||||
results = []
|
||||
for problem in tqdm.tqdm(problems, desc="Generating", unit="problem"):
|
||||
raw = generate_batch(
|
||||
engine,
|
||||
problem["prompt"],
|
||||
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, ...],
|
||||
problems: List[dict],
|
||||
num_samples: int,
|
||||
max_tokens: int,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
top_k: int,
|
||||
batch_size: int,
|
||||
k_values: Tuple[int, ...] = (1, 10, 100),
|
||||
) -> Dict:
|
||||
# filter to k <= n (peek first result to get n)
|
||||
first = next(results)
|
||||
results = itertools.chain([first], results)
|
||||
n = first[1]
|
||||
k_values = tuple(k for k in k_values if k <= n)
|
||||
results = {}
|
||||
all_pass_at_k = {k: [] for k in k_values}
|
||||
|
||||
scores = {k: [] for k in k_values}
|
||||
output = {}
|
||||
for task_id, n, passed in results:
|
||||
entry = {"task_id": task_id, "n": n, "passed": passed}
|
||||
for problem in tqdm.tqdm(problems, desc="HumanEval", unit="problem"):
|
||||
task_id = problem["task_id"]
|
||||
prompt = problem["prompt"]
|
||||
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)
|
||||
entry[f"pass@{k}"] = pk
|
||||
scores[k].append(pk)
|
||||
output[task_id] = entry
|
||||
result[f"pass@{k}"] = round(_pass_at_k(n, c, k), 4)
|
||||
all_pass_at_k[k].append(_pass_at_k(n, c, k))
|
||||
results[task_id] = result
|
||||
|
||||
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
|
||||
return output
|
||||
results["_summary"] = summary
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def run_pipeline(cfg: EvalConfig) -> Dict:
|
||||
if cfg.test_only:
|
||||
with open(cfg.test_only, encoding="utf-8") as f:
|
||||
generated = json.load(f)
|
||||
else:
|
||||
download(HUMANEVAL_URL, cfg.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",
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="HumanEval benchmark")
|
||||
parser.add_argument(
|
||||
"--param_path", type=str, default="./params", help="Model directory"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data_path",
|
||||
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(
|
||||
"--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)
|
||||
args = parser.parse_args()
|
||||
|
||||
return EvalConfig(
|
||||
param_path=args.param_path,
|
||||
data_path=args.data_path,
|
||||
output=args.output,
|
||||
test_only=args.test_only,
|
||||
generate_only=args.generate_only,
|
||||
_download_humaneval(args.data_path)
|
||||
problems = _load_problems(args.data_path)
|
||||
if args.problems:
|
||||
problems = [problems[i] for i in args.problems if i < len(problems)]
|
||||
|
||||
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,
|
||||
test_timeout=args.test_timeout,
|
||||
problem_indices=args.problems,
|
||||
k_values=(1, 10, 100),
|
||||
)
|
||||
|
||||
|
||||
def report(scored: Dict):
|
||||
summary = scored.pop("_summary", {})
|
||||
summary = results.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():
|
||||
cfg = parse_args()
|
||||
scored = run_pipeline(cfg)
|
||||
report(scored)
|
||||
if cfg.output:
|
||||
save_json(cfg.output, scored)
|
||||
print(f"Results saved to {cfg.output}")
|
||||
engine.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
pairs, model, device, max_len=2048, sentinel_ids=None, per_token=False
|
||||
):
|
||||
"""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.
|
||||
"""
|
||||
def _score_batch(pairs, model, device, max_len=2048):
|
||||
"""BFD-packed IFD: pack items into bins, one forward pass per bin."""
|
||||
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 = []
|
||||
doc_ids = []
|
||||
global_pos = [] # doc-reset position IDs for RoPE
|
||||
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)
|
||||
causal = torch.tril(
|
||||
torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
|
||||
)
|
||||
T = len(seq_ids)
|
||||
causal = torch.tril(torch.ones(T, T, 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(
|
||||
resp_logits, resp_targets, reduction="none"
|
||||
).cpu()
|
||||
result[orig_idx] = {
|
||||
"_cond_losses": cond_losses,
|
||||
"_rl": rl,
|
||||
"_ctx_len": ctx_len,
|
||||
}
|
||||
L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()
|
||||
result[orig_idx] = (L_cond, rl)
|
||||
|
||||
# ---- unconditional pass (sentinel-prefixed, batched 2D) ----
|
||||
valid_items = [
|
||||
(
|
||||
i,
|
||||
result[i]["_rl"],
|
||||
result[i]["_ctx_len"],
|
||||
result[i]["_cond_losses"],
|
||||
pairs[i][1],
|
||||
)
|
||||
# unconditional pass: batch all responses separately (sorted by length)
|
||||
resp_seqs = [
|
||||
(i, result[i][1], 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:
|
||||
return result
|
||||
|
||||
valid_items.sort(key=lambda x: -x[1])
|
||||
prefix_len = len(sentinel_ids)
|
||||
max_rl = prefix_len + max(rl for _, rl, _, _, _ in valid_items)
|
||||
bsz = len(valid_items)
|
||||
|
||||
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
|
||||
if resp_seqs:
|
||||
resp_seqs.sort(key=lambda x: -x[1])
|
||||
r_batch = torch.zeros(
|
||||
len(resp_seqs),
|
||||
max(len(r) for _, _, r in resp_seqs),
|
||||
dtype=torch.long,
|
||||
device=device,
|
||||
)
|
||||
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, ctx_len, cond_losses, _) in enumerate(valid_items):
|
||||
unp_logits = logits_resp[ri, prefix_len - 1 : prefix_len - 1 + rl]
|
||||
unp_targets = u_batch[ri, prefix_len : prefix_len + rl]
|
||||
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
|
||||
|
||||
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
|
||||
for ri, (orig_idx, rl, _) in enumerate(resp_seqs):
|
||||
L_cond = result[orig_idx][0]
|
||||
unp_logits = logits_resp[ri, : rl - 1]
|
||||
unp_targets = r_batch[ri, 1:rl]
|
||||
L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
|
||||
ifd = L_cond / L_uncond if L_uncond > 0 else None
|
||||
result[orig_idx] = {
|
||||
"L_cond": round(L_cond, 6),
|
||||
"L_uncond": round(L_uncond, 6),
|
||||
"ifd": round(ifd, 6) if ifd is not None else None,
|
||||
"resp_len": rl,
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
|
|
@ -204,40 +135,17 @@ def _trim(context_ids, resp_ids, max_len):
|
|||
return context_ids[overflow:], resp_ids
|
||||
|
||||
|
||||
def score_plain(
|
||||
model,
|
||||
tokenizer,
|
||||
instruction,
|
||||
response,
|
||||
device,
|
||||
max_len=2048,
|
||||
sentinel_ids=None,
|
||||
per_token=False,
|
||||
):
|
||||
def score_plain(model, tokenizer, instruction, response, device, max_len=2048):
|
||||
"""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 {
|
||||
"L_cond": None,
|
||||
"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]
|
||||
return {"L_cond": None, "L_uncond": None, "ifd": None, "error": "empty"}
|
||||
return _score_batch([(ctx_ids, resp_ids)], model, device, max_len)[0]
|
||||
|
||||
|
||||
def score_messages(
|
||||
model, tokenizer, messages, device, max_len=2048, sentinel_ids=None, per_token=False
|
||||
):
|
||||
def score_messages(model, tokenizer, messages, device, max_len=2048):
|
||||
"""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(
|
||||
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.get("ifd") is not None]
|
||||
raw_scores = _score_batch(turns, model, device, max_len)
|
||||
valid = [s for s in raw_scores if s is not None and s["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(
|
||||
{
|
||||
**item,
|
||||
"ifd": None,
|
||||
"skip_reason": "no valid assistant turns",
|
||||
"ifd_turns": [],
|
||||
}
|
||||
)
|
||||
results.append({**item, "ifd": None, "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(
|
||||
{
|
||||
**item,
|
||||
"ifd": None,
|
||||
"ifd_detail": {"skip_reason": "empty ctx or resp"},
|
||||
}
|
||||
)
|
||||
results.append({**item, "ifd": None, "ifd_detail": {"error": "empty"}})
|
||||
continue
|
||||
buffer.append((item, [(ctx_ids, resp_ids)], "plain"))
|
||||
|
||||
if len(buffer) >= batch_size:
|
||||
_flush_buffer(
|
||||
buffer,
|
||||
results,
|
||||
all_ifds,
|
||||
model,
|
||||
device,
|
||||
max_len,
|
||||
sentinel_ids,
|
||||
per_token,
|
||||
)
|
||||
_flush_buffer(buffer, results, all_ifds, model, device, max_len)
|
||||
|
||||
if buffer:
|
||||
_flush_buffer(
|
||||
buffer, results, all_ifds, model, device, max_len, sentinel_ids, per_token
|
||||
)
|
||||
_flush_buffer(buffer, results, all_ifds, model, device, max_len)
|
||||
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
for item in results:
|
||||
|
|
@ -358,22 +236,19 @@ def process_file(
|
|||
valid_ifd = [v for v in all_ifds if v is not None]
|
||||
if valid_ifd:
|
||||
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" Samples: {len(data)}")
|
||||
print(f" Valid IFD: {len(valid_ifd)}")
|
||||
print(f" Mean IFD: {statistics.mean(valid_ifd):.4f}")
|
||||
print(f" Median IFD: {statistics.median(valid_ifd):.4f}")
|
||||
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(
|
||||
buffer, results, all_ifds, model, device, max_len, sentinel_ids, per_token
|
||||
):
|
||||
def _flush_buffer(buffer, results, all_ifds, model, device, max_len=2048):
|
||||
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(
|
||||
all_pairs,
|
||||
model,
|
||||
device,
|
||||
max_len,
|
||||
sentinel_ids=sentinel_ids,
|
||||
per_token=per_token,
|
||||
)
|
||||
raw = _score_batch(all_pairs, model, device, max_len)
|
||||
|
||||
for item, turns, fmt, start, end in indices:
|
||||
turn_scores = raw[start:end]
|
||||
if fmt == "messages":
|
||||
valid = [
|
||||
s for s in turn_scores if s is not None and s.get("ifd") is not None
|
||||
]
|
||||
valid = [s for s in turn_scores if s is not None and s["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"))
|
||||
results.append({**item, "ifd": score.get("ifd"), "ifd_detail": score})
|
||||
all_ifds.append(score["ifd"])
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in New Issue