Upload 7 files

app.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+import argparse
+import logging
+import typing as t
+
+from gradio_ui import build_gradio_ui_for
+from koboldai_client import run_raw_inference_on_kai, wait_for_kai_server
+from parsing import parse_messages_from_str
+from prompting import build_prompt_for
+from utils import clear_stdout
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# For UI debugging purposes.
+DONT_USE_MODEL = False
+
+
+def main(server_port: int,
+         share_gradio_link: bool = False,
+         model_name: t.Optional[str] = None,
+         koboldai_url: t.Optional[str] = None) -> None:
+    '''Script entrypoint.'''
+    if model_name and not DONT_USE_MODEL:
+        from model import build_model_and_tokenizer_for, run_raw_inference
+        model, tokenizer = build_model_and_tokenizer_for(model_name)
+    else:
+        model, tokenizer = None, None
+
+    def inference_fn(history: t.List[str], user_input: str,
+                     generation_settings: t.Dict[str, t.Any],
+                     *char_settings: t.Any) -> str:
+        if DONT_USE_MODEL:
+            return "Mock response for UI tests."
+
+        # Brittle. Comes from the order defined in gradio_ui.py.
+        [
+            char_name,
+            _user_name,
+            char_persona,
+            char_greeting,
+            world_scenario,
+            example_dialogue,
+        ] = char_settings
+
+        # If we're just starting the conversation and the character has a greeting
+        # configured, return that instead. This is a workaround for the fact that
+        # Gradio assumed that a chatbot cannot possibly start a conversation, so we
+        # can't just have the greeting there automatically, it needs to be in
+        # response to a user message.
+        if len(history) == 0 and char_greeting is not None:
+            return f"{char_name}: {char_greeting}"
+
+        prompt = build_prompt_for(history=history,
+                                  user_message=user_input,
+                                  char_name=char_name,
+                                  char_persona=char_persona,
+                                  example_dialogue=example_dialogue,
+                                  world_scenario=world_scenario)
+
+        if model and tokenizer:
+            model_output = run_raw_inference(model, tokenizer, prompt,
+                                             user_input, **generation_settings)
+        elif koboldai_url:
+            model_output = f"{char_name}:"
+            model_output += run_raw_inference_on_kai(koboldai_url, prompt,
+                                                     **generation_settings)
+        else:
+            raise Exception(
+                "Not using local inference, but no Kobold instance URL was"
+                " given. Nowhere to perform inference on.")
+
+        generated_messages = parse_messages_from_str(model_output,
+                                                     ["You", char_name])
+        logger.debug("Parsed model response is: `%s`", generated_messages)
+        bot_message = generated_messages[0]
+
+        return bot_message
+
+    ui = build_gradio_ui_for(inference_fn, for_kobold=koboldai_url is not None)
+    ui.launch(server_port=server_port, share=share_gradio_link)
+
+
+def _parse_args_from_argv() -> argparse.Namespace:
+    '''Parses arguments coming in from the command line.'''
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-m",
+        "--model-name",
+        help=
+        "HuggingFace Transformers model name, if not using a KoboldAI instance as an inference server.",
+    )
+    parser.add_argument(
+        "-p",
+        "--port",
+        type=int,
+        default=3000,
+        help="Port to listen on.",
+    )
+    parser.add_argument(
+        "-k",
+        "--koboldai-url",
+        help="URL to a KoboldAI instance to use as an inference server.",
+    )
+    parser.add_argument(
+        "-s",
+        "--share",
+        action="store_true",
+        help="Enable to generate a public link for the Gradio UI.",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = _parse_args_from_argv()
+
+    if args.koboldai_url:
+        # I have no idea how long a safe wait time is, but we'd rather wait for
+        # too long rather than just cut the user off _right_ when the setup is
+        # about to finish, so let's pick something absurd here.
+        wait_for_kai_server(args.koboldai_url, max_wait_time_seconds=60 * 30)
+
+        # Clear out any Kobold logs so the user can clearly see the Gradio link
+        # that's about to show up afterwards.
+        clear_stdout()
+
+    main(model_name=args.model_name,
+         server_port=args.port,
+         koboldai_url=args.koboldai_url,
+         share_gradio_link=args.share)

gradio_ui.py

@@ -0,0 +1,442 @@
+import json
+import logging
+
+import gradio as gr
+
+
+def get_generation_defaults(for_kobold):
+    defaults = {
+        "do_sample": True,
+        "max_new_tokens": 196,
+        "temperature": 0.5,
+        "top_p": 0.9,
+        "top_k": 0,
+        "typical_p": 1.0,
+        "repetition_penalty": 1.05,
+    }
+
+    if for_kobold:
+        defaults.update({"max_context_length": 768})
+    else:
+        defaults.update({"penalty_alpha": 0.6})
+
+    return defaults
+
+
+logger = logging.getLogger(__name__)
+
+
+def build_gradio_ui_for(inference_fn, for_kobold):
+    '''
+    Builds a Gradio UI to interact with the model. Big thanks to TearGosling for
+    the initial version that inspired this.
+    '''
+    with gr.Blocks(title="Pygmalion", analytics_enabled=False) as interface:
+        history_for_gradio = gr.State([])
+        history_for_model = gr.State([])
+        generation_settings = gr.State(
+            get_generation_defaults(for_kobold=for_kobold))
+
+        def _update_generation_settings(
+            original_settings,
+            param_name,
+            new_value,
+        ):
+            '''
+            Merges `{param_name: new_value}` into `original_settings` and
+            returns a new dictionary.
+            '''
+            updated_settings = {**original_settings, param_name: new_value}
+            logging.debug("Generation settings updated to: `%s`",
+                         updated_settings)
+            return updated_settings
+
+        def _run_inference(
+            model_history,
+            gradio_history,
+            user_input,
+            generation_settings,
+            *char_setting_states,
+        ):
+            '''
+            Runs inference on the model, and formats the returned response for
+            the Gradio state and chatbot component.
+            '''
+            char_name = char_setting_states[0]
+            user_name = char_setting_states[1]
+
+            # If user input is blank, format it as if user was silent
+            if user_input is None or user_input.strip() == "":
+                user_input = "..."
+
+            inference_result = inference_fn(model_history, user_input,
+                                            generation_settings,
+                                            *char_setting_states)
+
+            inference_result_for_gradio = inference_result \
+                .replace(f"{char_name}:", f"**{char_name}:**") \
+                .replace("<USER>", user_name) \
+                .replace("\n", "<br>") # Gradio chatbot component can display br tag as linebreak
+
+            model_history.append(f"You: {user_input}")
+            model_history.append(inference_result)
+            gradio_history.append((user_input, inference_result_for_gradio))
+
+            return None, model_history, gradio_history, gradio_history
+
+        def _regenerate(
+            model_history,
+            gradio_history,
+            generation_settings,
+            *char_setting_states,
+        ):
+            '''Regenerates the last response.'''
+            return _run_inference(
+                model_history[:-2],
+                gradio_history[:-1],
+                model_history[-2].replace("You: ", ""),
+                generation_settings,
+                *char_setting_states,
+            )
+
+        def _undo_last_exchange(model_history, gradio_history):
+            '''Undoes the last exchange (message pair).'''
+            return model_history[:-2], gradio_history[:-1], gradio_history[:-1]
+
+        def _save_chat_history(model_history, *char_setting_states):
+            '''Saves the current chat history to a .json file.'''
+            char_name = char_setting_states[0]
+            with open(f"{char_name}_conversation.json", "w") as f:
+                f.write(json.dumps({"chat": model_history}))
+            return f"{char_name}_conversation.json"
+
+        def _load_chat_history(file_obj, *char_setting_states):
+            '''Loads up a chat history from a .json file.'''
+            # #############################################################################################
+            # TODO(TG): Automatically detect and convert any CAI dump files loaded in to Pygmalion format #
+            # #############################################################################################
+
+            # https://stackoverflow.com/questions/5389507/iterating-over-every-two-elements-in-a-list
+            def pairwise(iterable):
+                # "s -> (s0, s1), (s2, s3), (s4, s5), ..."
+                a = iter(iterable)
+                return zip(a, a)
+
+            char_name = char_setting_states[0]
+            user_name = char_setting_states[1]
+
+            file_data = json.loads(file_obj.decode('utf-8'))
+            model_history = file_data["chat"]
+            # Construct a new gradio history
+            new_gradio_history = []
+            for human_turn, bot_turn in pairwise(model_history):
+                # Handle the situation where convo history may be loaded before character defs
+                if char_name == "":
+                    # Grab char name from the model history
+                    char_name = bot_turn.split(":")[0]
+                # Format the user and bot utterances
+                user_turn = human_turn.replace("You: ", "")
+                bot_turn = bot_turn.replace(f"{char_name}:", f"**{char_name}:**")
+
+                # Somebody released a script on /g/ which tries to convert CAI dump logs
+                # to Pygmalion character settings and chats. The anonymization of the dumps, however, means that
+                # [NAME_IN_MESSAGE_REDACTED] is left in the conversational history. We obviously wouldn't want this
+                # This therefore accomodates users of that script, so that [NAME_IN_MESSAGE_REDACTED] doesn't have
+                # to be manually edited in the conversation JSON.
+                # The model shouldn't generate [NAME_IN_MESSAGE_REDACTED] by itself.
+                user_turn = user_turn.replace("[NAME_IN_MESSAGE_REDACTED]", user_name)
+                bot_turn = bot_turn.replace("[NAME_IN_MESSAGE_REDACTED]", user_name)
+
+                new_gradio_history.append((user_turn, bot_turn))
+
+            return model_history, new_gradio_history, new_gradio_history
+
+        with gr.Tab("Character Settings") as settings_tab:
+            charfile, char_setting_states = _build_character_settings_ui()
+
+        with gr.Tab("Chat Window"):
+            chatbot = gr.Chatbot(
+                label="Your conversation will show up here").style(
+                    color_map=("#326efd", "#212528"))
+
+            char_name, _user_name, char_persona, char_greeting, world_scenario, example_dialogue = char_setting_states
+            charfile.upload(
+                fn=_char_file_upload,
+                inputs=[charfile, history_for_model, history_for_gradio],
+                outputs=[history_for_model, history_for_gradio, chatbot, char_name, char_persona, char_greeting, world_scenario, example_dialogue]
+            )
+
+            message = gr.Textbox(
+                label="Your message (hit Enter to send)",
+                placeholder="Write a message...",
+            )
+            message.submit(
+                fn=_run_inference,
+                inputs=[
+                    history_for_model, history_for_gradio, message,
+                    generation_settings, *char_setting_states
+                ],
+                outputs=[
+                    message, history_for_model, history_for_gradio, chatbot
+                ],
+            )
+
+            with gr.Row():
+                send_btn = gr.Button("Send", variant="primary")
+                send_btn.click(
+                    fn=_run_inference,
+                    inputs=[
+                        history_for_model, history_for_gradio, message,
+                        generation_settings, *char_setting_states
+                    ],
+                    outputs=[
+                        message, history_for_model, history_for_gradio, chatbot
+                    ],
+                )
+
+                regenerate_btn = gr.Button("Regenerate")
+                regenerate_btn.click(
+                    fn=_regenerate,
+                    inputs=[
+                        history_for_model, history_for_gradio,
+                        generation_settings, *char_setting_states
+                    ],
+                    outputs=[
+                        message, history_for_model, history_for_gradio, chatbot
+                    ],
+                )
+
+                undo_btn = gr.Button("Undo last exchange")
+                undo_btn.click(
+                    fn=_undo_last_exchange,
+                    inputs=[history_for_model, history_for_gradio],
+                    outputs=[history_for_model, history_for_gradio, chatbot],
+                )
+
+            with gr.Row():
+                with gr.Column():
+                    chatfile = gr.File(type="binary", file_types=[".json"], interactive=True)
+                    chatfile.upload(
+                        fn=_load_chat_history,
+                        inputs=[chatfile, *char_setting_states],
+                        outputs=[history_for_model, history_for_gradio, chatbot]
+                    )
+
+                    save_char_btn = gr.Button(value="Save Conversation History")
+                    save_char_btn.click(_save_chat_history, inputs=[history_for_model, *char_setting_states], outputs=[chatfile])
+                with gr.Column():
+                    gr.Markdown("""
+                        ### To save a chat
+                        Click "Save Conversation History". The file will appear above the button and you can click to download it.
+
+                        ### To load a chat
+                        Drag a valid .json file onto the upload box, or click the box to browse.
+
+                        **Remember to fill out/load up your character definitions before resuming a chat!**
+                    """)
+
+
+
+        with gr.Tab("Generation Settings"):
+            _build_generation_settings_ui(
+                state=generation_settings,
+                fn=_update_generation_settings,
+                for_kobold=for_kobold,
+            )
+
+    return interface
+
+
+def _char_file_upload(file_obj, history_model, history_gradio):
+    file_data = json.loads(file_obj.decode('utf-8'))
+    char_name = file_data["char_name"]
+    greeting = file_data["char_greeting"]
+    empty_history = not history_model or (len(history_model) <= 2 and history_model[0] == '')
+    if empty_history and char_name and greeting:
+        # if chat history is empty so far, and there is a character greeting, add character greeting to the chat
+        s = f'{char_name}: {greeting}'
+        t = f'**{char_name}**: {greeting}'
+        history_model = ['', s]
+        history_gradio = [('', t)]
+    return history_model, history_gradio, history_gradio, char_name, file_data["char_persona"], greeting, file_data["world_scenario"], file_data["example_dialogue"]
+
+def _build_character_settings_ui():
+
+    def char_file_create(char_name, char_persona, char_greeting, world_scenario, example_dialogue):
+        with open(char_name + ".json", "w") as f:
+            f.write(json.dumps({"char_name": char_name, "char_persona": char_persona, "char_greeting": char_greeting, "world_scenario": world_scenario, "example_dialogue": example_dialogue}))
+        return char_name + ".json"
+
+    with gr.Column():
+        with gr.Row():
+            char_name = gr.Textbox(
+                label="Character Name",
+                placeholder="The character's name",
+            )
+            user_name = gr.Textbox(
+                label="Your Name",
+                placeholder="How the character should call you",
+            )
+
+        char_persona = gr.Textbox(
+            label="Character Persona",
+            placeholder=
+            "Describe the character's persona here. Think of this as CharacterAI's description + definitions in one box.",
+            lines=4,
+        )
+        char_greeting = gr.Textbox(
+            label="Character Greeting",
+            placeholder=
+            "Write the character's greeting here. They will say this verbatim as their first response.",
+            lines=3,
+        )
+
+        world_scenario = gr.Textbox(
+            label="Scenario",
+            placeholder=
+            "Optionally, describe the starting scenario in a few short sentences.",
+        )
+        example_dialogue = gr.Textbox(
+            label="Example Chat",
+            placeholder=
+            "Optionally, write in an example chat here. This is useful for showing how the character should behave, for example.",
+            lines=4,
+        )
+
+        with gr.Row():
+            with gr.Column():
+                charfile = gr.File(type="binary", file_types=[".json"])
+
+                save_char_btn = gr.Button(value="Generate Character File")
+                save_char_btn.click(char_file_create, inputs=[char_name, char_persona, char_greeting, world_scenario, example_dialogue], outputs=[charfile])
+            with gr.Column():
+                gr.Markdown("""
+                    ### To save a character
+                    Click "Generate Character File". The file will appear above the button and you can click to download it.
+
+                    ### To upload a character
+                    Drag a valid .json file onto the upload box, or click the box to browse.
+                """)
+
+    return charfile, (char_name, user_name, char_persona, char_greeting, world_scenario, example_dialogue)
+
+
+def _build_generation_settings_ui(state, fn, for_kobold):
+    generation_defaults = get_generation_defaults(for_kobold=for_kobold)
+
+    with gr.Row():
+        with gr.Column():
+            max_new_tokens = gr.Slider(
+                16,
+                512,
+                value=generation_defaults["max_new_tokens"],
+                step=4,
+                label="max_new_tokens",
+            )
+            max_new_tokens.change(
+                lambda state, value: fn(state, "max_new_tokens", value),
+                inputs=[state, max_new_tokens],
+                outputs=state,
+            )
+
+            temperature = gr.Slider(
+                0.1,
+                2,
+                value=generation_defaults["temperature"],
+                step=0.01,
+                label="temperature",
+            )
+            temperature.change(
+                lambda state, value: fn(state, "temperature", value),
+                inputs=[state, temperature],
+                outputs=state,
+            )
+
+            top_p = gr.Slider(
+                0.0,
+                1.0,
+                value=generation_defaults["top_p"],
+                step=0.01,
+                label="top_p",
+            )
+            top_p.change(
+                lambda state, value: fn(state, "top_p", value),
+                inputs=[state, top_p],
+                outputs=state,
+            )
+
+        with gr.Column():
+            typical_p = gr.Slider(
+                0.0,
+                1.0,
+                value=generation_defaults["typical_p"],
+                step=0.01,
+                label="typical_p",
+            )
+            typical_p.change(
+                lambda state, value: fn(state, "typical_p", value),
+                inputs=[state, typical_p],
+                outputs=state,
+            )
+
+            repetition_penalty = gr.Slider(
+                1.0,
+                3.0,
+                value=generation_defaults["repetition_penalty"],
+                step=0.01,
+                label="repetition_penalty",
+            )
+            repetition_penalty.change(
+                lambda state, value: fn(state, "repetition_penalty", value),
+                inputs=[state, repetition_penalty],
+                outputs=state,
+            )
+
+            top_k = gr.Slider(
+                0,
+                100,
+                value=generation_defaults["top_k"],
+                step=1,
+                label="top_k",
+            )
+            top_k.change(
+                lambda state, value: fn(state, "top_k", value),
+                inputs=[state, top_k],
+                outputs=state,
+            )
+
+            if not for_kobold:
+                penalty_alpha = gr.Slider(
+                    0,
+                    1,
+                    value=generation_defaults["penalty_alpha"],
+                    step=0.05,
+                    label="penalty_alpha",
+                )
+                penalty_alpha.change(
+                    lambda state, value: fn(state, "penalty_alpha", value),
+                    inputs=[state, penalty_alpha],
+                    outputs=state,
+                )
+
+    #
+    # Some of these explanations are taken from Kobold:
+    # https://github.com/KoboldAI/KoboldAI-Client/blob/main/gensettings.py
+    #
+    # They're passed directly into the `generate` call, so they should exist here:
+    # https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationConfig
+    #
+    with gr.Accordion(label="Helpful information", open=False):
+        gr.Markdown("""
+        Here's a basic rundown of each setting:
+
+        - `max_new_tokens`: Number of tokens the AI should generate. Higher numbers will take longer to generate.
+        - `temperature`: Randomness of sampling. High values can increase creativity but may make text less sensible. Lower values will make text more predictable but can become repetitious.
+        - `top_p`: Used to discard unlikely text in the sampling process. Lower values will make text more predictable but can become repetitious. (Put this value on 1 to disable its effect)
+        - `top_k`: Alternative sampling method, can be combined with top_p. The number of highest probability vocabulary tokens to keep for top-k-filtering. (Put this value on 0 to disable its effect)
+        - `typical_p`: Alternative sampling method described in the paper "Typical_p Decoding for Natural Language Generation" (10.48550/ARXIV.2202.00666). The paper suggests 0.2 as a good value for this setting. Set this setting to 1 to disable its effect.
+        - `repetition_penalty`: Used to penalize words that were already generated or belong to the context (Going over 1.2 breaks 6B models. Set to 1.0 to disable).
+        - `penalty_alpha`: The alpha coefficient when using contrastive search.
+
+        Some settings might not show up depending on which inference backend is being used.
+        """)

koboldai_client.py

@@ -0,0 +1,117 @@
+import datetime
+import logging
+import time
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+class KoboldApiServerException(Exception):
+    pass
+
+
+def wait_for_kai_server(koboldai_url: str, max_wait_time_seconds: int) -> None:
+    '''Blocks until the KAI server is up.'''
+    start_time = datetime.datetime.now()
+
+    while True:
+        try:
+            requests.head(koboldai_url, timeout=(5, 5))
+            break
+        except requests.exceptions.ConnectionError as ex:
+            if "Connection refused" not in str(ex):
+                raise ex
+
+            abort_at = start_time + datetime.timedelta(
+                seconds=max_wait_time_seconds)
+
+            if datetime.datetime.now() > abort_at:
+                raise TimeoutError(
+                    f"Waited for {max_wait_time_seconds} seconds but KoboldAI"
+                    " server is still not up, aborting.")
+
+            time.sleep(1)
+
+
+def run_raw_inference_on_kai(
+    koboldai_url: str,
+    prompt: str,
+    max_new_tokens: int,
+    do_sample: bool,
+    typical_p: float,
+    repetition_penalty: float,
+    **kwargs,
+) -> str:
+    endpoint = f"{koboldai_url}/api/v1/generate"
+    payload = {
+        "prompt": prompt,
+
+        # Incredibly low max len for reasons explained in the "while True" loop
+        # below.
+        "max_length": 32,
+
+        # Take care of parameters which are named differently between Kobold and
+        # HuggingFace.
+        "sampler_full_determinism": not do_sample,
+        "typical": typical_p,
+        "rep_pen": repetition_penalty,
+
+        # Disable any pre or post-processing on the KoboldAI side, we'd rather
+        # take care of things on our own.
+        "frmttriminc": False,
+        "frmtrmspch": False,
+        "frmtrmblln": False,
+        "frmtadsnsp": False,
+
+        # Append any other generation parameters that we didn't handle manually.
+        **kwargs,
+    }
+    generated_text = ""
+
+    # Currently, Kobold doesn't support custom stopping criteria, and their chat
+    # mode can't handle multi-line responses. To work around both of those, we
+    # use the regular adventure mode generation but keep asking for more tokens
+    # until the model starts trying to talk as the user, then we stop.
+    attempts = 0
+    max_extra_attempts = 4
+    while attempts < (payload["max_length"] /
+                      max_new_tokens) + max_extra_attempts:
+        attempts += 1
+        response = requests.post(endpoint, json=payload)
+        if not response.ok:
+            error_message = response.text
+            raise KoboldApiServerException(
+                "The KoboldAI API server returned an error"
+                f" (HTTP status code {response.status_code}): {error_message}")
+
+        inference_result = response.json()["results"][0]["text"]
+        generated_text += inference_result
+
+        # Model started to talk as us. Stop generating and return results, the
+        # rest of the code will take care of trimming it properly.
+        if "\nYou:" in generated_text:
+            logger.debug("Hit `\nYou:`: `%s`", generated_text)
+            return generated_text
+
+        # For SFT: hit an EOS token. Trim and return.
+        if generated_text.endswith("<|endoftext|>"):
+            logger.debug("Got EOS token: `%s`", generated_text)
+
+            # We add a fake generated "\nYou:" here so the trimming code doesn't
+            # need to handle SFT and UFT models differently.
+            return generated_text.replace("<|endoftext|>", "\nYou:")
+
+        # Hit the configured generation limit.
+        if len(generated_text.split()) >= max_new_tokens:
+            logger.debug("Hit max length: `%s`", generated_text)
+            return generated_text
+
+        # Model still hasn't finished what it had to say. Append its output to
+        # the prompt and feed it back in.
+        logger.debug("Got another %s tokens, but still not done: `%s`",
+                     payload["max_length"], generated_text)
+        payload["prompt"] += inference_result
+
+    logger.debug("Exhausted generation attempts: `%s`", generated_text)
+    return generated_text

model.py

@@ -0,0 +1,110 @@
+import logging
+import typing as t
+
+import torch
+import transformers
+
+logger = logging.getLogger(__name__)
+
+
+def build_model_and_tokenizer_for(
+    model_name: str
+) -> t.Tuple[transformers.AutoModelForCausalLM, transformers.AutoTokenizer]:
+    '''Sets up the model and accompanying objects.'''
+    logger.info(f"Loading tokenizer for {model_name}")
+    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
+
+    # NOTE(11b): non-OPT models support passing this in at inference time, might
+    # be worth refactoring for a debug version so we're able to experiment on
+    # the fly
+    bad_words_ids = [
+        tokenizer(bad_word, add_special_tokens=False).input_ids
+        for bad_word in _build_bad_words_list_for(model_name)
+    ]
+
+    logger.info(f"Loading the {model_name} model")
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        model_name, bad_words_ids=bad_words_ids)
+    model.eval().half().to("cuda")
+
+    logger.info("Model and tokenizer are ready")
+    return model, tokenizer
+
+
+def run_raw_inference(model: transformers.AutoModelForCausalLM,
+                      tokenizer: transformers.AutoTokenizer, prompt: str,
+                      user_message: str, **kwargs: t.Any) -> str:
+    '''
+    Runs inference on the model, and attempts to returns only the newly
+    generated text.
+
+    :param model: Model to perform inference with.
+    :param tokenizer: Tokenizer to tokenize input with.
+    :param prompt: Input to feed to the model.
+    :param user_message: The user's raw message, exactly as appended to the end
+        of `prompt`. Used for trimming the original input from the model output.
+    :return: Decoded model generation.
+    '''
+    tokenized_items = tokenizer(prompt, return_tensors="pt").to("cuda")
+
+    # Atrocious code to stop generation when the model outputs "\nYou: " in
+    # freshly generated text. Feel free to send in a PR if you know of a
+    # cleaner way to do this.
+    stopping_criteria_list = transformers.StoppingCriteriaList([
+        _SentinelTokenStoppingCriteria(
+            sentinel_token_ids=tokenizer(
+                "\nYou:",
+                add_special_tokens=False,
+                return_tensors="pt",
+            ).input_ids.to("cuda"),
+            starting_idx=tokenized_items.input_ids.shape[-1])
+    ])
+
+    logits = model.generate(stopping_criteria=stopping_criteria_list,
+                            **tokenized_items,
+                            **kwargs)
+    output = tokenizer.decode(logits[0], skip_special_tokens=True)
+
+    logger.debug("Before trimming, model output was: `%s`", output)
+
+    # Trim out the input prompt from the generated output.
+    if (idx := prompt.rfind(user_message)) != -1:
+        trimmed_output = output[idx + len(user_message) - 1:].strip()
+        logger.debug("After trimming, it became: `%s`", trimmed_output)
+
+        return trimmed_output
+    else:
+        raise Exception(
+            "Couldn't find user message in the model's output. What?")
+
+
+def _build_bad_words_list_for(_model_name: str) -> t.List[str]:
+    '''Builds a list of bad words for the given model.'''
+
+    # NOTE(11b): This was implemented as a function because each model size
+    # seems to have it quirks at the moment, but this is a rushed implementation
+    # so I'm not handling that, hence the dumb return here.
+    return ["Persona:", "Scenario:", "<START>"]
+
+
+class _SentinelTokenStoppingCriteria(transformers.StoppingCriteria):
+
+    def __init__(self, sentinel_token_ids: torch.LongTensor,
+                 starting_idx: int):
+        transformers.StoppingCriteria.__init__(self)
+        self.sentinel_token_ids = sentinel_token_ids
+        self.starting_idx = starting_idx
+
+    def __call__(self, input_ids: torch.LongTensor,
+                 _scores: torch.FloatTensor) -> bool:
+        for sample in input_ids:
+            trimmed_sample = sample[self.starting_idx:]
+            # Can't unfold, output is still too tiny. Skip.
+            if trimmed_sample.shape[-1] < self.sentinel_token_ids.shape[-1]:
+                continue
+
+            for window in trimmed_sample.unfold(
+                    0, self.sentinel_token_ids.shape[-1], 1):
+                if torch.all(torch.eq(self.sentinel_token_ids, window)):
+                    return True
+        return False

