Spaces:
Runtime error
Runtime error
| """ | |
| 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=30, | |
| 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 = 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=22, | |
| 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, | |
| # num_beams=5, | |
| # no_repeat_ngram_size=2, | |
| # early_stopping=True, | |
| # temperature=temperature, | |
| # # top_k=top_k, top_p=top_p, | |
| # num_return_sequences=num_return_sequences, | |
| # max_time=timeout, | |
| # return_full_text=full_text, | |
| # clean_up_tokenization_spaces=True, | |
| # ) | |
| 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, | |
| 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 | |
| ) | |
| if verbose: | |
| print(f"DEBUG: {bot_dialogue} was original response pre-SC") | |
| return bot_dialogue # | |
| def consolidate_texts(name_resp: str, model_resp: list, name_spk: str, verbose=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 | |
| 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] | |
| 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 is not None and name_spk.lower() in resline.lower(): | |
| break # the name of the speaker is in the line, so we're done | |
| if ":" in resline and name_counter > 0: | |
| if break_safe: | |
| # we know this is a response from the bot even tho ':' is in the line | |
| fn_resp.append( | |
| resline | |
| ) # TODO: revisit the logic here, other names besides the bot could be in the line | |
| break_safe = False | |
| else: | |
| # don't have confidence in the line, so don't add it to the list. break out of the loop | |
| break | |
| else: | |
| fn_resp.append(resline) | |
| break_safe = False | |
| if verbose: | |
| print("the full response is:\n") | |
| print("\n".join(fn_resp)) | |
| return fn_resp | |