Spaces:
Runtime error
Runtime error
File size: 2,992 Bytes
6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 292c2df 6926a80 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import gc
import traceback
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: list, starting_idx: int):
transformers.StoppingCriteria.__init__(self)
self.sentinel_token_ids = sentinel_token_ids
self.starting_idx = starting_idx
self.shortest = min([x.shape[-1] for x in sentinel_token_ids])
def __call__(self, input_ids: torch.LongTensor, _scores: torch.FloatTensor) -> bool:
for sample in input_ids:
trimmed_sample = sample[self.starting_idx:]
trimmed_len = trimmed_sample.shape[-1]
if trimmed_len < self.shortest:
continue
for sentinel in self.sentinel_token_ids:
sentinel_len = sentinel.shape[-1]
if trimmed_len < sentinel_len:
continue
window = trimmed_sample[-sentinel_len:]
if torch.all(torch.eq(sentinel, 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).
Adapted from: https://stackoverflow.com/a/9969000
"""
def __init__(self, func, kwargs=None, callback=None):
self.mfunc = func
self.c_callback = callback
self.q = Queue()
self.sentinel = object()
self.kwargs = kwargs or {}
self.stop_now = False
def _callback(val):
if self.stop_now or shared.stop_everything:
raise ValueError
self.q.put(val)
def gentask():
try:
ret = self.mfunc(callback=_callback, **self.kwargs)
except ValueError:
pass
except:
traceback.print_exc()
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()
|