mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-04-19 22:02:06 +00:00
25 lines
556 B
Python
25 lines
556 B
Python
from typing import Dict, Optional
|
|
|
|
from text_generation.models.types import Batch
|
|
|
|
|
|
class Cache:
|
|
def __init__(self):
|
|
self.cache: Dict[int, Batch] = {}
|
|
|
|
def pop(self, batch_id: int) -> Optional[Batch]:
|
|
return self.cache.pop(batch_id, None)
|
|
|
|
def set(self, entry: Batch):
|
|
if entry is not None:
|
|
self.cache[entry.batch_id] = entry
|
|
|
|
def delete(self, batch_id: int):
|
|
del self.cache[batch_id]
|
|
|
|
def clear(self):
|
|
self.cache.clear()
|
|
|
|
def __len__(self):
|
|
return len(self.cache.keys())
|