parsing.py

@@ -0,0 +1,40 @@
+import re
+import typing as t
+
+def parse_messages_from_str(string: str, names: t.List[str]) -> t.List[str]:
+    '''
+    Given a big string containing raw chat history, this function attempts to
+    parse it out into a list where each item is an individual message.
+    '''
+    sanitized_names = [
+        re.escape(name) for name in names
+    ]
+
+    speaker_regex = re.compile(rf"^({'|'.join(sanitized_names)}): ?",
+                               re.MULTILINE)
+
+    message_start_indexes = []
+    for match in speaker_regex.finditer(string):
+        message_start_indexes.append(match.start())
+
+    if len(message_start_indexes) < 2:
+        # Single message in the string.
+        return [string.strip()]
+
+    prev_start_idx = message_start_indexes[0]
+    messages = []
+
+    for start_idx in message_start_indexes[1:]:
+        message = string[prev_start_idx:start_idx].strip()
+        messages.append(message)
+        prev_start_idx = start_idx
+
+    # add the last message
+    messages.append(string[prev_start_idx:].strip())
+
+    return messages
+
+
+def serialize_chat_history(history: t.List[str]) -> str:
+    '''Given a structured chat history object, collapses it down to a string.'''
+    return "\n".join(history)

prompting.py

@@ -0,0 +1,54 @@
+import logging
+import typing as t
+
+from parsing import parse_messages_from_str
+
+logger = logging.getLogger(__name__)
+
+
+def build_prompt_for(
+    history: t.List[str],
+    user_message: str,
+    char_name: str,
+    char_persona: t.Optional[str] = None,
+    example_dialogue: t.Optional[str] = None,
+    world_scenario: t.Optional[str] = None,
+) -> str:
+    '''Converts all the given stuff into a proper input prompt for the model.'''
+
+    # If example dialogue is given, parse the history out from it and append
+    # that at the beginning of the dialogue history.
+    example_history = parse_messages_from_str(
+        example_dialogue, ["You", char_name]) if example_dialogue else []
+    concatenated_history = [*example_history, *history]
+
+    # Construct the base turns with the info we already have.
+    prompt_turns = [
+        # TODO(11b): Shouldn't be here on the original 350M.
+        "<START>",
+
+        # TODO(11b): Arbitrary limit. See if it's possible to vary this
+        # based on available context size and VRAM instead.
+        *concatenated_history[-8:],
+        f"You: {user_message}",
+        f"{char_name}:",
+    ]
+
+    # If we have a scenario or the character has a persona definition, add those
+    # to the beginning of the prompt.
+    if world_scenario:
+        prompt_turns.insert(
+            0,
+            f"Scenario: {world_scenario}",
+        )
+
+    if char_persona:
+        prompt_turns.insert(
+            0,
+            f"{char_name}'s Persona: {char_persona}",
+        )
+
+    # Done!
+    logger.debug("Constructed prompt is: `%s`", prompt_turns)
+    prompt_str = "\n".join(prompt_turns)
+    return prompt_str

utils.py

@@ -0,0 +1,11 @@
+def clear_stdout():
+    '''
+    Attempts to clear stdout, whether running in a notebook (IPython) or locally
+    in a Unix envirnoment.
+    '''
+    try:
+        from IPython.display import clear_output
+        clear_output(wait=True)
+    except ImportError:
+        import os
+        os.system("clear")