Refactor text_generation.py, add support for custom generation functions (#1817)

This commit is contained in:
oobabooga 2023-05-05 18:53:03 -03:00 committed by GitHub
parent 876fbb97c0
commit 8aafb1f796
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 289 additions and 195 deletions

View file

@ -75,7 +75,6 @@ if params['manage_VRAM']:
samplers = ['DDIM', 'DPM++ 2M Karras'] # TODO: get the availible samplers with http://{address}}/sdapi/v1/samplers
SD_models = ['NeverEndingDream'] # TODO: get with http://{address}}/sdapi/v1/sd-models and allow user to select
streaming_state = shared.args.no_stream # remember if chat streaming was enabled
picture_response = False # specifies if the next model response should appear as a picture
def remove_surrounded_chars(string):
@ -92,6 +91,13 @@ def triggers_are_in(string):
return bool(re.search('(?aims)(send|mail|message|me)\\b.+?\\b(image|pic(ture)?|photo|snap(shot)?|selfie|meme)s?\\b', string))
def state_modifier(state):
if picture_response:
state['stream'] = False
return state
def input_modifier(string):
"""
This function is applied to your text inputs before
@ -218,14 +224,13 @@ def bot_prefix_modifier(string):
def toggle_generation(*args):
global picture_response, shared, streaming_state
global picture_response, shared
if not args:
picture_response = not picture_response
else:
picture_response = args[0]
shared.args.no_stream = True if picture_response else streaming_state # Disable streaming cause otherwise the SD-generated picture would return as a dud
shared.processing_message = "*Is sending a picture...*" if picture_response else "*Is typing...*"