"""
converse.py - this script has functions for handling the conversation between the user and the bot.
https://huggingface.co/docs/transformers/v4.15.0/en/main_classes/model#transformers.generation_utils.GenerationMixin.generate.no_repeat_ngram_size
"""
import pprint as pp
import time
import torch
import transformers
from grammar_improve import remove_trailing_punctuation
def discussion(
prompt_text: str,
speaker: str,
responder: str,
pipeline,
timeout=45,
max_length=128,
top_p=0.95,
top_k=50,
temperature=0.7,
full_text=False,
num_return_sequences=1,
device=-1,
verbose=False,
):
"""
discussion - a function that takes in a prompt and generates a response. This function is meant to be used in a conversation loop, and is the main function for the bot.
Parameters
----------
prompt_text : str, the prompt to ask the bot, usually the user's question
speaker : str, the name of the person who is speaking the prompt
responder : str, the name of the person who is responding to the prompt
pipeline : transformers.Pipeline, the pipeline to use for generating the response
timeout : int, optional, the number of seconds to wait before timing out, by default 45
max_length : int, optional, the maximum number of tokens to generate, defaults to 128
top_p : float, optional, the top probability to use for sampling, defaults to 0.95
top_k : int, optional, the top k to use for sampling, defaults to 50
temperature : float, optional, the temperature to use for sampling, defaults to 0.7
full_text : bool, optional, whether to return the full text or just the generated text, defaults to False
num_return_sequences : int, optional, the number of sequences to return, defaults to 1
device : int, optional, the device to use for generation, defaults to -1 (CPU)
verbose : bool, optional, whether to print the generated text, defaults to False
Returns
-------
str, the generated text
"""
p_list = [] # track conversation
p_list.append(speaker.lower() + ":" + "\n")
p_list.append(prompt_text.lower() + "\n")
p_list.append("\n")
p_list.append(responder.lower() + ":" + "\n")
this_prompt = "".join(p_list)
if verbose:
print("overall prompt:\n")
pp.pprint(this_prompt, indent=4)
# call the model
print("\n... generating...")
bot_dialogue = gen_response(
this_prompt,
pipeline,
speaker,
responder,
timeout=timeout,
max_length=max_length,
top_p=top_p,
top_k=top_k,
temperature=temperature,
full_text=full_text,
num_return_sequences=num_return_sequences,
device=device,
verbose=verbose,
)
if isinstance(bot_dialogue, list) and len(bot_dialogue) > 1:
bot_resp = ", ".join(bot_dialogue)
elif isinstance(bot_dialogue, list) and len(bot_dialogue) == 1:
bot_resp = bot_dialogue[0]
else:
bot_resp = bot_dialogue
bot_resp = " ".join(bot_resp) if isinstance(bot_resp, list) else bot_resp
bot_resp = bot_resp.strip()
# remove the last ',' '.' chars
bot_resp = remove_trailing_punctuation(bot_resp)
if verbose:
print("\n... bot response:\n")
pp.pprint(bot_resp)
p_list.append(bot_resp + "\n")
p_list.append("\n")
print("\nfinished!")
# return the bot response and the full conversation
return {"out_text": bot_resp, "full_conv": p_list}
def gen_response(
query: str,
pipeline,
speaker: str,
responder: str,
timeout=45,
max_length=128,
top_p=0.95,
top_k=50,
temperature=0.7,
full_text=False,
num_return_sequences=1,
device=-1,
verbose=False,
**kwargs,
):
"""
gen_response - a function that takes in a prompt and generates a response using the pipeline. This operates underneath the discussion function.
Parameters
----------
query : str, the prompt to ask the bot, usually the user's question
speaker : str, the name of the person who is speaking the prompt
responder : str, the name of the person who is responding to the prompt
pipeline : transformers.Pipeline, the pipeline to use for generating the response
timeout : int, optional, the number of seconds to wait before timing out, by default 45
max_length : int, optional, the maximum number of tokens to generate, defaults to 128
top_p : float, optional, the top probability to use for sampling, defaults to 0.95
top_k : int, optional, the top k to use for sampling, defaults to 50
temperature : float, optional, the temperature to use for sampling, defaults to 0.7
full_text : bool, optional, whether to return the full text or just the generated text, defaults to False
num_return_sequences : int, optional, the number of sequences to return, defaults to 1
device : int, optional, the device to use for generation, defaults to -1 (CPU)
verbose : bool, optional, whether to print the generated text, defaults to False
Returns
-------
str, the generated text
"""
if max_length > 1024:
max_length = 1024
print("max_length is too large, setting to 1024")
st = time.perf_counter()
response = pipeline(
query,
max_length=max_length,
temperature=temperature,
top_k=top_k,
top_p=top_p,
num_return_sequences=num_return_sequences,
max_time=timeout,
return_full_text=full_text,
no_repeat_ngram_size=3,
length_penalty=0.3,
repetition_penalty=3.4,
clean_up_tokenization_spaces=True,
**kwargs,
) # the likely better beam-less method
rt = round(time.perf_counter() - st, 2)
if verbose:
print(f"took {rt} sec to respond")
if verbose:
print("\n[DEBUG] generated:\n")
pp.pprint(response) # for debugging
# process the full result to get the ~bot response~ piece
this_result = str(response[0]["generated_text"]).split(
"\n"
) # TODO: adjust hardcoded value for index to dynamic (if n>1)
bot_dialogue = consolidate_texts(
name_resp=responder,
model_resp=this_result,
name_spk=speaker,
verbose=verbose,
print_debug=True,
)
if verbose:
print(f"DEBUG: {bot_dialogue} was original response pre-SC")
return bot_dialogue #
def consolidate_texts(
model_resp: list,
name_resp: str = None,
name_spk: str = None,
verbose=False,
print_debug=False,
):
"""
consolidate_texts - given a list with speaker name followed by speaker text, returns all consecutive values of the first speaker name
Parameters:
name_resp (str): the name of the person who is responding
model_resp (list): the list of strings to consolidate (usually from the model)
name_spk (str): the name of the person who is speaking
verbose (bool): whether to print the results
print_debug (bool): whether to print the debug info during looping
Returns:
list, a list of all the consecutive messages of the first speaker name
"""
assert len(model_resp) > 0, "model_resp is empty"
if len(model_resp) == 1:
return model_resp[0]
name_resp = "person beta" if name_resp is None else name_resp
name_spk = "person alpha" if name_spk is None else name_spk
if verbose:
print("====" * 10)
print(
f"\n[DEBUG] initial model_resp has {len(model_resp)} lines: \n\t{model_resp}"
)
print(
f" the first element is \n\t{model_resp[0]} and it is {type(model_resp[0])}"
)
fn_resp = []
name_counter = 0
break_safe = False
for resline in model_resp:
if name_resp.lower() in resline:
name_counter += 1
break_safe = True # know the line is from bot as this line starts with the name of the bot
continue # don't add this line to the list
if name_spk.lower() in resline.lower():
if print_debug:
print(f"\nDEBUG: \n\t{resline}\ncaused the break")
break # the name of the speaker is in the line, so we're done
if (
any([": " in resline, ":\n" in resline])
and name_resp.lower() not in resline.lower()
):
if print_debug:
print(f"\nDEBUG: \n\t{resline}\ncaused the break")
break
else:
fn_resp.append(resline)
break_safe = False
if verbose:
print("--" * 10)
print("\nthe full response is:\n")
print("\n".join(fn_resp))
print("--" * 10)
return fn_resp