zhengjl/text-generation-webui
Public Access
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Make the code more like PEP8

Browse Source
This commit is contained in:
oobabooga
2023-04-07 00:00:16 -03:00
parent 8e4e8dd741
commit 7b1b1d8e39
30 changed files with 314 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
modules/callbacks.py
View File

@@ -30,6 +30,7 @@ class _SentinelTokenStoppingCriteria(transformers.StoppingCriteria):
return True
return False
class Stream(transformers.StoppingCriteria):
def __init__(self, callback_func=None):
self.callback_func = callback_func
@@ -39,6 +40,7 @@ class Stream(transformers.StoppingCriteria):
self.callback_func(input_ids[0])
return False
class Iteratorize:
"""
@@ -47,8 +49,8 @@ class Iteratorize:
"""
def __init__(self, func, kwargs={}, callback=None):
self.mfunc=func
self.c_callback=callback
self.mfunc = func
self.c_callback = callback
self.q = Queue()
self.sentinel = object()
self.kwargs = kwargs
@@ -80,7 +82,7 @@ class Iteratorize:
return self
def __next__(self):
obj = self.q.get(True,None)
obj = self.q.get(True, None)
if obj is self.sentinel:
raise StopIteration
else:
@@ -96,6 +98,7 @@ class Iteratorize:
self.stop_now = True
clear_torch_cache()
def clear_torch_cache():
gc.collect()
if not shared.args.cpu: