Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
KoboldAI
GitHub Repository: KoboldAI/KoboldAI-Client
Path: blob/main/structures.py
471 views
1
import collections
2
from typing import Iterable, Tuple
3
4
5
class KoboldStoryRegister(collections.OrderedDict):
6
'''
7
Complexity-optimized class for keeping track of story chunks
8
'''
9
10
def __init__(self, sequence: Iterable[Tuple[int, str]] = ()):
11
super().__init__(sequence)
12
self.__next_id: int = len(sequence)
13
14
def append(self, v: str) -> None:
15
self[self.__next_id] = v
16
self.increment_id()
17
18
def pop(self) -> str:
19
return self.popitem()[1]
20
21
def get_first_key(self) -> int:
22
if len(self) == 0:
23
return -1
24
else:
25
return next(iter(self))
26
27
def get_last_key(self) -> int:
28
if len(self) == 0:
29
return -1
30
else:
31
return next(reversed(self))
32
33
def __getitem__(self, k: int) -> str:
34
return super().__getitem__(k)
35
36
def __setitem__(self, k: int, v: str) -> None:
37
return super().__setitem__(k, v)
38
39
def increment_id(self) -> None:
40
self.__next_id += 1
41
42
def get_next_id(self) -> int:
43
return self.__next_id
44
45
def set_next_id(self, x: int) -> None:
46
self.__next_id = x
47
48