Create ai_single_response.py

ai_single_response.py

@@ -0,0 +1,383 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+ai_single_response.py - a script to generate a response to a prompt from a pretrained GPT model
+
+example:
+*\gpt2_chatbot> python ai_single_response.py --model "GPT2_conversational_355M_WoW10k" --prompt "hey, what's up?" --time
+
+query_gpt_model is used throughout the code, and is the "fundamental" building block of the bot and how everything works. I would recommend testing this function with a few different models.
+
+"""
+import argparse
+import pprint as pp
+import sys
+import time
+import warnings
+from datetime import datetime
+from pathlib import Path
+import logging
+
+logging.basicConfig(
+    filename=f"LOGFILE-{Path(__file__).stem}.log",
+    filemode="a",
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    level=logging.INFO,
+)
+
+from utils import DisableLogger, print_spacer, remove_trailing_punctuation
+
+with DisableLogger():
+    from cleantext import clean
+
+warnings.filterwarnings(action="ignore", message=".*gradient_checkpointing*")
+
+from aitextgen import aitextgen
+
+
+def extract_response(full_resp: list, plist: list, verbose: bool = False):
+    """
+    extract_response - helper fn for ai_single_response.py. By default aitextgen returns the prompt and the response, we just want the response
+
+    Args:
+        full_resp (list): the full response from aitextgen
+        plist (list): the prompt list
+        verbose (bool, optional): Defaults to False.
+
+    Returns:
+        response (str): the response, without the prompt
+    """
+    bot_response = []
+    for line in full_resp:
+        if line.lower() in plist and len(bot_response) < len(plist):
+            first_loc = plist.index(line)
+            del plist[first_loc]
+            continue
+        bot_response.append(line)
+    full_resp = [clean(ele, lower=False) for ele in bot_response]
+
+    if verbose:
+        print("the isolated responses are:\n")
+        pp.pprint(full_resp)
+        print_spacer()
+        print("the input prompt was:\n")
+        pp.pprint(plist)
+        print_spacer()
+    return full_resp  # list of only the model generated responses
+
+
+def get_bot_response(
+    name_resp: str, model_resp: list, name_spk: str, verbose: bool = False
+):
+    """
+    get_bot_response - gets the bot response to a prompt, checking to ensure that additional statements by the "speaker" are not included in the response.
+
+    Args:
+        name_resp (str): the name of the responder
+        model_resp (list): the model response
+        name_spk (str): the name of the speaker
+        verbose (bool, optional): Defaults to False.
+
+    Returns:
+        bot_response (str): the bot response, isolated down to just text without the "name tokens" or further messages from the speaker.
+    """
+
+    fn_resp = []
+
+    name_counter = 0
+    break_safe = False
+    for resline in model_resp:
+        if name_resp.lower() in resline.lower():
+            name_counter += 1
+            break_safe = True
+            continue
+        if ":" in resline and name_resp.lower() not in resline.lower():
+            break
+        if name_spk.lower() in resline.lower() and not break_safe:
+            break
+        else:
+            fn_resp.append(resline)
+    if verbose:
+        print("the full response is:\n")
+        print("\n".join(fn_resp))
+
+    return fn_resp
+
+
+def query_gpt_model(
+    folder_path: str or Path,
+    prompt_msg: str,
+    conversation_history: list = None,
+    speaker: str = None,
+    responder: str = None,
+    resp_length: int = 48,
+    kparam: int = 20,
+    temp: float = 0.4,
+    top_p: float = 0.9,
+    aitextgen_obj=None,
+    verbose: bool = False,
+    use_gpu: bool = False,
+):
+    """
+    query_gpt_model - queries the GPT model and returns the first response by <responder>
+
+    Args:
+        folder_path (str or Path): the path to the model folder
+        prompt_msg (str): the prompt message
+        conversation_history (list, optional): the conversation history. Defaults to None.
+        speaker (str, optional): the name of the speaker. Defaults to None.
+        responder (str, optional): the name of the responder. Defaults to None.
+        resp_length (int, optional): the length of the response in tokens. Defaults to 48.
+        kparam (int, optional): the k parameter for the top_k. Defaults to 40.
+        temp (float, optional): the temperature for the softmax. Defaults to 0.7.
+        top_p (float, optional): the top_p parameter for nucleus sampling. Defaults to 0.9.
+        aitextgen_obj (_type_, optional): a pre-loaded aitextgen object. Defaults to None.
+        verbose (bool, optional): Defaults to False.
+        use_gpu (bool, optional): Defaults to False.
+
+    Returns:
+        model_resp (dict): the model response, as a dict with the following keys: out_text (str) the generated text and full_conv (dict) the conversation history
+    """
+
+    try:
+        ai = (
+            aitextgen_obj
+            if aitextgen_obj
+            else aitextgen(
+                model_folder=folder_path,
+                to_gpu=use_gpu,
+            )
+        )
+    except Exception as e:
+        print(f"Unable to initialize aitextgen model: {e}")
+        print(
+            f"Check model folder: {folder_path}, run the download_models.py script to download the model files"
+        )
+        sys.exit(1)
+
+    mpath = Path(folder_path)
+    mpath_base = (
+        mpath.stem
+    )  # only want the base name of the model folder for check below
+    # these models used person alpha and person beta in training
+    mod_ids = ["natqa", "dd", "trivqa", "wow", "conversational"]
+    if any(substring in str(mpath_base).lower() for substring in mod_ids):
+        speaker = "person alpha" if speaker is None else speaker
+        responder = "person beta" if responder is None else responder
+    else:
+        if verbose:
+            print("speaker and responder not set - using default")
+        speaker = "person" if speaker is None else speaker
+        responder = "george robot" if responder is None else responder
+
+    prompt_list = (
+        conversation_history if conversation_history is not None else []
+    )  # track conversation
+    prompt_list.append(speaker.lower() + ":" + "\n")
+    prompt_list.append(prompt_msg.lower() + "\n")
+    prompt_list.append("\n")
+    prompt_list.append(responder.lower() + ":" + "\n")
+    this_prompt = "".join(prompt_list)
+    pr_len = len(this_prompt)
+    if verbose:
+        print("overall prompt:\n")
+        pp.pprint(prompt_list)
+    # call the model
+    print("\n... generating...")
+    this_result = ai.generate(
+        n=1,
+        top_k=kparam,
+        batch_size=128,
+        # the prompt input counts for text length constraints
+        max_length=resp_length + pr_len,
+        min_length=16 + pr_len,
+        prompt=this_prompt,
+        temperature=temp,
+        top_p=top_p,
+        do_sample=True,
+        return_as_list=True,
+        use_cache=True,
+    )
+    if verbose:
+        print("\n... generated:\n")
+        pp.pprint(this_result)  # for debugging
+    # process the full result to get the ~bot response~ piece
+    this_result = str(this_result[0]).split("\n")
+    input_prompt = this_prompt.split("\n")
+
+    diff_list = extract_response(
+        this_result, input_prompt, verbose=verbose
+    )  # isolate the responses from the prompts
+    # extract the bot response from the model generated text
+    bot_dialogue = get_bot_response(
+        name_resp=responder, model_resp=diff_list, name_spk=speaker, verbose=verbose
+    )
+    bot_resp = ", ".join(bot_dialogue)
+    bot_resp = remove_trailing_punctuation(
+        bot_resp.strip()
+    )  # remove trailing punctuation to seem more natural
+    if verbose:
+        print("\n... bot response:\n")
+        pp.pprint(bot_resp)
+    prompt_list.append(bot_resp + "\n")
+    prompt_list.append("\n")
+    conv_history = {}
+    for i, line in enumerate(prompt_list):
+        if i not in conv_history.keys():
+            conv_history[i] = line
+    if verbose:
+        print("\n... conversation history:\n")
+        pp.pprint(conv_history)
+    print("\nfinished!")
+
+    # return the bot response and the full conversation
+    return {"out_text": bot_resp, "full_conv": conv_history}
+
+
+# Set up the parsing of command-line arguments
+def get_parser():
+    """
+    get_parser [a helper function for the argparse module]
+
+    Returns: argparse.ArgumentParser
+    """
+
+    parser = argparse.ArgumentParser(
+        description="submit a message and have a pretrained GPT model respond"
+    )
+    parser.add_argument(
+        "-p",
+        "--prompt",
+        required=True,  # MUST HAVE A PROMPT
+        type=str,
+        help="the message the bot is supposed to respond to. Prompt is said by speaker, answered by responder.",
+    )
+    parser.add_argument(
+        "-m",
+        "--model",
+        required=False,
+        type=str,
+        default="distilgpt2-tiny-conversational",
+        help="folder - with respect to git directory of your repo that has the model files in it (pytorch.bin + "
+        "config.json). You can also pass the huggingface model name (e.g. distilgpt2)",
+    )
+
+    parser.add_argument(
+        "-s",
+        "--speaker",
+        required=False,
+        default=None,
+        help="Who the prompt is from (to the bot). Primarily relevant to bots trained on multi-individual chat data",
+    )
+    parser.add_argument(
+        "-r",
+        "--responder",
+        required=False,
+        default="person beta",
+        help="who the responder is. Primarily relevant to bots trained on multi-individual chat data",
+    )
+
+    parser.add_argument(
+        "--topk",
+        required=False,
+        type=int,
+        default=20,
+        help="how many responses to sample (positive integer). lower = more random responses",
+    )
+
+    parser.add_argument(
+        "--temp",
+        required=False,
+        type=float,
+        default=0.4,
+        help="specify temperature hyperparam (0-1). roughly considered as 'model creativity'",
+    )
+
+    parser.add_argument(
+        "--topp",
+        required=False,
+        type=float,
+        default=0.9,
+        help="nucleus sampling frac (0-1). aka: what fraction of possible options are considered?",
+    )
+
+    parser.add_argument(
+        "--resp_length",
+        required=False,
+        type=int,
+        default=50,
+        help="max length of the response (positive integer)",
+    )
+
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        default=False,
+        action="store_true",
+        help="pass this argument if you want all the printouts",
+    )
+
+    parser.add_argument(
+        "-rt",
+        "--time",
+        default=False,
+        action="store_true",
+        help="pass this argument if you want to know runtime",
+    )
+
+    parser.add_argument(
+        "--use_gpu",
+        required=False,
+        action="store_true",
+        help="use gpu if available",
+    )
+
+    return parser
+
+
+if __name__ == "__main__":
+    # parse the command line arguments
+    args = get_parser().parse_args()
+    query = args.prompt
+    model_dir = str(args.model)
+    model_loc = Path.cwd() / model_dir if "/" not in model_dir else model_dir
+    spkr = args.speaker
+    rspndr = args.responder
+    k_results = args.topk
+    my_temp = args.temp
+    my_top_p = args.topp
+    resp_length = args.resp_length
+    assert resp_length > 0, "response length must be positive"
+    want_verbose = args.verbose
+    want_rt = args.time
+    use_gpu = args.use_gpu
+
+    st = time.perf_counter()
+
+    resp = query_gpt_model(
+        folder_path=model_loc,
+        prompt_msg=query,
+        speaker=spkr,
+        responder=rspndr,
+        kparam=k_results,
+        temp=my_temp,
+        top_p=my_top_p,
+        resp_length=resp_length,
+        verbose=want_verbose,
+        use_gpu=use_gpu,
+    )
+
+    output = resp["out_text"]
+    pp.pprint(output, indent=4)
+
+    rt = round(time.perf_counter() - st, 1)
+
+    if want_rt:
+        print("took {runtime} seconds to generate. \n".format(runtime=rt))
+
+    if want_verbose:
+        print("finished - ", datetime.now())
+        p_list = resp["full_conv"]
+        print("A transcript of your chat is as follows: \n")
+        p_list = [item.strip() for item in p_list]
+        pp.pprint(p_list)