| import gc |
| from queue import Queue |
| from threading import Thread |
|
|
| import torch |
| import transformers |
|
|
| import modules.shared as shared |
|
|
| |
| class _SentinelTokenStoppingCriteria(transformers.StoppingCriteria): |
|
|
| def __init__(self, sentinel_token_ids: torch.LongTensor, |
| starting_idx: int): |
| transformers.StoppingCriteria.__init__(self) |
| self.sentinel_token_ids = sentinel_token_ids |
| self.starting_idx = starting_idx |
|
|
| def __call__(self, input_ids: torch.LongTensor, |
| _scores: torch.FloatTensor) -> bool: |
| for sample in input_ids: |
| trimmed_sample = sample[self.starting_idx:] |
| |
| if trimmed_sample.shape[-1] < self.sentinel_token_ids.shape[-1]: |
| continue |
|
|
| for window in trimmed_sample.unfold( |
| 0, self.sentinel_token_ids.shape[-1], 1): |
| if torch.all(torch.eq(self.sentinel_token_ids, window)): |
| return True |
| return False |
|
|
| class Stream(transformers.StoppingCriteria): |
| def __init__(self, callback_func=None): |
| self.callback_func = callback_func |
|
|
| def __call__(self, input_ids, scores) -> bool: |
| if self.callback_func is not None: |
| self.callback_func(input_ids[0]) |
| return False |
|
|
| class Iteratorize: |
|
|
| """ |
| Transforms a function that takes a callback |
| into a lazy iterator (generator). |
| """ |
|
|
| def __init__(self, func, kwargs={}, callback=None): |
| self.mfunc=func |
| self.c_callback=callback |
| self.q = Queue() |
| self.sentinel = object() |
| self.kwargs = kwargs |
| self.stop_now = False |
|
|
| def _callback(val): |
| if self.stop_now: |
| raise ValueError |
| self.q.put(val) |
|
|
| def gentask(): |
| try: |
| ret = self.mfunc(callback=_callback, **self.kwargs) |
| except ValueError: |
| pass |
| clear_torch_cache() |
| self.q.put(self.sentinel) |
| if self.c_callback: |
| self.c_callback(ret) |
|
|
| self.thread = Thread(target=gentask) |
| self.thread.start() |
|
|
| def __iter__(self): |
| return self |
|
|
| def __next__(self): |
| obj = self.q.get(True,None) |
| if obj is self.sentinel: |
| raise StopIteration |
| else: |
| return obj |
|
|
| def __del__(self): |
| clear_torch_cache() |
|
|
| def __enter__(self): |
| return self |
|
|
| def __exit__(self, exc_type, exc_val, exc_tb): |
| self.stop_now = True |
| clear_torch_cache() |
|
|
| def clear_torch_cache(): |
| gc.collect() |
| if not shared.args.cpu: |
| torch.cuda.empty_cache() |
|
|
