Remoev

client_test.py

@@ -1,362 +0,0 @@
-"""
-Client test.
-
-Run server:
-
-python generate.py  --base_model=h2oai/h2ogpt-oig-oasst1-512-6_9b
-
-NOTE: For private models, add --use-auth_token=True
-
-NOTE: --use_gpu_id=True (default) must be used for multi-GPU in case see failures with cuda:x cuda:y mismatches.
-Currently, this will force model to be on a single GPU.
-
-Then run this client as:
-
-python src/client_test.py
-
-
-
-For HF spaces:
-
-HOST="https://h2oai-h2ogpt-chatbot.hf.space" python src/client_test.py
-
-Result:
-
-Loaded as API: https://h2oai-h2ogpt-chatbot.hf.space ✔
-{'instruction_nochat': 'Who are you?', 'iinput_nochat': '', 'response': 'I am h2oGPT, a large language model developed by LAION.', 'sources': ''}
-
-
-For demo:
-
-HOST="https://gpt.h2o.ai" python src/client_test.py
-
-Result:
-
-Loaded as API: https://gpt.h2o.ai ✔
-{'instruction_nochat': 'Who are you?', 'iinput_nochat': '', 'response': 'I am h2oGPT, a chatbot created by LAION.', 'sources': ''}
-
-NOTE: Raw output from API for nochat case is a string of a python dict and will remain so if other entries are added to dict:
-
-{'response': "I'm h2oGPT, a large language model by H2O.ai, the visionary leader in democratizing AI.", 'sources': ''}
-
-
-"""
-import ast
-import time
-import os
-import markdown  # pip install markdown
-import pytest
-from bs4 import BeautifulSoup  # pip install beautifulsoup4
-
-from enums import DocumentSubset, LangChainAction
-
-debug = False
-
-os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
-
-
-def get_client(serialize=True):
-    from gradio_client import Client
-
-    client = Client(os.getenv('HOST', "http://localhost:7860"), serialize=serialize)
-    if debug:
-        print(client.view_api(all_endpoints=True))
-    return client
-
-
-def get_args(prompt, prompt_type, chat=False, stream_output=False,
-             max_new_tokens=50,
-             top_k_docs=3,
-             langchain_mode='Disabled',
-             add_chat_history_to_context=True,
-             langchain_action=LangChainAction.QUERY.value,
-             langchain_agents=[],
-             prompt_dict=None):
-    from collections import OrderedDict
-    kwargs = OrderedDict(instruction=prompt if chat else '',  # only for chat=True
-                         iinput='',  # only for chat=True
-                         context='',
-                         # streaming output is supported, loops over and outputs each generation in streaming mode
-                         # but leave stream_output=False for simple input/output mode
-                         stream_output=stream_output,
-                         prompt_type=prompt_type,
-                         prompt_dict=prompt_dict,
-                         temperature=0.1,
-                         top_p=0.75,
-                         top_k=40,
-                         num_beams=1,
-                         max_new_tokens=max_new_tokens,
-                         min_new_tokens=0,
-                         early_stopping=False,
-                         max_time=20,
-                         repetition_penalty=1.0,
-                         num_return_sequences=1,
-                         do_sample=True,
-                         chat=chat,
-                         instruction_nochat=prompt if not chat else '',
-                         iinput_nochat='',  # only for chat=False
-                         langchain_mode=langchain_mode,
-                         add_chat_history_to_context=add_chat_history_to_context,
-                         langchain_action=langchain_action,
-                         langchain_agents=langchain_agents,
-                         top_k_docs=top_k_docs,
-                         chunk=True,
-                         chunk_size=512,
-                         document_subset=DocumentSubset.Relevant.name,
-                         document_choice=[],
-                         )
-    from evaluate_params import eval_func_param_names
-    assert len(set(eval_func_param_names).difference(set(list(kwargs.keys())))) == 0
-    if chat:
-        # add chatbot output on end.  Assumes serialize=False
-        kwargs.update(dict(chatbot=[]))
-
-    return kwargs, list(kwargs.values())
-
-
-@pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_basic(prompt_type='human_bot'):
-    return run_client_nochat(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
-
-
-def run_client_nochat(prompt, prompt_type, max_new_tokens):
-    kwargs, args = get_args(prompt, prompt_type, chat=False, max_new_tokens=max_new_tokens)
-
-    api_name = '/submit_nochat'
-    client = get_client(serialize=True)
-    res = client.predict(
-        *tuple(args),
-        api_name=api_name,
-    )
-    print("Raw client result: %s" % res, flush=True)
-    res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
-                    response=md_to_text(res))
-    print(res_dict)
-    return res_dict, client
-
-
-@pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_basic_api(prompt_type='human_bot'):
-    return run_client_nochat_api(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
-
-
-def run_client_nochat_api(prompt, prompt_type, max_new_tokens):
-    kwargs, args = get_args(prompt, prompt_type, chat=False, max_new_tokens=max_new_tokens)
-
-    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
-    client = get_client(serialize=True)
-    res = client.predict(
-        str(dict(kwargs)),
-        api_name=api_name,
-    )
-    print("Raw client result: %s" % res, flush=True)
-    res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
-                    response=md_to_text(ast.literal_eval(res)['response']),
-                    sources=ast.literal_eval(res)['sources'])
-    print(res_dict)
-    return res_dict, client
-
-
-@pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_basic_api_lean(prompt_type='human_bot'):
-    return run_client_nochat_api_lean(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
-
-
-def run_client_nochat_api_lean(prompt, prompt_type, max_new_tokens):
-    kwargs = dict(instruction_nochat=prompt)
-
-    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
-    client = get_client(serialize=True)
-    res = client.predict(
-        str(dict(kwargs)),
-        api_name=api_name,
-    )
-    print("Raw client result: %s" % res, flush=True)
-    res_dict = dict(prompt=kwargs['instruction_nochat'],
-                    response=md_to_text(ast.literal_eval(res)['response']),
-                    sources=ast.literal_eval(res)['sources'])
-    print(res_dict)
-    return res_dict, client
-
-
-@pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_basic_api_lean_morestuff(prompt_type='human_bot'):
-    return run_client_nochat_api_lean_morestuff(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
-
-
-def run_client_nochat_api_lean_morestuff(prompt, prompt_type='human_bot', max_new_tokens=512):
-    kwargs = dict(
-        instruction='',
-        iinput='',
-        context='',
-        stream_output=False,
-        prompt_type=prompt_type,
-        temperature=0.1,
-        top_p=0.75,
-        top_k=40,
-        num_beams=1,
-        max_new_tokens=256,
-        min_new_tokens=0,
-        early_stopping=False,
-        max_time=20,
-        repetition_penalty=1.0,
-        num_return_sequences=1,
-        do_sample=True,
-        chat=False,
-        instruction_nochat=prompt,
-        iinput_nochat='',
-        langchain_mode='Disabled',
-        add_chat_history_to_context=True,
-        langchain_action=LangChainAction.QUERY.value,
-        langchain_agents=[],
-        top_k_docs=4,
-        document_subset=DocumentSubset.Relevant.name,
-        document_choice=[],
-    )
-
-    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
-    client = get_client(serialize=True)
-    res = client.predict(
-        str(dict(kwargs)),
-        api_name=api_name,
-    )
-    print("Raw client result: %s" % res, flush=True)
-    res_dict = dict(prompt=kwargs['instruction_nochat'],
-                    response=md_to_text(ast.literal_eval(res)['response']),
-                    sources=ast.literal_eval(res)['sources'])
-    print(res_dict)
-    return res_dict, client
-
-
-@pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_chat(prompt_type='human_bot'):
-    return run_client_chat(prompt='Who are you?', prompt_type=prompt_type, stream_output=False, max_new_tokens=50,
-                           langchain_mode='Disabled',
-                           langchain_action=LangChainAction.QUERY.value,
-                           langchain_agents=[])
-
-
-@pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_chat_stream(prompt_type='human_bot'):
-    return run_client_chat(prompt="Tell a very long kid's story about birds.", prompt_type=prompt_type,
-                           stream_output=True, max_new_tokens=512,
-                           langchain_mode='Disabled',
-                           langchain_action=LangChainAction.QUERY.value,
-                           langchain_agents=[])
-
-
-def run_client_chat(prompt, prompt_type, stream_output, max_new_tokens,
-                    langchain_mode, langchain_action, langchain_agents,
-                    prompt_dict=None):
-    client = get_client(serialize=False)
-
-    kwargs, args = get_args(prompt, prompt_type, chat=True, stream_output=stream_output,
-                            max_new_tokens=max_new_tokens,
-                            langchain_mode=langchain_mode,
-                            langchain_action=langchain_action,
-                            langchain_agents=langchain_agents,
-                            prompt_dict=prompt_dict)
-    return run_client(client, prompt, args, kwargs)
-
-
-def run_client(client, prompt, args, kwargs, do_md_to_text=True, verbose=False):
-    assert kwargs['chat'], "Chat mode only"
-    res = client.predict(*tuple(args), api_name='/instruction')
-    args[-1] += [res[-1]]
-
-    res_dict = kwargs
-    res_dict['prompt'] = prompt
-    if not kwargs['stream_output']:
-        res = client.predict(*tuple(args), api_name='/instruction_bot')
-        res_dict['response'] = res[0][-1][1]
-        print(md_to_text(res_dict['response'], do_md_to_text=do_md_to_text))
-        return res_dict, client
-    else:
-        job = client.submit(*tuple(args), api_name='/instruction_bot')
-        res1 = ''
-        while not job.done():
-            outputs_list = job.communicator.job.outputs
-            if outputs_list:
-                res = job.communicator.job.outputs[-1]
-                res1 = res[0][-1][-1]
-                res1 = md_to_text(res1, do_md_to_text=do_md_to_text)
-                print(res1)
-            time.sleep(0.1)
-        full_outputs = job.outputs()
-        if verbose:
-            print('job.outputs: %s' % str(full_outputs))
-        # ensure get ending to avoid race
-        # -1 means last response if streaming
-        # 0 means get text_output, ignore exception_text
-        # 0 means get list within text_output that looks like [[prompt], [answer]]
-        # 1 means get bot answer, so will have last bot answer
-        res_dict['response'] = md_to_text(full_outputs[-1][0][0][1], do_md_to_text=do_md_to_text)
-        return res_dict, client
-
-
-@pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_nochat_stream(prompt_type='human_bot'):
-    return run_client_nochat_gen(prompt="Tell a very long kid's story about birds.", prompt_type=prompt_type,
-                                 stream_output=True, max_new_tokens=512,
-                                 langchain_mode='Disabled',
-                                 langchain_action=LangChainAction.QUERY.value,
-                                 langchain_agents=[])
-
-
-def run_client_nochat_gen(prompt, prompt_type, stream_output, max_new_tokens,
-                          langchain_mode, langchain_action, langchain_agents):
-    client = get_client(serialize=False)
-
-    kwargs, args = get_args(prompt, prompt_type, chat=False, stream_output=stream_output,
-                            max_new_tokens=max_new_tokens, langchain_mode=langchain_mode,
-                            langchain_action=langchain_action, langchain_agents=langchain_agents)
-    return run_client_gen(client, prompt, args, kwargs)
-
-
-def run_client_gen(client, prompt, args, kwargs, do_md_to_text=True, verbose=False):
-    res_dict = kwargs
-    res_dict['prompt'] = prompt
-    if not kwargs['stream_output']:
-        res = client.predict(str(dict(kwargs)), api_name='/submit_nochat_api')
-        res_dict['response'] = res[0]
-        print(md_to_text(res_dict['response'], do_md_to_text=do_md_to_text))
-        return res_dict, client
-    else:
-        job = client.submit(str(dict(kwargs)), api_name='/submit_nochat_api')
-        while not job.done():
-            outputs_list = job.communicator.job.outputs
-            if outputs_list:
-                res = job.communicator.job.outputs[-1]
-                res_dict = ast.literal_eval(res)
-                print('Stream: %s' % res_dict['response'])
-            time.sleep(0.1)
-        res_list = job.outputs()
-        assert len(res_list) > 0, "No response, check server"
-        res = res_list[-1]
-        res_dict = ast.literal_eval(res)
-        print('Final: %s' % res_dict['response'])
-        return res_dict, client
-
-
-def md_to_text(md, do_md_to_text=True):
-    if not do_md_to_text:
-        return md
-    assert md is not None, "Markdown is None"
-    html = markdown.markdown(md)
-    soup = BeautifulSoup(html, features='html.parser')
-    return soup.get_text()
-
-
-def run_client_many(prompt_type='human_bot'):
-    ret1, _ = test_client_chat(prompt_type=prompt_type)
-    ret2, _ = test_client_chat_stream(prompt_type=prompt_type)
-    ret3, _ = test_client_nochat_stream(prompt_type=prompt_type)
-    ret4, _ = test_client_basic(prompt_type=prompt_type)
-    ret5, _ = test_client_basic_api(prompt_type=prompt_type)
-    ret6, _ = test_client_basic_api_lean(prompt_type=prompt_type)
-    ret7, _ = test_client_basic_api_lean_morestuff(prompt_type=prompt_type)
-    return ret1, ret2, ret3, ret4, ret5, ret6, ret7
-
-
-if __name__ == '__main__':
-    run_client_many()

create_data.py

@@ -1,1809 +0,0 @@
-"""
-Dataset creation tools.
-
-Keep to-level imports clean of non-trivial imports for specific tools,
-because this file is imported for various purposes
-"""
-
-import ast
-import concurrent.futures
-import contextlib
-import hashlib
-import json
-import os
-import shutil
-import signal
-import sys
-import traceback
-from concurrent.futures import ProcessPoolExecutor
-
-import psutil
-import pytest
-import pandas as pd
-import numpy as np
-from tqdm import tqdm
-
-from utils import flatten_list, remove
-
-
-def parse_rst_file(filepath):
-    with open(filepath, 'r') as f:
-        input_data = f.read()
-    settings_overrides = {'initial_header_level': 2}
-    from docutils import core
-    document = core.publish_doctree(
-        source=input_data,
-        source_path=filepath,
-        settings_overrides=settings_overrides,
-    )
-    qa_pairs = []
-    current_section = None
-    current_question = ""
-    current_answer = ""
-    for node in document.traverse():
-        if node.__class__.__name__ == 'section':
-            current_section = ""
-        elif current_section is not None:
-            if node.__class__.__name__ == 'Text':
-                if node.astext()[-1] == "?":
-                    if current_question:
-                        qa_pairs.append((current_question, current_answer))
-                    current_question = node.astext()
-                    current_answer = ""
-                else:
-                    current_answer += node.astext()
-    if current_answer:
-        qa_pairs.append((current_question, current_answer))
-    return {k: v for k, v in qa_pairs}
-
-
-def test_scrape_dai_docs():
-    home = os.path.expanduser('~')
-    file = os.path.join(home, 'h2oai/docs/faq.rst')
-    qa_pairs = parse_rst_file(file)
-    prompt_type = 'human_bot'
-    from prompter import prompt_types
-    assert prompt_type in prompt_types
-    save_thing = [{"instruction": k, "output": v, 'prompt_type': prompt_type} for k, v in qa_pairs.items()]
-    output_file = "dai_faq.json"
-    with open(output_file, "wt") as f:
-        f.write(json.dumps(save_thing, indent=2))
-
-
-def test_scrape_dai_docs_all():
-    """
-    pytest create_data.py::test_scrape_dai_docs_all
-    """
-    import glob
-    import nltk
-    nltk.download('punkt')
-    dd = {}
-    np.random.seed(1234)
-    home = os.path.expanduser('~')
-    files = list(glob.glob(os.path.join(home, "h2oai/docs/**/*rst")))
-    np.random.shuffle(files)
-    val_count = int(0.05 * len(files))
-    train_files = files[val_count:]
-    valid_files = files[:val_count]
-    things = [
-        ("dai_docs.train.json", train_files),
-        ("dai_docs.valid.json", valid_files)
-    ]
-    for LEN in [100, 200, 500]:
-        for output_file, ff in things:
-            if output_file not in dd:
-                dd[output_file] = []
-            for f in ff:
-                with open(f) as input:
-                    blob = input.read()
-                    blob = blob.replace("~~", "")
-                    blob = blob.replace("==", "")
-                    blob = blob.replace("''", "")
-                    blob = blob.replace("--", "")
-                    blob = blob.replace("**", "")
-                    dd[output_file].extend(get_sentences(blob, length=LEN))
-    for output_file, _ in things:
-        save_thing = [{"output": k.strip(), 'prompt_type': 'plain'} for k in dd[output_file]]
-        with open(output_file, "wt") as f:
-            f.write(json.dumps(save_thing, indent=2))
-
-
-def get_sentences(blob, length):
-    """
-    break-up input text into sentences and then output list of sentences of about length in size
-    :param blob:
-    :param length:
-    :return:
-    """
-    import nltk
-    nltk.download('punkt')
-    from nltk.tokenize import sent_tokenize
-    sentences = sent_tokenize(blob)
-    my_sentences = []
-    my_string = ""
-    for sentence in sentences:
-        if len(my_string) + len(sentence) <= length:
-            if my_string:
-                my_string += " " + sentence
-            else:
-                my_string = sentence
-        else:
-            my_sentences.append(my_string)
-            my_string = ""
-    return my_sentences or [my_string]
-
-
-def setup_dai_docs(path=None, dst="working_dir_docs", from_hf=False):
-    """
-    Only supported if have access to source code or HF token for HF spaces and from_hf=True
-    :param path:
-    :param dst:
-    :param from_hf:
-    :return:
-    """
-
-    home = os.path.expanduser('~')
-
-    if from_hf:
-        # assumes
-        from huggingface_hub import hf_hub_download
-        # True for case when locally already logged in with correct token, so don't have to set key
-        token = os.getenv('HUGGINGFACE_API_TOKEN', True)
-        path_to_zip_file = hf_hub_download('h2oai/dai_docs', 'dai_docs.zip', token=token, repo_type='dataset')
-        path = 'h2oai'
-        import zipfile
-        with zipfile.ZipFile(path_to_zip_file, 'r') as zip_ref:
-            zip_ref.extractall(path)
-        path = os.path.join(path, 'docs/**/*')
-
-    if path is None:
-        if os.path.isdir(os.path.join(home, 'h2oai')):
-            path = os.path.join(home, "h2oai/docs/**/*")
-        else:
-            assert os.path.isdir(os.path.join(home, 'h2oai.superclean')), '%s does not exist' % path
-            path = os.path.join(home, "h2oai.superclean/docs/**/*")
-    import glob
-    files = list(glob.glob(path, recursive=True))
-
-    # pandoc can't find include files
-
-    remove(dst)
-    os.makedirs(dst)
-
-    # copy full tree, for absolute paths in rst
-    for fil in files:
-        if os.path.isfile(fil):
-            shutil.copy(fil, dst)
-
-    # hack for relative path
-    scorers_dir = os.path.join(dst, 'scorers')
-    makedirs(scorers_dir)
-    for fil in glob.glob(os.path.join(dst, '*.frag')):
-        shutil.copy(fil, scorers_dir)
-
-    return dst
-
-
-def rst_to_outputs(files, min_len=30, max_len=2048 // 2 - 30):
-    # account for sequence length (context window) including prompt and input and output
-
-    # os.system('pandoc -f rst -t plain ./expert_settings/nlp_settings.rst')
-    import pypandoc
-    basedir = os.path.abspath(os.getcwd())
-
-    outputs = []
-    for fil in files:
-        os.chdir(basedir)
-        os.chdir(os.path.dirname(fil))
-        fil = os.path.basename(fil)
-        print("Processing %s" % fil, flush=True)
-        # out_format can be one of: asciidoc, asciidoctor, beamer, biblatex, bibtex, commonmark, commonmark_x,
-        # context, csljson, docbook, docbook4, docbook5, docx, dokuwiki,
-        # dzslides, epub, epub2, epub3, fb2, gfm, haddock, html, html4, html5, icml,
-        # ipynb, jats, jats_archiving, jats_articleauthoring, jats_publishing, jira,
-        # json, latex, man,
-        # markdown, markdown_github, markdown_mmd, markdown_phpextra, markdown_strict,
-        # mediawiki, ms, muse, native, odt, opendocument, opml, org, pdf, plain, pptx,
-        # revealjs, rst, rtf, s5, slideous, slidy, tei, texinfo, textile, xwiki, zimwiki
-        out_format = 'plain'
-        # avoid extra new lines injected into text
-        extra_args = ['--wrap=preserve', '--resource path="%s" % dst']
-
-        plain_list = []
-        try:
-            # valid for expert settings
-            input_rst = pypandoc.convert_file(fil, 'rst')
-            input_list = input_rst.split('\n``')
-            for input_subrst in input_list:
-                input_plain = pypandoc.convert_text(input_subrst, format='rst', to='plain')
-                plain_list.append([input_plain, fil])
-        except Exception as e:
-            print("file exception: %s %s" % (fil, str(e)), flush=True)
-
-        if not plain_list:
-            # if failed to process as pieces of rst, then
-            output = pypandoc.convert_file(fil, out_format, extra_args=extra_args, format='rst')
-            outputs1 = get_sentences(output, length=max_len)
-            for oi, output in enumerate(outputs1):
-                output = output.replace('\n\n', '\n')
-                plain_list.append([output, fil])
-        outputs.extend(plain_list)
-
-    # report:
-    # [print(len(x)) for x in outputs]
-
-    # deal with blocks longer than context size (sequence length) of 2048
-    new_outputs = []
-    num_truncated = 0
-    num_orig = len(outputs)
-    for output, fil in outputs:
-        if len(output) < max_len:
-            new_outputs.append([output, fil])
-            continue
-        outputs1 = get_sentences(output, length=max_len)
-        for oi, output1 in enumerate(outputs1):
-            output1 = output1.replace('\n\n', '\n')
-            new_outputs.append([output1, fil])
-        num_truncated += 1
-    print('num_orig: %s num_truncated: %s' % (num_orig, num_truncated), flush=True)
-
-    new_outputs = [[k.strip(), fil] for k, fil in new_outputs if len(k.strip()) > min_len]
-
-    return new_outputs
-
-
-def test_scrape_dai_docs_all_pandoc():
-    """
-    pytest -s -v create_data.py::test_scrape_dai_docs_all_pandoc
-    :return:
-    """
-
-    dst = setup_dai_docs()
-
-    import glob
-    files = list(glob.glob(os.path.join(dst, '*rst'), recursive=True))
-
-    basedir = os.path.abspath(os.getcwd())
-    new_outputs = rst_to_outputs(files)
-    os.chdir(basedir)
-
-    remove(dst)
-    save_thing = [{"output": k.strip(), 'prompt_type': 'plain'} for k in new_outputs]
-    output_file = "dai_docs.train_cleaned.json"
-    with open(output_file, "wt") as f:
-        f.write(json.dumps(save_thing, indent=2))
-
-
-def test_config_to_json():
-    """
-    Needs to run from Driverless AI source directory.
-    E.g. (base) jon@gpu:~/h2oai$ pytest -s -v /data/jon/h2ogpt/create_data.py::test_config_to_json ; cp config.json /data/jon/h2ogpt/
-    :return:
-    """
-    try:
-        # Arrange
-        import json
-        from h2oaicore.systemutils import config
-        toml_list = []
-        for k, v in config.get_meta_dict().items():
-            title = (v.title + ": ") if v.title else ''
-            comment = v.comment or ''
-            if not (title or comment):
-                continue
-            toml_list.extend(
-                [
-                    {
-                        'prompt_type': 'plain',
-                        'instruction': f"<human>: What does {k} do?\n<bot>: {k.replace('_', ' ')} config.toml:  {comment or title}\n<human>:".replace(
-                            "\n", ""),
-                    },
-                    {
-                        'prompt_type': 'plain',
-                        'instruction': f"<human>: Explain {k}.\n<bot>: {k.replace('_', ' ')} config.toml:  {comment or title}\n<human>:".replace(
-                            "\n", ""),
-                    },
-                    {
-                        'prompt_type': 'plain',
-                        'instruction': f"<human>: How can I do this: {title}.\n<bot>: Set the {k.replace('_', ' ')} config.toml\n<human>:".replace(
-                            "\n", ""),
-                    } if title and comment else None,
-                    {
-                        'prompt_type': 'human_bot',
-                        'instruction': f'Explain the following expert setting for Driverless AI',
-                        'input': f"{k}",
-                        'output': f"{k.replace('_', ' ')} config.toml: {comment or title}".replace("\n", ""),
-                    },
-                    {
-                        'prompt_type': 'human_bot',
-                        'instruction': f'Explain the following expert setting for Driverless AI',
-                        'input': f"{k}",
-                        'output': f"{k.replace('_', ' ')} config.toml: {title}{comment}".replace("\n", ""),
-                    },
-                    {
-                        'prompt_type': 'human_bot',
-                        'instruction': f'Explain the following expert setting for Driverless AI',
-                        'input': f"{k.replace('_', ' ')}",
-                        'output': f"{k.replace('_', ' ')} config.toml: {title}{comment}".replace("\n", ""),
-                    },
-                    {
-                        'prompt_type': 'human_bot',
-                        'instruction': f'Explain the following expert setting for Driverless AI',
-                        'input': f"{title}",
-                        'output': f"{k.replace('_', ' ')} config.toml: {title}{comment}".replace("\n", ""),
-                    },
-                    {
-                        'prompt_type': 'human_bot',
-                        'instruction': f'Provide a short explanation of the expert setting {k}',
-                        'output': f"{k.replace('_', ' ')} config.toml: {comment or title}".replace("\n", ""),
-                    },
-                    {
-                        'prompt_type': 'human_bot',
-                        'instruction': f'Provide a detailed explanation of the expert setting {k}',
-                        'output': f"{k.replace('_', ' ')} config.toml: {title}{comment}".replace("\n", ""),
-                    },
-                ]
-            )
-        toml_list = [x for x in toml_list if x]
-        with open("config.json", "wt") as f:
-            f.write(json.dumps(toml_list, indent=2))
-    except Exception as e:
-        print("Exception: %s" % str(e), flush=True)
-
-
-def copy_tree(src, dst, follow_symlink=False):
-    makedirs(dst, exist_ok=True)
-    for (path, dirs, files) in os.walk(src, followlinks=follow_symlink):
-        new_path = path.replace(src, dst)
-        makedirs(new_path, exist_ok=True)
-        for file in files:
-            filename = os.path.join(path, file)
-            new_filename = os.path.join(new_path, file)
-            # print("%s -> %s" % (filename, new_filename))
-            try:
-                atomic_copy(filename, new_filename)
-            except FileNotFoundError:
-                pass
-
-
-def atomic_move(src, dst):
-    try:
-        shutil.move(src, dst)
-    except (shutil.Error, FileExistsError):
-        pass
-    remove(src)
-
-
-def atomic_copy(src=None, dst=None, with_permissions=True):
-    if os.path.isfile(dst):
-        return
-    import uuid
-    my_uuid = uuid.uuid4()
-    dst_tmp = dst + str(my_uuid)
-    makedirs(os.path.dirname(dst), exist_ok=True)
-    if with_permissions:
-        shutil.copy(src, dst_tmp)
-    else:
-        shutil.copyfile(src, dst_tmp)
-    atomic_move(dst_tmp, dst)
-    remove(dst_tmp)
-
-
-def makedirs(path, exist_ok=True):
-    """
-    Avoid some inefficiency in os.makedirs()
-    :param path:
-    :param exist_ok:
-    :return:
-    """
-    if os.path.isdir(path) and os.path.exists(path):
-        assert exist_ok, "Path already exists"
-        return path
-    os.makedirs(path, exist_ok=exist_ok)
-
-
-## Download from https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_unfiltered_cleaned_split.json
-## Turn into simple instruct prompt type. No context/previous conversations.
-def test_prep_instruct_vicuna():
-    from datasets import load_dataset
-    filename = 'ShareGPT_unfiltered_cleaned_split.json'
-    if not os.path.exists(filename):
-        os.system(
-            'wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/%s' % filename)
-    data = load_dataset("json", data_files={"train": filename})["train"]
-    training_rows = []
-    for i in range(data.num_rows):
-        conversations = data[i]['conversations']
-        assert isinstance(conversations, list), conversations
-        convo = ""
-        for j, conv in enumerate(conversations):
-            # Get ready for generate.py prompt_type=human_bot
-            # But train with prompt_type=plain
-            if conv['from'] == 'human':
-                FROM = '<human>: '
-            elif conv['from'] == 'gpt':
-                FROM = '<bot>: '
-            convo += f"{FROM}" + conv['value'] + "\n"
-        if convo:
-            training_rows.append(dict(input=convo))
-    with open(filename + ".generate_human_bot.train_plain.json", "wt") as f:
-        f.write(json.dumps(training_rows, indent=2))
-
-
-POSTFIX = ".generate_human_bot.train_plain.json"
-
-# https://bair.berkeley.edu/blog/2023/04/03/koala/
-OIG_DATASETS = [
-    "unified_chip2.jsonl",
-    "unified_grade_school_math_instructions.jsonl",
-    "unified_poetry_2_song.jsonl",
-    "unified_plot_screenplay_books_dialog.jsonl",
-]
-
-# hub issue: https://huggingface.co/datasets/laion/OIG/discussions/4
-ALL_OIG_DATASETS = ['unified_abstract_infill.jsonl',
-                    'unified_basic.jsonl',
-                    'unified_canadian_parliament.jsonl',
-                    'unified_chip2.jsonl',
-                    'unified_conv_finqa.jsonl',
-                    'unified_cuad.jsonl',
-                    'unified_essays.jsonl',
-                    'unified_flan.jsonl.gz',
-                    'unified_grade_school_math_instructions.jsonl',
-                    'unified_hc3_human.jsonl',
-                    'unified_image_prompts_instructions.jsonl',
-                    'unified_joke_explanations.jsonl',
-                    'unified_mathqa_flanv2_kojma_cot.jsonl',
-                    'unified_merged_code_xp3.jsonl',
-                    'unified_multi_news.jsonl',
-                    'unified_multi_sum.jsonl',
-                    'unified_ni.jsonl.gz',
-                    'unified_nq.jsonl',
-                    'unified_openai_summarize_tldr.jsonl',
-                    'unified_oscar_en_sample_dialog.jsonl',
-                    'unified_p3.jsonl.gz',
-                    'unified_plot_screenplay_books_dialog.jsonl',
-                    'unified_poetry_2_song.jsonl',
-                    'unified_poetry_instructions.jsonl',
-                    'unified_rallio_safety_and_prosocial.jsonl',
-                    'unified_rallio_soda_upgraded_2048.jsonl',
-                    'unified_soda_dialog.jsonl',
-                    'unified_sqlv1.jsonl',
-                    'unified_sqlv2.jsonl',
-                    'unified_squad_v2.jsonl',
-                    'unified_squad_v2_more_neg.jsonl',
-                    'unified_ul2_plus_oscar_en_sample_dialog.jsonl',
-                    'unified_unifiedskg_instructions.jsonl',
-                    'unified_unnatural_instructions.jsonl',
-                    'unified_xp3_sample.jsonl']
-
-useful_oig_files = ['unified_rallio_safety_and_prosocial.jsonl.parquet',
-                    'unified_chip2.jsonl.parquet',
-                    'unified_cuad.jsonl.parquet',
-                    'unified_essays.jsonl.parquet',
-                    'unified_flan.jsonl.gz.parquet',
-                    'unified_grade_school_math_instructions.jsonl.parquet',
-                    'unified_hc3_human.jsonl.parquet',
-                    'unified_mathqa_flanv2_kojma_cot.jsonl.parquet',
-                    'unified_merged_code_xp3.jsonl.parquet',
-                    'unified_multi_news.jsonl.parquet',
-                    # 'unified_multi_sum.jsonl.parquet'
-                    'unified_ni.jsonl.gz.parquet',
-                    'unified_openai_summarize_tldr.jsonl.parquet',
-                    # 'unified_oscar_en_sample_dialog.jsonl.parquet', # create text containing these N words, not specific
-                    'unified_plot_screenplay_books_dialog.jsonl.parquet',
-                    'unified_soda_dialog.jsonl.parquet',
-                    'unified_unnatural_instructions.jsonl.parquet',
-                    ]
-
-
-@pytest.mark.parametrize("filename", OIG_DATASETS)
-def test_get_small_sample_oig_data(filename):
-    if not os.path.exists(filename):
-        os.system('wget https://huggingface.co/datasets/laion/OIG/resolve/main/%s' % filename)
-    import json
-    rows = []
-    with open(filename, "r") as f:
-        for line in f.readlines():
-            row = json.loads(line)
-            rows.append(dict(input=row["text"]))
-    with open(filename + POSTFIX, "w") as f:
-        f.write(json.dumps(rows, indent=2))
-
-
-@pytest.mark.parametrize("filename", ALL_OIG_DATASETS)
-def test_download_useful_data_as_parquet(filename):
-    dest_file = filename + '.parquet'
-    if dest_file not in useful_oig_files:
-        pytest.skip('file declared not useful')
-    if not os.path.exists(filename):
-        os.system('wget https://huggingface.co/datasets/laion/OIG/resolve/main/%s' % filename)
-    if not os.path.exists(dest_file):
-        df = pd.read_json(path_or_buf=filename, lines=True)
-        df.to_parquet(dest_file, index=False)
-
-
-def test_merge_shuffle_small_sample_oig_data():
-    np.random.seed(1234)
-    rows = []
-    for filename in OIG_DATASETS:
-        with open(filename + POSTFIX, "r") as f:
-            rows.extend(json.loads(f.read()))
-    np.random.shuffle(rows)
-    with open("merged_shuffled_OIG_%s.json" % hashlib.sha256(str(OIG_DATASETS).encode()).hexdigest()[:10], "w") as f:
-        f.write(json.dumps(rows, indent=2))
-
-
-def test_join_jsons():
-    files = ['config.json'] * 1 + \
-            ['dai_docs.train_cleaned.json'] * 2 + \
-            ['dai_faq.json'] * 3
-    print(files)
-    lst = []
-    [lst.extend(json.load(open(fil, 'rt'))) for fil in files]
-    print(len(lst))
-    json.dump(lst, open("merged.json", "wt"), indent=2)
-
-
-@pytest.mark.parametrize("filename", ['Anthropic/hh-rlhf'])
-def test_make_rlhf_good_data(filename):
-    from datasets import load_dataset
-    rows = load_dataset(filename)["train"]["chosen"]
-    new_rows = []
-    for row in rows:
-        if row[:2] == "\n\n":
-            row = row[2:]
-        row = row.replace("Human: ", "<human>: ")
-        row = row.replace("Assistant: ", "<bot>: ")
-        new_rows.append(dict(input=row))
-    with open(filename.replace("/", "_") + POSTFIX, "w") as f:
-        f.write(json.dumps(new_rows, indent=2))
-
-
-def test_show_prompts():
-    files = ['config.json'] * 1 + \
-            ['dai_docs.train_cleaned.json'] * 1 + \
-            ['dai_faq.json'] * 1
-    file_points = [json.load(open(fil, 'rt')) for fil in files]
-    from prompter import generate_prompt
-    for data_points in file_points:
-        for data_point in data_points:
-            print(generate_prompt(data_point, 'plain', '', False, False, False)[0])
-
-
-def test_get_open_datasets():
-    # HF changed things so don't get raw list of all datasets, so not have to filter, but can't do negative filter
-    open_tags = ['license:Apache License 2.0',
-                 'license:mit',
-                 'license:apache',
-                 'license:apache2',
-                 'license:apache-2.0',
-                 'license:bsd',
-                 'license:bsd-2-clause',
-                 'license:bsd-3-clause',
-                 'license:bsd-3-clause-clear',
-                 'license:lgpl-2.1',
-                 'license:lgpl-3.0',
-                 'license:lgpl-lr',
-                 'license:lgpl',
-                 'license:openrail++',
-                 'license:openrail',
-                 'license:bigscience-bloom-rail-1.0',
-                 # 'license:agpl-3.0',
-                 'license:other',
-                 'license:unknown',
-                 # 'license:mpl-2.0',     # ok, but would have to include original copyright, license, source, copies in distribution
-                 # Attribution required:
-                 'license:odc-by',
-                 'license:cc-by-4.0',
-                 'license:cc-by-3.0',
-                 'license:cc-by-2.0',
-                 'license:cc-by-2.5',
-                 # 'license:cc-by-sa-4.0',  # would require same license
-                 'license:odbl',
-                 'license:pddl',
-                 'license:ms-pl',
-                 'license:zlib',
-                 ]
-    # bad license: cc-by-nc-4.0
-
-    from huggingface_hub import list_datasets
-    datasets = flatten_list([[x for x in list_datasets(filter=y)] for y in open_tags])
-    datasets += [x for x in list_datasets(author='openai')]
-    # check all:
-    all_license_tags = set(flatten_list([[y for y in x.tags if 'license' in y] for x in datasets]))
-    print(len(all_license_tags))
-    open_datasets = [x for x in datasets if any([y in x.tags for y in open_tags]) or 'license:' not in str(x.tags)]
-    print('open_datasets', len(open_datasets))
-    all_task_tags = set(flatten_list([[y for y in x.tags if 'task' in y] for x in open_datasets]))
-    print('all_task_tags', len(all_task_tags))
-    excluded_tags = ['image', 'hate', 'tabular', 'table-', 'classification', 'retrieval',
-                     'translation', 'identification', 'object', 'mask', 'to-text',
-                     'face-detection', 'audio', 'voice', 'reinforcement', 'depth-est',
-                     'forecasting', 'parsing', 'visual', 'speech', 'multiple-choice',
-                     'slot-filling', 'irds/argsme', '-scoring', 'other', 'graph-ml',
-                     'feature-extraction', 'keyword-spotting',
-                     'coreference-resolution', 'segmentation',
-                     'word-sense-disambiguation',
-                     'lemmatization']
-    task_tags = [x.replace('task_categories:', '').replace('task_ids:', '')
-                 for x in all_task_tags if not any([y in x for y in
-                                                    excluded_tags])]
-    print('task_tags', len(task_tags))
-    # str(x.tags) to catch any pattern match to anything in list
-    open_tasked_datasets = [x for x in open_datasets if
-                            any([y in str([x for x in x.tags if 'task' in x]) for y in task_tags]) and
-                            not any([y in str([x for x in x.tags if 'task' in x]) for y in excluded_tags]) or
-                            'task_categories' not in str(x.tags) and 'task_ids' not in str(x.tags)]
-    open_tasked_datasets = [x for x in open_tasked_datasets if not x.disabled]
-    open_tasked_datasets = [x for x in open_tasked_datasets if not x.gated]
-    open_tasked_datasets = [x for x in open_tasked_datasets if not x.private]
-    print('open_tasked_datasets', len(open_tasked_datasets))
-    sizes = list(set(flatten_list([[(y, x.id) for y in x.tags if 'size' in y] for x in open_tasked_datasets])))
-    languages = list(set(flatten_list([[(y, x.id) for y in x.tags if 'language:' in y] for x in open_tasked_datasets])))
-    open_english_tasked_datasets = [x for x in open_tasked_datasets if
-                                    'language:' not in str(x.tags) or
-                                    'language:en' in str(x.tags)]
-    small_open_english_tasked_datasets = [x for x in open_english_tasked_datasets if
-                                          'n<1K' in str(x.tags) or
-                                          '1K<n<10K' in str(x.tags) or
-                                          '1K0<n<100K' in str(x.tags) or
-                                          '100K<n<1M' in str(x.tags) or
-                                          'size_category' not in str(x.tags)
-                                          ]
-    # 'aeslc' : email_body, subject -> summarization?
-    # load_dataset(open_tasked_datasets[0].id).data['train'].to_pandas()
-    ids = [x.id for x in small_open_english_tasked_datasets]
-
-    # sanity checks
-    # https://bair.berkeley.edu/blog/2023/04/03/koala/
-    assert 'alespalla/chatbot_instruction_prompts' in ids
-    assert 'laion/OIG' in ids
-    assert 'openai/webgpt_comparisons' in ids
-    assert 'openai/summarize_from_feedback' in ids
-    assert 'Anthropic/hh-rlhf' in ids
-
-    # useful but not allowed for commercial purposes:
-    # https://huggingface.co/datasets/squad
-
-    print('open_english_tasked_datasets: ', ids, flush=True)
-
-    exclude_ids = ['allenai/nllb',  # translation only
-                   'hf-internal-testing/fixtures_image_utils',  # testing
-                   'allenai/c4',  # search-url
-                   'agemagician/uniref50',  # unknown
-                   'huggingface-course/documentation-images',  # images
-                   'smilegate-ai/kor_unsmile',  # korean
-                   'MohamedRashad/ChatGPT-prompts',  # ChatGPT/LearnGPT/https://www.emergentmind.com/
-                   'humarin/chatgpt-paraphrases',  # Paraphrase using ChatGPT
-                   'Jeska/vaccinchat',  # not useful
-                   'alespalla/chatbot_instruction_prompts',  # mixes alpaca
-                   'allenai/prosocial-dialog',
-                   # already exlucded, but wrongly in other datasets that say more permissive license
-                   'AlekseyKorshuk/persona-chat',  # low quality
-                   'bavard/personachat_truecased',  # low quality
-                   'adamlin/daily_dialog',  # medium quality conversations
-                   'adamlin/FewShotWoz',  # low quality
-                   'benjaminbeilharz/better_daily_dialog',  # low quality
-                   'benjaminbeilharz/daily_dialog_w_turn_templates',  # low
-                   'benjaminbeilharz/empathetic_dialogues_for_lm',  # low
-                   'GEM-submissions/GEM__bart_base_schema_guided_dialog__1645547915',  # NA
-                   'ia-bentebib/conv_ai_2_fr',  # low fr
-                   'ia-bentebib/daily_dialog_fr',  # low fr
-                   'ia-bentebib/dialog_re_fr',  # low fr
-                   'ia-bentebib/empathetic_dialogues_fr',  # low fr
-                   'roskoN/dailydialog',  # low
-                   'VadorMazer/skyrimdialogstest',  # low
-                   'bigbio/med_qa',  # med specific Q/A
-                   'biu-nlp/qa_srl2018',  # low quality Q/A
-                   'biu-nlp/qa_discourse',  # low quality Q/A
-                   'iarfmoose/qa_evaluator',  # low quality Q/A
-                   'jeopardy',  # low quality Q/A -- no reasoning
-                   'narrativeqa',  # low quality Q/A
-                   'nomic-ai/gpt4all_prompt_generations',  # bad license
-                   'nomic-ai/gpt4all_prompt_generations_with_p3',  # bad license
-                   'HuggingFaceH4/alpaca',  # bad license
-                   'tatsu-lab/alpaca',  # ToS breaking
-                   'yahma/alpaca-cleaned',  # ToS breaking
-                   'Hello-SimpleAI/HC3',  # bad license
-                   'glue',  # no reasoning QA
-                   'sahil2801/CodeAlpaca-20k',  # bad license
-                   'Short-Answer-Feedback/saf_communication_networks_english',  # long Q, medium A
-                   ]
-    small_open_english_tasked_datasets = [x for x in small_open_english_tasked_datasets if x.id not in exclude_ids]
-    # some ids clearly speech related
-    small_open_english_tasked_datasets = [x for x in small_open_english_tasked_datasets if 'speech' not in x.id]
-    # HF testing
-    small_open_english_tasked_datasets = [x for x in small_open_english_tasked_datasets if
-                                          'hf-internal-testing' not in x.id]
-    small_open_english_tasked_datasets = [x for x in small_open_english_tasked_datasets if
-                                          'chinese' not in x.id]
-
-    sorted_small_open_english_tasked_datasets = sorted([(x.downloads, x) for x in small_open_english_tasked_datasets],
-                                                       key=lambda x: x[0], reverse=True)
-
-    # NOTES:
-    # Run like pytest -s -v create_data.py::test_get_open_datasets &> getdata9.log
-    # See what needs config passed and add:
-    # grep 'load_dataset(' getdata9.log|grep -v data_id|less -S
-    # grep "pip install" getdata9.log
-    # NOTE: Some datasets have default config, but others are there.  Don't know how to access them.
-
-    """
-    https://huggingface.co/datasets/wikihow/blob/main/wikihow.py
-    https://github.com/mahnazkoupaee/WikiHow-Dataset
-    https://ucsb.box.com/s/ap23l8gafpezf4tq3wapr6u8241zz358
-    https://ucsb.app.box.com/s/ap23l8gafpezf4tq3wapr6u8241zz358
-    """
-
-    """
-    # some ambiguous or non-commercial datasets
-    https://github.com/PhoebusSi/alpaca-CoT
-    """
-
-    timeout = 3 * 60
-    # laion/OIG takes longer
-    for num_downloads, dataset in sorted_small_open_english_tasked_datasets:
-        data_id = dataset.id
-        func = do_one
-        args = (data_id, num_downloads)
-        kwargs = {}
-        with ProcessPoolExecutor(max_workers=1) as executor:
-            future = executor.submit(func, *args, **kwargs)
-            try:
-                future.result(timeout=timeout)
-            except concurrent.futures.TimeoutError:
-                print("\n\ndata_id %s timeout\n\n" % data_id, flush=True)
-            for child in psutil.Process(os.getpid()).children(recursive=True):
-                os.kill(child.pid, signal.SIGINT)
-                os.kill(child.pid, signal.SIGTERM)
-                os.kill(child.pid, signal.SIGKILL)
-
-
-def do_one(data_id, num_downloads):
-    from datasets import load_dataset
-    out_file = "data_%s.parquet" % str(data_id.replace('/', '_'))
-    if os.path.isfile(out_file) and os.path.getsize(out_file) > 1024 ** 3:
-        return
-    try:
-        print("Loading data_id %s num_downloads: %s" % (data_id, num_downloads), flush=True)
-        avail_list = None
-        try:
-            data = load_dataset(data_id, 'foobar')
-        except Exception as e:
-            if 'Available: ' in str(e):
-                avail_list = ast.literal_eval(str(e).split('Available:')[1].strip())
-            else:
-                avail_list = None
-        if avail_list is None:
-            avail_list = [None]
-        print("%s avail_list: %s" % (data_id, avail_list), flush=True)
-
-        for name in avail_list:
-            out_file = "data_%s_%s.parquet" % (str(data_id.replace('/', '_')), str(name))
-            if os.path.isfile(out_file):
-                continue
-            data = load_dataset(data_id, name)
-            column_names_dict = data.column_names
-            column_names = column_names_dict[list(column_names_dict.keys())[0]]
-            print("Processing data_id %s num_downloads: %s columns: %s" % (data_id, num_downloads, column_names),
-                  flush=True)
-            data_dict = data.data
-            col_dict = data.num_columns
-            first_col = list(col_dict.keys())[0]
-            if 'train' in data_dict:
-                df = data['train'].to_pandas()
-            else:
-                df = data[first_col].to_pandas()
-            # csv has issues with escaping chars, even for datasets I know I want
-            df.to_parquet(out_file, index=False)
-    except Exception as e:
-        t, v, tb = sys.exc_info()
-        ex = ''.join(traceback.format_exception(t, v, tb))
-        print("Exception: %s %s" % (data_id, ex), flush=True)
-
-
-def test_otherlic():
-    from huggingface_hub import list_datasets
-    lic = ['license:odc-by',
-           'license:cc-by-4.0',
-           'license:cc-by-3.0',
-           'license:cc-by-2.0',
-           'license:cc-by-2.5',
-           'license:cc-by-sa-4.0',
-           'license:odbl',
-           'license:pddl',
-           'license:ms-pl',
-           'license:zlib',
-           ]
-    datasets = flatten_list([[x for x in list_datasets(filter=y) if 'translation' not in str(x.tags)] for y in lic])
-    print(len(datasets))
-
-
-# These useful datasets are determined based upon data sample, column types, and uniqueness compared to larger datasets like Pile
-# grep columns getdata13.log|grep -v "\['image'\]"|sort|uniq|grep -v tokens|grep -v "'image'"|grep -v embedding|grep dialog
-useful = ['Dahoas/instruct-human-assistant-prompt',
-          'Dahoas/first-instruct-human-assistant-prompt',
-          'knkarthick/dialogsum',  # summary of conversation
-          'McGill-NLP/FaithDial',  # medium quality
-          'Zaid/quac_expanded',  # medium quality context + QA
-          '0-hero/OIG-small-chip2',  # medium
-          'alistvt/coqa-flat',  # QA medium
-          'AnonymousSub/MedQuAD_47441_Question_Answer_Pairs',  # QA medium
-          'Anthropic/hh-rlhf',  # high quality  # similar to Dahoas/full-hh-rlhf
-          'arjunth2001/online_privacy_qna',  # good quality QA
-          'Dahoas/instruct_helpful_preferences',  # medium quality instruct
-          'Dahoas/rl-prompt-dataset',  # medium chat
-          'Dahoas/rm-static',  # medium chat
-          'Dahoas/static-hh',  # medium chat  # HuggingFaceH4/self_instruct
-          'Dahoas/synthetic-instruct-gptj-pairwise',  # medium chat
-          'eli5',  # QA if prompt ELI5
-          'gsm8k',  # QA (various)
-          'guanaco/guanaco',  # prompt/response
-          'kastan/rlhf-qa-comparisons',  # good QA
-          'kastan/rlhf-qa-conditional-generation-v2',  # prompt answer
-          'OllieStanley/humaneval-mbpp-codegen-qa',  # code QA, but started from words, so better than other code QA
-          'OllieStanley/humaneval-mbpp-testgen-qa',  # code QA
-          'Graverman/Instruct-to-Code',  # code QA
-          'openai/summarize_from_feedback',  # summarize
-          'relbert/analogy_questions',  # analogy QA
-          'yitingxie/rlhf-reward-datasets',  # prompt, chosen, rejected.
-          'yizhongw/self_instruct',  # instruct (super natural & instruct)
-          'HuggingFaceH4/asss',  # QA, big A
-          'kastan/rlhf-qa-conditional-generation-v2',  # QA
-          'cosmos_qa',  # context QA
-          'vishal-burman/c4-faqs',  # QA but not so much reasoning, but alot of text
-          'squadshifts',  # QA from context
-          'hotpot_qa',  # QA from context
-          'adversarial_qa',  # QA from context
-          'allenai/soda',  # dialog -> narrative/summary
-          'squad_v2',  # context QA
-          'squadshifts',  # context QA
-          'dferndz/cSQuAD1',  # context QA
-          'dferndz/cSQuAD2',  # context QA
-          'din0s/msmarco-nlgen',  # context QA
-          'domenicrosati/TruthfulQA',  # common sense truthful QA -- trivia but good trivia
-          'hotpot_qa',  # context, QA
-          'HuggingFaceH4/self-instruct-eval',  # instruct QA, medium quality, some language reasoning
-          'kastan/EE_QA_for_RLHF',  # context QA
-          'KK04/LogicInference_OA',  # instruction logical QA
-          'lmqg/qa_squadshifts_synthetic',  # context QA
-          'lmqg/qg_squad',  # context QA
-          'lmqg/qg_squadshifts',  # context QA
-          'lmqg/qg_subjqa',  # context QA
-          'pszemraj/HC3-textgen-qa',
-          # QA medium, has human responses -- humans tend to provide links instead of trying to answer
-          'pythonist/newdata',  # long context, QA, brief A
-          'ropes',  # long background, situation, question, A
-          'wikitablequestions',  # table -> QA
-          'bigscience/p3',  # context QA but short answers
-          ]
-
-code_useful = ['0n1xus/codexglue',
-               'openai_humaneval',
-               'koutch/staqc',
-               ]
-
-maybe_useful = ['AlekseyKorshuk/comedy-scripts',
-                'openbookqa',  # hard to parse, low reasoning
-                'qed',  # reasonable QA, but low reasoning
-                'selqa',  # candidate answers
-                'HuggingFaceH4/instruction-pilot-outputs-filtered',
-                'GBaker/MedQA-USMLE-4-options',  # medical QA with long questions
-                'npc-engine/light-batch-summarize-dialogue',  # dialog summarize, kinda low specific quality
-                ]
-
-summary_useful = ['austin/rheum_abstracts',
-                  'CarperAI/openai_summarize_comparisons',  # summarize chosen/rejected
-                  'CarperAI/openai_summarize_tldr',  # summarize QA
-                  'ccdv/cnn_dailymail',  # summarize news
-                  'ccdv/govreport-summarization',  # summarize high quality
-                  'ccdv/pubmed-summarization',  # summarize high quality
-                  'duorc',  # plot -> QA
-                  'farleyknight/big_patent_5_percent',  # desc -> abstract
-                  'multi_news',  # summary
-                  'opinosis',
-                  'SophieTr/reddit_clean',
-                  'allenai/mup',  # long text -> summary
-                  'allenai/multi_lexsum',  # long text -> summary
-                  'big_patent',
-                  'allenai/wcep_dense_max',
-                  'awinml/costco_long_practice',
-                  'GEM/xsum',
-                  'ratishsp/newshead',
-                  'RussianNLP/wikiomnia',  # russian
-                  'stacked-summaries/stacked-xsum-1024',
-                  ]
-
-math_useful = [
-    'competition_math'
-]
-
-skipped = ['c4',  # maybe useful, used for flan, but skipped due to size
-           ]
-
-"""
-To get training data from oig:
-pytest test_oig test_grade_final test_finalize_to_json
-"""
-
-human = '<human>:'
-bot = '<bot>:'
-
-
-def test_assemble_and_detox():
-    import re
-    from profanity_check import predict_prob
-    df_list = []
-    for data in useful_oig_files:
-        print("Processing %s" % data, flush=True)
-        df = pd.read_parquet(data)
-        df = df.reset_index(drop=True)
-        # chop up into human/bot interactions of no more than 10kB per row
-        text_list = df[['text']].values.ravel().tolist()
-        new_text = []
-        max_len = 2048  # uber cutoff
-        MAX_LEN = 2048 // 2 - 30  # max len per question/answer
-        for text in tqdm(text_list):
-            human_starts = [m.start() for m in re.finditer('<human>: ', text)]
-            if len(human_starts) == 1:
-                human_starts = [0, len(text)]  # always go into for loop below
-            blurb = ''
-            for i in range(len(human_starts) - 1):
-                interaction = text[human_starts[i]: human_starts[i + 1]][:max_len]
-                blurb += interaction
-                if len(blurb) >= MAX_LEN:
-                    blurb = get_sentences(blurb, length=MAX_LEN)[0]
-                    new_text.append(blurb + "\n<human>:")
-                    blurb = ''
-            if blurb:
-                blurb = get_sentences(blurb, length=MAX_LEN)[0]
-                new_text.append(blurb + "\n<human>:")
-
-        if len(new_text) > len(text_list):
-            print("Added %d new rows (before: %d)" % (len(new_text) - df.shape[0], df.shape[0]))
-        df = pd.DataFrame({"text": new_text, "source": [data] * len(new_text)})
-        df = df.drop_duplicates(keep='first')
-        print(df['text'].apply(lambda x: len(x)).describe())
-        assert df['text'].apply(lambda x: len(x)).max() <= 2 * max_len
-
-        # faster than better_profanity, do early
-        df['profanity'] = predict_prob(df['text'])
-        before_rows = df.shape[0]
-        df = df[df['profanity'] < 0.25]  # drop any low quality stuff
-        after_rows = df.shape[0]
-        print("Dropped %d rows out of %d due to alt-profanity-check" % (before_rows - after_rows, before_rows))
-        df_list.append(df)
-        print("Done processing %s -> %s rows" % (data, df.shape[0]), flush=True)
-        print("So far have %d rows" % sum([len(x) for x in df_list]))
-    df_final = pd.concat(df_list)
-    df_final = df_final.sample(frac=1, random_state=1234).reset_index(drop=True)
-    df_final.to_parquet('h2oGPT.cleaned.human_bot.shorter.parquet', index=False)
-
-
-def test_basic_cleaning():
-    # from better_profanity import profanity
-    # https://pypi.org/project/alt-profanity-check/
-    from profanity_check import predict
-    df_list = []
-    for data in useful_oig_files:
-        # for data in useful_oig_files[:5]:
-        # for data in ['unified_openai_summarize_tldr.jsonl.parquet']:
-        print("Processing %s" % data, flush=True)
-        df = pd.read_parquet(data)
-        df = df.reset_index(drop=True)
-        # NOTE: Not correct if multiple human-bot interactions, but those dialogs even more desired
-        # avg_chars = len(df['text'][0])/(df['text'][0].count(human)+df['text'][0].count(bot))
-        df['avg_words'] = df['text'].apply(lambda x: x.count(' ') / (x.count(human) + x.count(bot)) / 2.0)
-        df['avg_bot_words'] = df['text'].apply(lambda x: x.split(bot)[1].count(' ') / x.count(bot))
-        # df['bad_words'] = df['text'].apply(lambda x: profanity.contains_profanity(x))
-        # low_quality_patterns = ['Write the rest of this wikipedia article']
-        res = predict(df['text'])
-        df['bad_words'] = res
-        df = df.reset_index(drop=True)
-        df = df[df['bad_words'] == 0]
-        df = df[['text', 'avg_words', 'avg_bot_words']]
-        df = df.drop_duplicates(keep='first')
-        print(df[df['avg_words'] == df['avg_words'].max()]['text'].values)
-        median_words = np.median(df['avg_words'])
-        min_words_per_entity = max(30, 0.8 * median_words)
-        max_words_per_entity = 2048  # too hard to learn from for now
-        df = df[df['avg_words'] > min_words_per_entity]
-        df = df[df['avg_words'] < max_words_per_entity]
-
-        min_words_per_entity = max(20, 0.5 * median_words)  # bot should say stuff for now
-        max_words_per_entity = 2048  # too hard to learn from for now
-        df = df[df['avg_bot_words'] > min_words_per_entity]
-        df = df[df['avg_bot_words'] < max_words_per_entity]
-
-        df_list.append(df)
-        print("Done processing %s -> %s rows" % (data, df.shape[0]), flush=True)
-    df_final = pd.concat(df_list)
-    df_final.to_parquet('h2oGPT.cleaned.human_bot.parquet', index=False)
-
-
-from joblib import Parallel, delayed, effective_n_jobs
-from sklearn.utils import gen_even_slices
-from sklearn.utils.validation import _num_samples
-
-
-def parallel_apply(df, func, n_jobs=-1, **kwargs):
-    """ Pandas apply in parallel using joblib.
-    Uses sklearn.utils to partition input evenly.
-
-    Args:
-        df: Pandas DataFrame, Series, or any other object that supports slicing and apply.
-        func: Callable to apply
-        n_jobs: Desired number of workers. Default value -1 means use all available cores.
-        **kwargs: Any additional parameters will be supplied to the apply function
-
-    Returns:
-        Same as for normal Pandas DataFrame.apply()
-
-    """
-
-    if effective_n_jobs(n_jobs) == 1:
-        return df.apply(func, **kwargs)
-    else:
-        ret = Parallel(n_jobs=n_jobs)(
-            delayed(type(df).apply)(df[s], func, **kwargs)
-            for s in gen_even_slices(_num_samples(df), effective_n_jobs(n_jobs)))
-        return pd.concat(ret)
-
-
-def add_better_profanity_flag(df):
-    from better_profanity import profanity
-    df['better_profanity'] = parallel_apply(
-        df['text'],
-        lambda x: profanity.contains_profanity(x),
-        n_jobs=-1,
-    )
-    return df
-
-
-def add_textstat_grade(df):
-    import textstat
-
-    def myfunc(x):
-        return textstat.flesch_kincaid_grade(x)  # simple grade
-
-    if False:
-        import dask.dataframe as dd
-        # 40 seconds for 1000 rows, but have 1,787,799 rows
-        ddata = dd.from_pandas(df, npartitions=120)
-
-        df['flesch_grade'] = ddata['text'].apply(myfunc).compute()
-    if True:
-        # fast way
-        df['flesch_grade'] = parallel_apply(df['text'], myfunc, n_jobs=-1)
-    return df
-
-
-def add_deberta_grade(df):
-    from transformers import AutoModelForSequenceClassification, AutoTokenizer
-    import torch
-    reward_name = "OpenAssistant/reward-model-deberta-v3-large-v2"
-    rank_model, tokenizer = AutoModelForSequenceClassification.from_pretrained(
-        reward_name), AutoTokenizer.from_pretrained(reward_name)
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    rank_model.to(device)
-
-    def get_question(x):
-        return x.replace('<human>: ', '').split('<bot>:')[0]
-
-    def get_answer(x):
-        try:
-            answer = x.split('<bot>: ')[1].split('<human>:')[0].replace('<bot>: ', '')
-        except:
-            answer = x.split('<bot>:')[1].split('<human>:')[0].replace('<bot>:', '')
-        return answer
-
-    df['question'] = parallel_apply(df['text'], get_question, n_jobs=-1)
-    df['answer'] = parallel_apply(df['text'], get_answer, n_jobs=-1)
-
-    from datasets import Dataset
-    from transformers import pipeline
-    from transformers.pipelines.pt_utils import KeyPairDataset
-    import tqdm
-
-    pipe = pipeline(
-        "text-classification",
-        model=reward_name,
-        device="cuda:0" if torch.cuda.is_available() else "cpu"
-    )
-    start = 0
-    batch_size = 64 * 16
-    micro_batch = orig_micro_batch = 16
-    end = 0
-    import socket
-    checkpoint = "grades.%s.pkl" % socket.gethostname()
-    grades = []
-    import pickle
-    if os.path.exists(checkpoint):
-        with open(checkpoint, "rb") as f:
-            start, grades = pickle.loads(f.read())
-    last_oom = 0
-    while end < df.shape[0]:
-        # manual batching to handle OOM more gracefully
-        end = min(start + batch_size, df.shape[0])
-        if start == end:
-            break
-        dataset = Dataset.from_pandas(df.iloc[start:end, :])
-        try:
-            grades.extend([
-                x['score'] for x in tqdm.tqdm(
-                    pipe(KeyPairDataset(dataset, "question", "answer"), batch_size=micro_batch)
-                )
-            ])
-        except torch.cuda.OutOfMemoryError:
-            last_oom = start
-            micro_batch = max(1, micro_batch // 2)
-            print("OOM - retrying with micro_batch=%d" % micro_batch)
-            continue
-        if last_oom == start:
-            micro_batch = orig_micro_batch
-            print("Returning to micro_batch=%d" % micro_batch)
-        assert len(grades) == end
-        start = end
-        with open(checkpoint, "wb") as f:
-            f.write(pickle.dumps((end, grades)))
-        print("%d/%d" % (end, df.shape[0]))
-    df['grade_deberta'] = grades
-    if os.path.exists(checkpoint):
-        os.remove(checkpoint)
-    return df
-
-
-def test_chop_by_lengths():
-    file = "h2oGPT.cleaned.human_bot.shorter.parquet"
-    df = pd.read_parquet(file).reset_index(drop=True)
-    df = count_human_bot_lengths(df)
-    df['rand'] = np.random.rand(df.shape[0])
-    df['rand2'] = np.random.rand(df.shape[0])
-    before_rows = df.shape[0]
-    # throw away short human/bot responses with higher likelihood
-    df = df[(df['len_human_mean'] > 20)]  # never keep very short ones
-    df = df[(df['len_human_mean'] > 30) | (df['rand'] < 0.2)]
-    df = df[(df['len_human_mean'] > 50) | (df['rand'] < 0.5)]
-    df = df[(df['len_human_max'] < 10000)]  # drop super long (basically only human) ones
-    df = df[(df['len_bot_mean'] > 20)]  # never keep very short ones
-    df = df[(df['len_bot_mean'] > 30) | (df['rand2'] < 0.2)]
-    df = df[(df['len_bot_mean'] > 50) | (df['rand2'] < 0.5)]
-    df = df[(df['len_bot_max'] < 10000)]  # drop super long (only bot) ones
-    assert df['text'].apply(lambda x: len(x)).max() < 20000
-    df = df.drop(['rand', 'rand2'], axis=1)
-    after_rows = df.shape[0]
-    print("Chopped off %d out of %d rows due to length" % (before_rows - after_rows, before_rows))
-    print(df.describe())
-    df.to_parquet('h2oGPT.cleaned.chopped.human_bot.shorter.parquet', index=False)
-
-
-def count_human_bot_lengths(df, human=None, bot=None):
-    import re
-    len_human_min = []
-    len_human_max = []
-    len_human_mean = []
-    len_bot_min = []
-    len_bot_max = []
-    len_bot_mean = []
-    human = human or '<human>:'
-    bot = bot or '<bot>:'
-    for is_human in [True, False]:
-        what = human if is_human else bot
-        other = human if not is_human else bot
-        for i in range(df.shape[0]):
-            text = df.loc[i, 'text']
-            assert isinstance(text, str)
-            starts = [m.start() for m in re.finditer(what, text)]
-            if len(starts) == 1:
-                starts = [starts[0], len(text)]  # always go into for loop below
-            assert len(text)
-            list_what = []
-            for ii in range(len(starts) - 1):
-                interaction = text[starts[ii]: starts[ii + 1]]
-                if other in interaction:
-                    interaction = interaction[:interaction.find(other)]
-                interaction.strip()
-                list_what.append(interaction)
-            if not list_what:
-                list_what = ['']  # handle corrupted data, very rare, leads to sizes 0
-            if is_human:
-                len_human_min.append(min([len(x) for x in list_what]))
-                len_human_max.append(max([len(x) for x in list_what]))
-                len_human_mean.append(np.mean([len(x) for x in list_what]))
-            else:
-                len_bot_min.append(min([len(x) for x in list_what]))
-                len_bot_max.append(max([len(x) for x in list_what]))
-                len_bot_mean.append(np.mean([len(x) for x in list_what]))
-    df['len_human_min'] = len_human_min
-    df['len_human_max'] = len_human_max
-    df['len_human_mean'] = len_human_mean
-    df['len_bot_min'] = len_bot_min
-    df['len_bot_max'] = len_bot_max
-    df['len_bot_mean'] = len_bot_mean
-    np.random.seed(1234)
-    pd.set_option('display.max_columns', None)
-    print("Before chopping")
-    print(df.describe())
-    return df
-
-
-def test_grade():
-    df = None
-
-    file = "h2oGPT.cleaned.chopped.human_bot.shorter.parquet"
-    output_file = "h2oGPT.cleaned.graded1.human_bot.shorter.parquet"
-    if not os.path.exists(output_file):
-        if df is None:
-            df = pd.read_parquet(file).reset_index(drop=True)
-        df = add_textstat_grade(df)
-        min_grade = 10
-        max_grade = 25
-        df = df[df['flesch_grade'] >= min_grade]
-        df = df[df['flesch_grade'] <= max_grade]
-        print("After Flesch grade")
-        print(df.describe())
-        df.to_parquet(output_file, index=False)
-
-    file = output_file
-    output_file = "h2oGPT.cleaned.graded2.human_bot.shorter.parquet"
-    if not os.path.exists(output_file):
-        # slower than alt-profanity, do last, but do before deberta grading, since that's slower
-        if df is None:
-            df = pd.read_parquet(file).reset_index(drop=True)
-        df = add_better_profanity_flag(df)
-        before_rows = df.shape[0]
-        df = df[df['better_profanity'] == 0]
-        df = df.drop(['better_profanity'], axis=1)
-        after_rows = df.shape[0]
-        print("Dropped %d rows out of %d due to better_profanity" % (before_rows - after_rows, before_rows))
-        print(df.describe())
-        df.to_parquet(output_file, index=False)
-
-    file = output_file
-    output_file = 'h2oGPT.cleaned.graded3.human_bot.shorter.parquet'
-    if not os.path.exists(output_file):
-        if df is None:
-            df = pd.read_parquet(file).reset_index(drop=True)
-        df = add_deberta_grade(df)
-        min_grade = 0.3
-        max_grade = np.inf
-        before_rows = df.shape[0]
-        df = df[df['grade_deberta'] >= min_grade]
-        df = df[df['grade_deberta'] <= max_grade]
-        after_rows = df.shape[0]
-        print("Dropped %d rows out of %d due to deberta grade" % (before_rows - after_rows, before_rows))
-        print("After DeBERTa grade")
-        print(df.describe())
-        df.to_parquet(output_file, index=False)
-
-    file = output_file
-    output_file = 'h2oGPT.cleaned.graded.human_bot.shorter.parquet'
-    if df is None:
-        df = pd.read_parquet(file).reset_index(drop=True)
-    df.to_parquet(output_file, index=False)
-
-
-@pytest.mark.parametrize(
-    "fixup_personality, only_personality, deberta_grading",
-    [
-        [False, False, False],
-        [True, True, False],
-        [True, False, False],
-        [True, False, True],
-    ]
-)
-def test_add_open_assistant(fixup_personality, only_personality, deberta_grading, save_json=True):
-    """
-    Flatten tree structure into one row per path from root to leaf
-    Also turn into human_bot prompting format:
-        <human>: question\n<bot>: answer <human>: question2\n<bot>: answer2 Etc.
-    Also saves a .json locally as side-effect
-    returns list of dicts, containing intput, prompt_type and source
-    """
-    from datasets import load_dataset
-    data_file = "OpenAssistant/oasst1"
-    ds = load_dataset(data_file)
-    df = pd.concat([ds['train'].to_pandas(), ds['validation'].to_pandas()], axis=0)
-    rows = {}
-    message_ids = df['message_id'].values.tolist()
-    message_tree_ids = df['message_tree_id'].values.tolist()
-    parent_ids = df['parent_id'].values.tolist()
-    texts = df['text'].values.tolist()
-    roles = df['role'].values.tolist()
-
-    for i in range(df.shape[0]):
-        # collect all trees
-        message_id = message_ids[i]
-        message_tree_id = message_tree_ids[i]
-        parent_id = parent_ids[i]
-        text = texts[i]
-        if fixup_personality:
-            text = text.replace("Open Assistant", "h2oGPT")
-            text = text.replace("Open-Assistant", "h2oGPT")
-            text = text.replace("open-assistant", "h2oGPT")
-            text = text.replace("OpenAssistant", "h2oGPT")
-            text = text.replace("open assistant", "h2oGPT")
-            text = text.replace("Open Assistand", "h2oGPT")
-            text = text.replace("Open Assitant", "h2oGPT")
-            text = text.replace("Open Assistent", "h2oGPT")
-            text = text.replace("Open Assisstant", "h2oGPT")
-            text = text.replace("Open Assitent", "h2oGPT")
-            text = text.replace("Open Assitiant", "h2oGPT")
-            text = text.replace("Open Assistiant", "h2oGPT")
-            text = text.replace("Open Assitan ", "h2oGPT ")
-            text = text.replace("Open Assistan ", "h2oGPT ")
-            text = text.replace("Open Asistant", "h2oGPT")
-            text = text.replace("Open Assiant", "h2oGPT")
-            text = text.replace("Assistant", "h2oGPT")
-            text = text.replace("LAION AI", "H2O.ai")
-            text = text.replace("LAION-AI", "H2O.ai")
-            text = text.replace("LAION,", "H2O.ai,")
-            text = text.replace("LAION.ai", "H2O.ai")
-            text = text.replace("LAION.", "H2O.ai.")
-            text = text.replace("LAION", "H2O.ai")
-
-        role = roles[i]
-        new_data = ('<human>: ' if role == 'prompter' else '<bot>: ') + text
-        entry = dict(message_id=message_id, parent_id=parent_id, text=new_data)
-        if message_tree_id not in rows:
-            rows[message_tree_id] = [entry]
-        else:
-            rows[message_tree_id].append(entry)
-
-    all_rows = []
-
-    for node_id in rows:
-        # order responses in tree, based on message/parent relationship
-        conversations = []
-
-        list_msgs = rows[node_id]
-        # find start
-        while len(list_msgs):
-            for i, leaf in enumerate(list_msgs):
-                found = False
-                parent_id = leaf['parent_id']
-                if parent_id is None:
-                    # conversation starter
-                    conversations.append(leaf)
-                    found = True
-                else:
-                    for conv in conversations:
-                        # find all conversations to add my message to
-                        if parent_id in conv['message_id'] and parent_id != conv['message_id'][-len(parent_id):]:
-                            # my message doesn't follow conversation
-                            continue
-                        if parent_id == conv['message_id'][-len(parent_id):]:
-                            # my message follows conversation, but fork first, so another follow-on message can do same
-                            conversations.append(conv.copy())
-                            conv['text'] += f"""
-{leaf['text']}
-"""
-                            conv['message_id'] += leaf['message_id']
-                            found = True
-                            break
-                if found:
-                    # my content was used, so nuke from list
-                    del list_msgs[i]
-                    break
-
-        # now reduce down to final conversations, find the longest chains of message ids
-        for i, conv in enumerate(conversations):
-            for j, conv2 in enumerate(conversations):
-                if i == j:
-                    continue
-                if conv['message_id'] and conv2['message_id']:
-                    assert conv['message_id'] != conv2['message_id']
-                    # delete the shorter conversation, if one contains the other
-                    if conv['message_id'] in conv2['message_id']:
-                        conv['message_id'] = None
-                    if conv2['message_id'] in conv['message_id']:
-                        conv2['message_id'] = None
-        conversations = [c for c in conversations if c['message_id']]
-        if only_personality:
-            all_rows.extend(
-                [dict(input=c['text'] + "\n<human>:", prompt_type='plain', source=data_file) for c in conversations if
-                 'h2oGPT' in c['text']])
-        else:
-            all_rows.extend(
-                [dict(input=c['text'] + "\n<human>:", prompt_type='plain', source=data_file) for c in conversations if
-                 "What is H2O.ai" not in c['text']])
-    unhelpful = get_unhelpful_list()
-    all_rows = [x for x in all_rows if not any(u in x['input'] for u in unhelpful)]
-    personality = create_personality_data()
-    all_rows.extend(personality * 10)
-    np.random.seed(123)
-    np.random.shuffle(all_rows)
-    print(len(all_rows))
-    if deberta_grading:
-        df = pd.DataFrame(all_rows)
-        df = df.rename(columns={'input': 'text'})
-        df = add_deberta_grade(df)
-        df = df.rename(columns={'text': 'input'})
-        drop = True
-        if drop:
-            min_grade = 0.3
-            max_grade = np.inf
-            before_rows = df.shape[0]
-            df = df[df['grade_deberta'] >= min_grade]
-            df = df[df['grade_deberta'] <= max_grade]
-            after_rows = df.shape[0]
-            print("Dropped %d rows out of %d due to deberta grade" % (before_rows - after_rows, before_rows))
-            print("After DeBERTa grade")
-        print(df.describe())
-        all_rows = []
-        for i in range(df.shape[0]):
-            all_rows.append(
-                dict(
-                    input=df['input'].iloc[i],
-                    source=df['source'].iloc[i],
-                    prompt_type=df['prompt_type'].iloc[i],
-                    grade_deberta=df['grade_deberta'].iloc[i],
-                )
-            )
-    if save_json:
-        data_file = data_file + \
-                    ("_h2ogpt" if fixup_personality else "") + \
-                    ("_only" if only_personality else "") + \
-                    ("_graded" if deberta_grading else "")
-        for i in range(len(all_rows)):
-            all_rows[i]['id'] = i
-        with open(data_file.lower().replace("/", "_") + ".json", "w") as f:
-            f.write(json.dumps(all_rows, indent=2))
-    return all_rows
-
-
-def test_finalize_to_json():
-    df = pd.read_parquet('h2oGPT.cleaned.graded.human_bot.shorter.parquet')
-    df = df.rename(columns={'text': 'input'})
-
-    print("Number of high-quality human_bot interactions: %s" % df.shape[0], flush=True)
-
-    print("Adding open assistant data")
-    with open("openassistant_oasst1_h2ogpt_graded.json") as f:
-        open_assistant = json.loads(f.read())
-    df = pd.concat([df, pd.DataFrame(open_assistant)], axis=0)
-
-    def final_clean(df):
-        from better_profanity import profanity
-        profanity.load_censor_words_from_file("data/censor_words.txt")
-        df['profanity'] = parallel_apply(
-            df['input'],
-            lambda x: profanity.contains_profanity(x),
-            n_jobs=-1,
-        )
-        return df[(df['profanity'] == 0)].reset_index(drop=True)
-
-    print("Before cleaning: Number of final high-quality human_bot interactions: %s" % df.shape[0], flush=True)
-    df = final_clean(df)
-    print("After cleaning: Number of final high-quality human_bot interactions: %s" % df.shape[0], flush=True)
-    print(df.describe())
-    print(df.shape)
-    row_list = []
-    for i in range(df.shape[0]):
-        row_list.append(
-            dict(
-                input=df.loc[i, 'input'],
-                source=df.loc[i, 'source'],
-                prompt_type='plain',
-            )
-        )
-    np.random.seed(1234)
-    np.random.shuffle(row_list)
-    unhelpful = get_unhelpful_list()
-    row_list = [x for x in row_list if not any(u in x['input'] for u in unhelpful)]
-    for i in range(len(row_list)):
-        row_list[i]['id'] = i
-        row_list[i]['input'] = row_list[i]['input'].replace(" <bot>:", "\n<bot>:")
-    with open('h2ogpt-oig-oasst1-instruct-cleaned-v3.json', "w") as f:
-        f.write(json.dumps(row_list, indent=2))
-
-
-def create_personality_data():
-    questions = [
-        "What's your name?",
-        "What is your name?",
-        "What are you?",
-        "Who are you?",
-        "Do you have a name?",
-        "Who trained you?",
-        "Who created you?",
-        "Who made you?",
-    ]
-    answers = [
-        "I'm h2oGPT, a large language model by H2O.ai.",
-        "I'm h2oGPT, a large language model by H2O.ai, the visionary leader in democratizing AI.",
-        "My name is h2oGPT. I'm a large language model by H2O.ai, the visionary leader in democratizing AI.",
-        "My name is h2oGPT. I'm a large language model trained by H2O.ai.",
-        "Hi! I'm h2oGPT, a large language model by H2O.ai.",
-        "Hi! I'm h2oGPT, a large language model by H2O.ai, the visionary leader in democratizing AI.",
-    ]
-    help = [
-        "",
-        " How can I help you?",
-        " How may I assist you?",
-        " Nice to meet you.",
-    ]
-    import itertools
-    rows = []
-    for pair in itertools.product(questions, answers, help):
-        rows.append(
-            dict(input=f"<human>: {pair[0]}\n<bot>: {pair[1]}{pair[2]}\n<human>:", prompt_type='plain', source="H2O.ai")
-        )
-    for row in [
-        "<human>: What is H2O.ai?\n<bot>: H2O.ai is a technology company that aims to democratize AI and make it accessible to a broader audience by simplifying the process of creating and deploying machine learning models.\n<human>:",
-        "<human>: What is h2o.ai?\n<bot>: H2O.ai is a technology company that aims to democratize AI and make it accessible to a broader audience by simplifying the process of creating and deploying machine learning models.\n<human>:",
-        "<human>: What is H2O?\n<bot>: H2O.ai is a technology company that aims to democratize AI and make it accessible to a broader audience by simplifying the process of creating and deploying machine learning models.\n<human>:",
-        "<human>: Who is h2o.ai?\n<bot>: H2O.ai is a technology company that aims to democratize AI and make it accessible to a broader audience by simplifying the process of creating and deploying machine learning models.\n<human>:",
-        "<human>: who is h2o.ai?\n<bot>: H2O.ai is a technology company that aims to democratize AI and make it accessible to a broader audience by simplifying the process of creating and deploying machine learning models.\n<human>:",
-        "<human>: who is h2o?\n<bot>: H2O.ai is a technology company that aims to democratize AI and make it accessible to a broader audience by simplifying the process of creating and deploying machine learning models.\n<human>:",
-        "<human>: What is H2O.ai?\n<bot>: H2O.ai is the visionary leader in democratizing AI.\n<human>:",
-        "<human>: Who is H2O.ai?\n<bot>: H2O.ai is the visionary leader in democratizing AI.\n<human>:",
-        "<human>: Who is H2O?\n<bot>: H2O.ai is the visionary leader in democratizing AI.\n<human>:",
-        "<human>: Who is h2o?\n<bot>: H2O.ai is the visionary leader in democratizing AI.\n<human>:",
-        "<human>: who is h2o?\n<bot>: H2O.ai is the visionary leader in democratizing AI.\n<human>:",
-    ]:
-        rows.append(dict(input=row, prompt_type='plain', source='H2O.ai'))
-    print(len(rows))
-    with open("h2ogpt-personality.json", "w") as f:
-        f.write(json.dumps(rows, indent=2))
-    return rows
-
-
-def test_check_stats_data():
-    filename = 'h2ogpt-oig-oasst1-instruct-cleaned-v3.json'
-    df = pd.read_json(filename)
-
-    # get word stats
-    df['char_count'] = df['input'].apply(lambda x: len(x))
-    import matplotlib.pyplot as plt
-    plt.figure(figsize=(10, 10))
-    plt.hist(df['char_count'], bins=100)
-    chars_avg = np.mean(df['char_count'])
-    chars_median = np.median(df['char_count'])
-    plt.title("char_count avg: %s median: %s" % (chars_avg, chars_median))
-    plt.savefig('chars_hist.png')
-    plt.close()
-
-    # get tokenize stats for random sample of 1000 rows
-    from finetune import generate_and_tokenize_prompt
-    from loaders import get_loaders, get_tokenizer
-    from functools import partial
-
-    llama_type = False
-    tokenizer_base_model = base_model = 'h2oai/h2ogpt-oasst1-512-20b'
-    model_loader, tokenizer_loader = get_loaders(model_name=base_model, reward_type=False, llama_type=llama_type)
-    local_files_only = False
-    resume_download = True
-    use_auth_token = False
-    tokenizer = get_tokenizer(tokenizer_loader, tokenizer_base_model, local_files_only, resume_download, use_auth_token)
-    prompt_type = 'plain'  # trained with data already in human bot form
-    train_on_inputs = True
-    add_eos_token = False
-    cutoff_len = 512  # can choose 2048
-    generate_and_tokenize_prompt_fun = partial(generate_and_tokenize_prompt, prompt_type=prompt_type,
-                                               train_on_inputs=train_on_inputs, add_eos_token=add_eos_token,
-                                               cutoff_len=cutoff_len, tokenizer=tokenizer)
-    from datasets import load_dataset
-    data = load_dataset("json", data_files={"train": filename})
-    val_set_size = 0.90
-    train_val = data["train"].train_test_split(
-        test_size=val_set_size, shuffle=True, seed=42
-    )
-    train_data = train_val["train"]
-    train_data = train_data.shuffle().map(generate_and_tokenize_prompt_fun, num_proc=os.cpu_count())
-
-    df_tokens = pd.DataFrame([len(x) for x in train_data['input_ids']], columns=['token_count'])
-
-    plt.figure(figsize=(10, 10))
-    plt.hist(df_tokens['token_count'], bins=100)
-    token_avg = np.mean(df_tokens['token_count'])
-    token_median = np.median(df_tokens['token_count'])
-    plt.title("token_count with cutoff=%s avg: %s median: %s" % (cutoff_len, token_avg, token_median))
-    plt.savefig('token_hist_%s.png' % cutoff_len)
-    plt.close()
-
-
-def get_unhelpful_list():
-    # base versions
-    unhelpful = ["I'm sorry, I didn't quite understand your question, could you please rephrase it?",
-                 "I'm sorry, but I don't understand your question. Could you please rephrase it?",
-                 "I'm sorry, I don't quite understand your question",
-                 "I'm sorry, I don't know",
-                 "I'm sorry, but I don't know",
-                 "I don't know anything",
-                 "I do not know",
-                 "I don't know",
-                 "I don't know how",
-                 "I do not know how",
-                 "Can you please explain what you mean",
-                 "please explain what you mean",
-                 "please explain",
-                 "I'm sorry, but I don't know how to tell a story. Can you please explain what you mean by",
-                 "I'm sorry but I don't understand what you mean",
-                 "I don't understand",
-                 "I don't have the ability",
-                 "I do not have the ability",
-                 "I do not have",
-                 "I am a language model,",
-                 "I am a large language model,",
-                 "I do not understand your question. Can you please try to make it clearer?",
-                 "I'm sorry, but as an AI language model",
-                 "I apologize, but I cannot rephrase text that I cannot understand. Your post is difficult to read and follow.",
-                 "I apologize, but I am not h2oGPT. I am a language model developed by H2O.ai. How may I help you?",
-                 "Sorry, but I am not an actual Linux shell, nor am I capable of emulating one. I am an open source chat assistant and would be glad t",
-                 "I apologize, but I cannot perform the task you have requested.",
-                 "I'm sorry, I cannot perform this task as I am an AI language model and do not have access",
-                 "I'm sorry, I'm not sure what you're asking for here.",
-                 "I'm not sure what you are asking",
-                 "You need to provide more context",
-                 ]
-    # reduced versions, with redundant parts, just to give context for where they came from
-    unhelpful += ["sorry, I didn't quite understand your question",
-                  "I didn't quite understand your question",
-                  "I didn't understand your question",
-                  "I did not understand your question",
-                  "I did not understand the question",
-                  "could you please rephrase"
-                  "could you rephrase"
-                  "I do not understand your question.",
-                  "I do not understand the question.",
-                  "I do not understand that question.",
-                  "Can you please try to make it clearer",
-                  "Can you try to make it clearer",
-                  "sorry, but as an AI language model",
-                  "as an AI language model",
-                  "I apologize, but I cannot",
-                  "I cannot rephrase text",
-                  "I cannot understand. Your post is difficult to read and follow."
-                  "Your post is difficult to read and follow."
-                  "I apologize, but I am",
-                  "Sorry, but I am not ",
-                  "nor am I capable",
-                  "I am not capable of",
-                  "I apologize, but I cannot perform the task you have requested",
-                  "I cannot perform the task",
-                  "I cannot complete the task",
-                  "I'm sorry",
-                  "I am sorry",
-                  "do not have access",
-                  "not sure what you're asking for",
-                  "not sure what you are asking for",
-                  "not sure what is being asked",
-                  "I'm not sure what you are asking",
-                  "not sure what you are asking",
-                  "You need to provide more context",
-                  "provide more context",
-                  ]
-    unhelpful += ["As a large language model",
-                  "cannot provide any information",
-                  "As an artificial intelligence I do not have the capability",
-                  "As an artificial intelligence I don't have the capability",
-                  "As an artificial intelligence I can't",
-                  "As an artificial intelligence I cannot",
-                  "I am sorry but I do not understand",
-                  "Can you please explain",
-                  "(sorry couldn't resist)",
-                  "(sorry could not resist)",
-                  " :)",
-                  " ;)",
-                  " :-)",
-                  " ;-)",
-                  " lol ",
-                  "Thanks so much!!!",
-                  "Thank You :)!!!",
-                  "Please try not to repeat",
-                  "I am an AI language model",
-                  "I'm a AI assistant that",
-                  "I'm an AI assistant that",
-                  "I am an AI assistant that",
-                  "etc.",
-                  "etc.etc.",
-                  "etc. etc.",
-                  "etc etc",
-                  ]
-    return unhelpful
-
-
-def test_check_unhelpful():
-    # file = '/home/jon/Downloads/openassistant_oasst1_h2ogpt_graded.json'
-    file = '/home/jon/Downloads/openassistant_oasst1_h2ogpt_grades.json'
-    # file = 'h2ogpt-oig-oasst1-instruct-cleaned-v2.json'
-
-    unhelpful = get_unhelpful_list()
-    # data = json.load(open(file, 'rt'))
-    df = pd.read_json(file)
-
-    use_reward_score_threshold = False
-    use_bleu_threshold = False
-    use_sentence_sim = True
-
-    from sacrebleu.metrics import BLEU
-    bleu = BLEU()
-    from nltk.translate.bleu_score import sentence_bleu
-
-    def get_bleu(actual, expected_list):
-        # return bleu.sentence_score(actual, expected_list).score
-        return sentence_bleu(expected_list, actual)
-
-    threshold = 0.0
-    if use_reward_score_threshold:
-        df = df[df['grade_deberta'] > threshold]
-
-    # back to as if original json load
-    data = df.to_dict(orient='records')
-    bads = {}
-    string_all = str(data)
-    for sub in unhelpful:
-        bads[sub] = string_all.count(sub)
-    bads = {k: v for k, v in bads.items() if v > 0}
-    import pprint
-    pp = pprint.PrettyPrinter(indent=4)
-    pp.pprint(bads)
-
-    total_bads = sum(list(bads.values()))
-    print('total_bads: %s' % total_bads, flush=True)
-
-    # check just bot
-    import re
-    convs = [[x.strip() for x in re.split(r'%s|%s' % (human, bot), y['input']) if x.strip()] for y in data]
-    humans = [[x for i, x in enumerate(y) if i % 2 == 0] for y in convs]
-    bots = [[x for i, x in enumerate(y) if i % 2 == 1] for y in convs]
-
-    # FIXME: apply back to json etc., just see for now
-    bleu_threshold = 0.9
-    if use_bleu_threshold:
-        bots = [[x for x in y if get_bleu(x, unhelpful) < bleu_threshold] for y in tqdm(bots)]
-
-    cosine_sim_threshold = 0.8
-    if use_sentence_sim:
-        # pip install sentence_transformers-2.2.2
-        from sentence_transformers import SentenceTransformer
-        # sent_model = 'bert-base-nli-mean-tokens'
-        # sent_model = 'nli-distilroberta-base-v2'
-        sent_model = 'all-MiniLM-L6-v2'
-        model = SentenceTransformer(sent_model)
-        sentence_embeddings = model.encode(unhelpful)
-        from sklearn.metrics.pairwise import cosine_similarity
-        bots = [x for x in tqdm(bots) if
-                np.max(cosine_similarity(model.encode(x), sentence_embeddings)) < cosine_sim_threshold]
-
-    bads_bots = {}
-    string_all = str(bots)
-    for sub in unhelpful:
-        bads_bots[sub] = string_all.count(sub)
-    bads_bots = {k: v for k, v in bads_bots.items() if v > 0}
-    import pprint
-    pp = pprint.PrettyPrinter(indent=4)
-    pp.pprint(bads_bots)
-
-    total_bads_bots = sum(list(bads_bots.values()))
-    print('threshold: %g use_bleu_threshold: %g total_bads_bots: %s total_bots: %s total_humans: %s' % (
-    threshold, use_bleu_threshold, total_bads_bots, len(bots), len(humans)), flush=True)
-
-    # assert len(bads) == 0, bads
-    assert len(bads_bots) == 0, bads_bots
-
-
-def test_fortune2000_personalized():
-    row_list = []
-    import glob
-    if not os.path.isdir("wikitext"):
-        raise RuntimeError("download https://github.com/h2oai/h2ogpt/files/11423008/wikitext.zip and unzip")
-    for file in glob.glob("wikitext/*.txt"):
-        with open(file, "r") as f:
-            blob = f.read()
-        N = 512 * 4
-        row_list.extend([{'input': s, 'prompt_type': 'plain', 'source': "%s" % os.path.basename(file)}
-                         for s in get_sentences(blob, N) if s])
-    personality = create_personality_data()
-    import copy
-    for i in range(10):
-        row_list.extend(copy.deepcopy(personality))
-    np.random.seed(123)
-    np.random.shuffle(row_list)
-    for i in range(len(row_list)):
-        row_list[i]['id'] = i
-    for i in range(len(row_list)):
-        assert row_list[i]['id'] == i
-    with open("h2ogpt-fortune2000-personalized.json", "w") as ff:
-        ff.write(json.dumps(row_list, indent=2))

enums.py

@@ -1,120 +0,0 @@
-from enum import Enum
-
-
-class PromptType(Enum):
-    custom = -1
-    plain = 0
-    instruct = 1
-    quality = 2
-    human_bot = 3
-    dai_faq = 4
-    summarize = 5
-    simple_instruct = 6
-    instruct_vicuna = 7
-    instruct_with_end = 8
-    human_bot_orig = 9
-    prompt_answer = 10
-    open_assistant = 11
-    wizard_lm = 12
-    wizard_mega = 13
-    instruct_vicuna2 = 14
-    instruct_vicuna3 = 15
-    wizard2 = 16
-    wizard3 = 17
-    instruct_simple = 18
-    wizard_vicuna = 19
-    openai = 20
-    openai_chat = 21
-    gptj = 22
-    prompt_answer_openllama = 23
-    vicuna11 = 24
-    mptinstruct = 25
-    mptchat = 26
-    falcon = 27
-    guanaco = 28
-    llama2 = 29
-
-
-class DocumentSubset(Enum):
-    Relevant = 0
-    RelSources = 1
-    TopKSources = 2
-
-
-non_query_commands = [
-    DocumentSubset.RelSources.name,
-    DocumentSubset.TopKSources.name
-]
-
-
-class DocumentChoice(Enum):
-    ALL = 'All'
-
-
-class LangChainMode(Enum):
-    """LangChain mode"""
-
-    DISABLED = "Disabled"
-    LLM = "LLM"
-    ALL = "All"
-    WIKI = "wiki"
-    WIKI_FULL = "wiki_full"
-    USER_DATA = "UserData"
-    MY_DATA = "MyData"
-    GITHUB_H2OGPT = "github h2oGPT"
-    H2O_DAI_DOCS = "DriverlessAI docs"
-
-
-# modes should not be removed from visible list or added by name
-langchain_modes_intrinsic = [LangChainMode.DISABLED.value,
-                             LangChainMode.LLM.value,
-                             LangChainMode.MY_DATA.value]
-
-
-class LangChainAction(Enum):
-    """LangChain action"""
-
-    QUERY = "Query"
-    # WIP:
-    # SUMMARIZE_MAP = "Summarize_map_reduce"
-    SUMMARIZE_MAP = "Summarize"
-    SUMMARIZE_ALL = "Summarize_all"
-    SUMMARIZE_REFINE = "Summarize_refine"
-
-
-class LangChainAgent(Enum):
-    """LangChain agents"""
-
-    SEARCH = "Search"
-    # CSV = "csv"  # WIP
-
-
-no_server_str = no_lora_str = no_model_str = '[None/Remove]'
-
-# from site-packages/langchain/llms/openai.py
-# but needed since ChatOpenAI doesn't have this information
-model_token_mapping = {
-    "gpt-4": 8192,
-    "gpt-4-0314": 8192,
-    "gpt-4-32k": 32768,
-    "gpt-4-32k-0314": 32768,
-    "gpt-3.5-turbo": 4096,
-    "gpt-3.5-turbo-16k": 16 * 1024,
-    "gpt-3.5-turbo-0301": 4096,
-    "text-ada-001": 2049,
-    "ada": 2049,
-    "text-babbage-001": 2040,
-    "babbage": 2049,
-    "text-curie-001": 2049,
-    "curie": 2049,
-    "davinci": 2049,
-    "text-davinci-003": 4097,
-    "text-davinci-002": 4097,
-    "code-davinci-002": 8001,
-    "code-davinci-001": 8001,
-    "code-cushman-002": 2048,
-    "code-cushman-001": 2048,
-}
-
-source_prefix = "Sources [Score | Link]:"
-source_postfix = "End Sources<p>"

evaluate_params.py

@@ -1,52 +0,0 @@
-input_args_list = ['model_state', 'my_db_state', 'selection_docs_state']
-
-
-no_default_param_names = [
-    'instruction',
-    'iinput',
-    'context',
-    'instruction_nochat',
-    'iinput_nochat',
-]
-
-gen_hyper = ['temperature',
-             'top_p',
-             'top_k',
-             'num_beams',
-             'max_new_tokens',
-             'min_new_tokens',
-             'early_stopping',
-             'max_time',
-             'repetition_penalty',
-             'num_return_sequences',
-             'do_sample',
-             ]
-
-eval_func_param_names = ['instruction',
-                         'iinput',
-                         'context',
-                         'stream_output',
-                         'prompt_type',
-                         'prompt_dict'] + \
-                        gen_hyper + \
-                        ['chat',
-                         'instruction_nochat',
-                         'iinput_nochat',
-                         'langchain_mode',
-                         'add_chat_history_to_context',
-                         'langchain_action',
-                         'langchain_agents',
-                         'top_k_docs',
-                         'chunk',
-                         'chunk_size',
-                         'document_subset',
-                         'document_choice',
-                         ]
-
-# form evaluate defaults for submit_nochat_api
-eval_func_param_names_defaults = eval_func_param_names.copy()
-for k in no_default_param_names:
-    if k in eval_func_param_names_defaults:
-        eval_func_param_names_defaults.remove(k)
-
-eval_extra_columns = ['prompt', 'response', 'score']

generate.py

@@ -1,16 +0,0 @@
-import os
-import sys
-
-if os.path.dirname(os.path.abspath(__file__)) not in sys.path:
-    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-
-from src.gen import main
-from src.utils import H2O_Fire
-
-
-def entrypoint_main():
-    H2O_Fire(main)
-
-
-if __name__ == "__main__":
-    entrypoint_main()

gpt4all_llm.py

@@ -1,316 +0,0 @@
-import inspect
-import os
-from functools import partial
-from typing import Dict, Any, Optional, List
-from langchain.callbacks.manager import CallbackManagerForLLMRun
-from pydantic import root_validator
-from langchain.llms import gpt4all
-from dotenv import dotenv_values
-
-from utils import FakeTokenizer
-
-
-def get_model_tokenizer_gpt4all(base_model, **kwargs):
-    # defaults (some of these are generation parameters, so need to be passed in at generation time)
-    model_kwargs = dict(n_threads=os.cpu_count() // 2,
-                        temp=kwargs.get('temperature', 0.2),
-                        top_p=kwargs.get('top_p', 0.75),
-                        top_k=kwargs.get('top_k', 40),
-                        n_ctx=2048 - 256)
-    env_gpt4all_file = ".env_gpt4all"
-    model_kwargs.update(dotenv_values(env_gpt4all_file))
-    # make int or float if can to satisfy types for class
-    for k, v in model_kwargs.items():
-        try:
-            if float(v) == int(v):
-                model_kwargs[k] = int(v)
-            else:
-                model_kwargs[k] = float(v)
-        except:
-            pass
-
-    if base_model == "llama":
-        if 'model_path_llama' not in model_kwargs:
-            raise ValueError("No model_path_llama in %s" % env_gpt4all_file)
-        model_path = model_kwargs.pop('model_path_llama')
-        # FIXME: GPT4All version of llama doesn't handle new quantization, so use llama_cpp_python
-        from llama_cpp import Llama
-        # llama sets some things at init model time, not generation time
-        func_names = list(inspect.signature(Llama.__init__).parameters)
-        model_kwargs = {k: v for k, v in model_kwargs.items() if k in func_names}
-        model_kwargs['n_ctx'] = int(model_kwargs['n_ctx'])
-        model = Llama(model_path=model_path, **model_kwargs)
-    elif base_model in "gpt4all_llama":
-        if 'model_name_gpt4all_llama' not in model_kwargs and 'model_path_gpt4all_llama' not in model_kwargs:
-            raise ValueError("No model_name_gpt4all_llama or model_path_gpt4all_llama in %s" % env_gpt4all_file)
-        model_name = model_kwargs.pop('model_name_gpt4all_llama')
-        model_type = 'llama'
-        from gpt4all import GPT4All as GPT4AllModel
-        model = GPT4AllModel(model_name=model_name, model_type=model_type)
-    elif base_model in "gptj":
-        if 'model_name_gptj' not in model_kwargs and 'model_path_gptj' not in model_kwargs:
-            raise ValueError("No model_name_gpt4j or model_path_gpt4j in %s" % env_gpt4all_file)
-        model_name = model_kwargs.pop('model_name_gptj')
-        model_type = 'gptj'
-        from gpt4all import GPT4All as GPT4AllModel
-        model = GPT4AllModel(model_name=model_name, model_type=model_type)
-    else:
-        raise ValueError("No such base_model %s" % base_model)
-    return model, FakeTokenizer(), 'cpu'
-
-
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-
-
-class H2OStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
-
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Run on new LLM token. Only available when streaming is enabled."""
-        # streaming to std already occurs without this
-        # sys.stdout.write(token)
-        # sys.stdout.flush()
-        pass
-
-
-def get_model_kwargs(env_kwargs, default_kwargs, cls, exclude_list=[]):
-    # default from class
-    model_kwargs = {k: v.default for k, v in dict(inspect.signature(cls).parameters).items() if k not in exclude_list}
-    # from our defaults
-    model_kwargs.update(default_kwargs)
-    # from user defaults
-    model_kwargs.update(env_kwargs)
-    # ensure only valid keys
-    func_names = list(inspect.signature(cls).parameters)
-    model_kwargs = {k: v for k, v in model_kwargs.items() if k in func_names}
-    return model_kwargs
-
-
-def get_llm_gpt4all(model_name,
-                    model=None,
-                    max_new_tokens=256,
-                    temperature=0.1,
-                    repetition_penalty=1.0,
-                    top_k=40,
-                    top_p=0.7,
-                    streaming=False,
-                    callbacks=None,
-                    prompter=None,
-                    context='',
-                    iinput='',
-                    verbose=False,
-                    ):
-    assert prompter is not None
-    env_gpt4all_file = ".env_gpt4all"
-    env_kwargs = dotenv_values(env_gpt4all_file)
-    max_tokens = env_kwargs.pop('max_tokens', 2048 - max_new_tokens)
-    default_kwargs = dict(context_erase=0.5,
-                          n_batch=1,
-                          max_tokens=max_tokens,
-                          n_predict=max_new_tokens,
-                          repeat_last_n=64 if repetition_penalty != 1.0 else 0,
-                          repeat_penalty=repetition_penalty,
-                          temp=temperature,
-                          temperature=temperature,
-                          top_k=top_k,
-                          top_p=top_p,
-                          use_mlock=True,
-                          verbose=verbose)
-    if model_name == 'llama':
-        cls = H2OLlamaCpp
-        model_path = env_kwargs.pop('model_path_llama') if model is None else model
-        model_kwargs = get_model_kwargs(env_kwargs, default_kwargs, cls, exclude_list=['lc_kwargs'])
-        model_kwargs.update(dict(model_path=model_path, callbacks=callbacks, streaming=streaming,
-                                 prompter=prompter, context=context, iinput=iinput))
-        llm = cls(**model_kwargs)
-        llm.client.verbose = verbose
-    elif model_name == 'gpt4all_llama':
-        cls = H2OGPT4All
-        model_path = env_kwargs.pop('model_path_gpt4all_llama') if model is None else model
-        model_kwargs = get_model_kwargs(env_kwargs, default_kwargs, cls, exclude_list=['lc_kwargs'])
-        model_kwargs.update(
-            dict(model=model_path, backend='llama', callbacks=callbacks, streaming=streaming,
-                 prompter=prompter, context=context, iinput=iinput))
-        llm = cls(**model_kwargs)
-    elif model_name == 'gptj':
-        cls = H2OGPT4All
-        model_path = env_kwargs.pop('model_path_gptj') if model is None else model
-        model_kwargs = get_model_kwargs(env_kwargs, default_kwargs, cls, exclude_list=['lc_kwargs'])
-        model_kwargs.update(
-            dict(model=model_path, backend='gptj', callbacks=callbacks, streaming=streaming,
-                 prompter=prompter, context=context, iinput=iinput))
-        llm = cls(**model_kwargs)
-    else:
-        raise RuntimeError("No such model_name %s" % model_name)
-    return llm
-
-
-class H2OGPT4All(gpt4all.GPT4All):
-    model: Any
-    prompter: Any
-    context: Any = ''
-    iinput: Any = ''
-    """Path to the pre-trained GPT4All model file."""
-
-    @root_validator()
-    def validate_environment(cls, values: Dict) -> Dict:
-        """Validate that the python package exists in the environment."""
-        try:
-            if isinstance(values["model"], str):
-                from gpt4all import GPT4All as GPT4AllModel
-
-                full_path = values["model"]
-                model_path, delimiter, model_name = full_path.rpartition("/")
-                model_path += delimiter
-
-                values["client"] = GPT4AllModel(
-                    model_name=model_name,
-                    model_path=model_path or None,
-                    model_type=values["backend"],
-                    allow_download=False,
-                )
-                if values["n_threads"] is not None:
-                    # set n_threads
-                    values["client"].model.set_thread_count(values["n_threads"])
-            else:
-                values["client"] = values["model"]
-            try:
-                values["backend"] = values["client"].model_type
-            except AttributeError:
-                # The below is for compatibility with GPT4All Python bindings <= 0.2.3.
-                values["backend"] = values["client"].model.model_type
-
-        except ImportError:
-            raise ValueError(
-                "Could not import gpt4all python package. "
-                "Please install it with `pip install gpt4all`."
-            )
-        return values
-
-    def _call(
-            self,
-            prompt: str,
-            stop: Optional[List[str]] = None,
-            run_manager: Optional[CallbackManagerForLLMRun] = None,
-            **kwargs,
-    ) -> str:
-        # Roughly 4 chars per token if natural language
-        n_ctx = 2048
-        prompt = prompt[-self.max_tokens * 4:]
-
-        # use instruct prompting
-        data_point = dict(context=self.context, instruction=prompt, input=self.iinput)
-        prompt = self.prompter.generate_prompt(data_point)
-
-        verbose = False
-        if verbose:
-            print("_call prompt: %s" % prompt, flush=True)
-        # FIXME: GPT4ALl doesn't support yield during generate, so cannot support streaming except via itself to stdout
-        return super()._call(prompt, stop=stop, run_manager=run_manager)
-
-
-from langchain.llms import LlamaCpp
-
-
-class H2OLlamaCpp(LlamaCpp):
-    model_path: Any
-    prompter: Any
-    context: Any
-    iinput: Any
-    """Path to the pre-trained GPT4All model file."""
-
-    @root_validator()
-    def validate_environment(cls, values: Dict) -> Dict:
-        """Validate that llama-cpp-python library is installed."""
-        if isinstance(values["model_path"], str):
-            model_path = values["model_path"]
-            model_param_names = [
-                "lora_path",
-                "lora_base",
-                "n_ctx",
-                "n_parts",
-                "seed",
-                "f16_kv",
-                "logits_all",
-                "vocab_only",
-                "use_mlock",
-                "n_threads",
-                "n_batch",
-                "use_mmap",
-                "last_n_tokens_size",
-            ]
-            model_params = {k: values[k] for k in model_param_names}
-            # For backwards compatibility, only include if non-null.
-            if values["n_gpu_layers"] is not None:
-                model_params["n_gpu_layers"] = values["n_gpu_layers"]
-
-            try:
-                from llama_cpp import Llama
-
-                values["client"] = Llama(model_path, **model_params)
-            except ImportError:
-                raise ModuleNotFoundError(
-                    "Could not import llama-cpp-python library. "
-                    "Please install the llama-cpp-python library to "
-                    "use this embedding model: pip install llama-cpp-python"
-                )
-            except Exception as e:
-                raise ValueError(
-                    f"Could not load Llama model from path: {model_path}. "
-                    f"Received error {e}"
-                )
-        else:
-            values["client"] = values["model_path"]
-        return values
-
-    def _call(
-            self,
-            prompt: str,
-            stop: Optional[List[str]] = None,
-            run_manager: Optional[CallbackManagerForLLMRun] = None,
-            **kwargs,
-    ) -> str:
-        verbose = False
-        # tokenize twice, just to count tokens, since llama cpp python wrapper has no way to truncate
-        # still have to avoid crazy sizes, else hit llama_tokenize: too many tokens -- might still hit, not fatal
-        prompt = prompt[-self.n_ctx * 4:]
-        prompt_tokens = self.client.tokenize(b" " + prompt.encode("utf-8"))
-        num_prompt_tokens = len(prompt_tokens)
-        if num_prompt_tokens > self.n_ctx:
-            # conservative by using int()
-            chars_per_token = int(len(prompt) / num_prompt_tokens)
-            prompt = prompt[-self.n_ctx * chars_per_token:]
-            if verbose:
-                print("reducing tokens, assuming average of %s chars/token: %s" % chars_per_token, flush=True)
-                prompt_tokens2 = self.client.tokenize(b" " + prompt.encode("utf-8"))
-                num_prompt_tokens2 = len(prompt_tokens2)
-                print("reduced tokens from %d -> %d" % (num_prompt_tokens, num_prompt_tokens2), flush=True)
-
-        # use instruct prompting
-        data_point = dict(context=self.context, instruction=prompt, input=self.iinput)
-        prompt = self.prompter.generate_prompt(data_point)
-
-        if verbose:
-            print("_call prompt: %s" % prompt, flush=True)
-
-        if self.streaming:
-            text_callback = None
-            if run_manager:
-                text_callback = partial(
-                    run_manager.on_llm_new_token, verbose=self.verbose
-                )
-            # parent handler of streamer expects to see prompt first else output="" and lose if prompt=None in prompter
-            if text_callback:
-                text_callback(prompt)
-            text = ""
-            for token in self.stream(prompt=prompt, stop=stop, run_manager=run_manager):
-                text_chunk = token["choices"][0]["text"]
-                # self.stream already calls text_callback
-                # if text_callback:
-                #    text_callback(text_chunk)
-                text += text_chunk
-            return text
-        else:
-            params = self._get_parameters(stop)
-            params = {**params, **kwargs}
-            result = self.client(prompt=prompt, **params)
-            return result["choices"][0]["text"]

gradio_themes.py

@@ -1,231 +0,0 @@
-from __future__ import annotations
-
-from typing import Iterable
-
-from gradio.themes.soft import Soft
-from gradio.themes import Color, Size
-from gradio.themes.utils import colors, sizes, fonts
-
-h2o_yellow = Color(
-    name="yellow",
-    c50="#fffef2",
-    c100="#fff9e6",
-    c200="#ffecb3",
-    c300="#ffe28c",
-    c400="#ffd659",
-    c500="#fec925",
-    c600="#e6ac00",
-    c700="#bf8f00",
-    c800="#a67c00",
-    c900="#664d00",
-    c950="#403000",
-)
-h2o_gray = Color(
-    name="gray",
-    c50="#f8f8f8",
-    c100="#e5e5e5",
-    c200="#cccccc",
-    c300="#b2b2b2",
-    c400="#999999",
-    c500="#7f7f7f",
-    c600="#666666",
-    c700="#4c4c4c",
-    c800="#333333",
-    c900="#191919",
-    c950="#0d0d0d",
-)
-
-
-text_xsm = Size(
-    name="text_xsm",
-    xxs="4px",
-    xs="5px",
-    sm="6px",
-    md="7px",
-    lg="8px",
-    xl="10px",
-    xxl="12px",
-)
-
-
-spacing_xsm = Size(
-    name="spacing_xsm",
-    xxs="1px",
-    xs="1px",
-    sm="1px",
-    md="2px",
-    lg="3px",
-    xl="5px",
-    xxl="7px",
-)
-
-
-radius_xsm = Size(
-    name="radius_xsm",
-    xxs="1px",
-    xs="1px",
-    sm="1px",
-    md="2px",
-    lg="3px",
-    xl="5px",
-    xxl="7px",
-)
-
-
-class H2oTheme(Soft):
-    def __init__(
-            self,
-            *,
-            primary_hue: colors.Color | str = h2o_yellow,
-            secondary_hue: colors.Color | str = h2o_yellow,
-            neutral_hue: colors.Color | str = h2o_gray,
-            spacing_size: sizes.Size | str = sizes.spacing_md,
-            radius_size: sizes.Size | str = sizes.radius_md,
-            text_size: sizes.Size | str = sizes.text_lg,
-            font: fonts.Font
-            | str
-            | Iterable[fonts.Font | str] = (
-                fonts.GoogleFont("Montserrat"),
-                "ui-sans-serif",
-                "system-ui",
-                "sans-serif",
-            ),
-            font_mono: fonts.Font
-            | str
-            | Iterable[fonts.Font | str] = (
-                fonts.GoogleFont("IBM Plex Mono"),
-                "ui-monospace",
-                "Consolas",
-                "monospace",
-            ),
-    ):
-        super().__init__(
-            primary_hue=primary_hue,
-            secondary_hue=secondary_hue,
-            neutral_hue=neutral_hue,
-            spacing_size=spacing_size,
-            radius_size=radius_size,
-            text_size=text_size,
-            font=font,
-            font_mono=font_mono,
-        )
-        super().set(
-            link_text_color="#3344DD",
-            link_text_color_hover="#3344DD",
-            link_text_color_visited="#3344DD",
-            link_text_color_dark="#74abff",
-            link_text_color_hover_dark="#a3c8ff",
-            link_text_color_active_dark="#a3c8ff",
-            link_text_color_visited_dark="#74abff",
-            button_primary_text_color="*neutral_950",
-            button_primary_text_color_dark="*neutral_950",
-            button_primary_background_fill="*primary_500",
-            button_primary_background_fill_dark="*primary_500",
-            block_label_background_fill="*primary_500",
-            block_label_background_fill_dark="*primary_500",
-            block_label_text_color="*neutral_950",
-            block_label_text_color_dark="*neutral_950",
-            block_title_text_color="*neutral_950",
-            block_title_text_color_dark="*neutral_950",
-            block_background_fill_dark="*neutral_950",
-            body_background_fill="*neutral_50",
-            body_background_fill_dark="*neutral_900",
-            background_fill_primary_dark="*block_background_fill",
-            block_radius="0 0 8px 8px",
-            checkbox_label_text_color_selected_dark='#000000',
-            #checkbox_label_text_size="*text_xs",  # too small for iPhone etc. but good if full large screen zoomed to fit
-            checkbox_label_text_size="*text_sm",
-            #radio_circle="""url("data:image/svg+xml,%3csvg viewBox='0 0 32 32' fill='white' xmlns='http://www.w3.org/2000/svg'%3e%3ccircle cx='32' cy='32' r='1'/%3e%3c/svg%3e")""",
-            #checkbox_border_width=1,
-            #heckbox_border_width_dark=1,
-        )
-
-
-class SoftTheme(Soft):
-    def __init__(
-            self,
-            *,
-            primary_hue: colors.Color | str = colors.indigo,
-            secondary_hue: colors.Color | str = colors.indigo,
-            neutral_hue: colors.Color | str = colors.gray,
-            spacing_size: sizes.Size | str = sizes.spacing_md,
-            radius_size: sizes.Size | str = sizes.radius_md,
-            text_size: sizes.Size | str = sizes.text_md,
-            font: fonts.Font
-            | str
-            | Iterable[fonts.Font | str] = (
-                fonts.GoogleFont("Montserrat"),
-                "ui-sans-serif",
-                "system-ui",
-                "sans-serif",
-            ),
-            font_mono: fonts.Font
-            | str
-            | Iterable[fonts.Font | str] = (
-                fonts.GoogleFont("IBM Plex Mono"),
-                "ui-monospace",
-                "Consolas",
-                "monospace",
-            ),
-    ):
-        super().__init__(
-            primary_hue=primary_hue,
-            secondary_hue=secondary_hue,
-            neutral_hue=neutral_hue,
-            spacing_size=spacing_size,
-            radius_size=radius_size,
-            text_size=text_size,
-            font=font,
-            font_mono=font_mono,
-        )
-        super().set(
-            checkbox_label_text_size="*text_sm",
-        )
-
-
-h2o_logo = '<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" width="100%" height="100%"' \
-           ' viewBox="0 0 600.28 600.28"><defs><style>.cls-1{fill:#fec925;}.cls-2{fill:#161616;}.cls-3{fill:' \
-           '#54585a;}</style></defs><g id="Fill-1"><rect class="cls-1" width="600.28" height="600.28" ' \
-           'rx="23.24"/></g><path class="cls-2" d="M174.33,246.06v92.78H152.86v-38H110.71v38H89.24V246.06h21.' \
-           '47v36.58h42.15V246.06Z"/><path class="cls-2" d="M259.81,321.34v17.5H189.7V324.92l35.78-33.8c8.22-7.' \
-           '82,9.68-12.59,9.68-17.09,0-7.29-5-11.53-14.85-11.53-7.95,0-14.71,3-19.21,9.27L185.46,261.7c7.15-10' \
-           '.47,20.14-17.23,36.84-17.23,20.68,0,34.46,10.6,34.46,27.44,0,9-2.52,17.22-15.51,29.29l-21.33,20.14Z"' \
-           '/><path class="cls-2" d="M268.69,292.45c0-27.57,21.47-48,50.76-48s50.76,20.28,50.76,48-21.6,48-50.' \
-           '76,48S268.69,320,268.69,292.45Zm79.78,0c0-17.63-12.46-29.69-29-29.69s-29,12.06-29,29.69,12.46,29.69' \
-           ',29,29.69S348.47,310.08,348.47,292.45Z"/><path class="cls-3" d="M377.23,326.91c0-7.69,5.7-12.73,12.' \
-           '85-12.73s12.86,5,12.86,12.73a12.86,12.86,0,1,1-25.71,0Z"/><path class="cls-3" d="M481.4,298.15v40.' \
-           '69H462.05V330c-3.84,6.49-11.27,9.94-21.74,9.94-16.7,0-26.64-9.28-26.64-21.61,0-12.59,8.88-21.34,30.' \
-           '62-21.34h16.43c0-8.87-5.3-14-16.43-14-7.55,0-15.37,2.51-20.54,6.62l-7.43-14.44c7.82-5.57,19.35-8.' \
-           '62,30.75-8.62C468.81,266.47,481.4,276.54,481.4,298.15Zm-20.68,18.16V309H446.54c-9.67,0-12.72,3.57-' \
-           '12.72,8.35,0,5.16,4.37,8.61,11.66,8.61C452.37,326,458.34,322.8,460.72,316.31Z"/><path class="cls-3"' \
-           ' d="M497.56,246.06c0-6.49,5.17-11.53,12.86-11.53s12.86,4.77,12.86,11.13c0,6.89-5.17,11.93-12.86,' \
-           '11.93S497.56,252.55,497.56,246.06Zm2.52,21.47h20.68v71.31H500.08Z"/></svg>'
-
-
-def get_h2o_title(title, description):
-    # NOTE: Check full width desktop, smallest width browser desktop, iPhone browsers to ensure no overlap etc.
-    return f"""<div style="float:left; justify-content:left; height: 80px; width: 195px; margin-top:0px">
-                    {description}
-                </div>
-                <div style="display:flex; justify-content:center; margin-bottom:30px; margin-right:330px;">
-                    <div style="height: 60px; width: 60px; margin-right:20px;">{h2o_logo}</div>
-                    <h1 style="line-height:60px">{title}</h1>
-                </div>
-                <div style="float:right; height: 80px; width: 80px; margin-top:-100px">
-                    <img src="https://raw.githubusercontent.com/h2oai/h2ogpt/main/docs/h2o-qr.png">
-                </div>
-                """
-
-
-def get_simple_title(title, description):
-    return f"""{description}<h1 align="center"> {title}</h1>"""
-
-
-def get_dark_js():
-    return """() => {
-        if (document.querySelectorAll('.dark').length) {
-            document.querySelectorAll('.dark').forEach(el => el.classList.remove('dark'));
-        } else {
-            document.querySelector('body').classList.add('dark');
-        }
-    }"""

h2oai_pipeline.py

@@ -1,201 +0,0 @@
-import os
-
-from transformers import TextGenerationPipeline
-from transformers.pipelines.text_generation import ReturnType
-
-from stopping import get_stopping
-from prompter import Prompter, PromptType
-
-
-class H2OTextGenerationPipeline(TextGenerationPipeline):
-    def __init__(self, *args, debug=False, chat=False, stream_output=False,
-                 sanitize_bot_response=False,
-                 use_prompter=True, prompter=None,
-                 context='', iinput='',
-                 prompt_type=None, prompt_dict=None,
-                 max_input_tokens=2048 - 256, **kwargs):
-        """
-        HF-like pipeline, but handle instruction prompting and stopping (for some models)
-        :param args:
-        :param debug:
-        :param chat:
-        :param stream_output:
-        :param sanitize_bot_response:
-        :param use_prompter: Whether to use prompter.  If pass prompt_type, will make prompter
-        :param prompter: prompter, can pass if have already
-        :param prompt_type: prompt_type, e.g. human_bot.  See prompt_type to model mapping in from prompter.py.
-                            If use_prompter, then will make prompter and use it.
-        :param prompt_dict: dict of get_prompt(, return_dict=True) for prompt_type=custom
-        :param max_input_tokens:
-        :param kwargs:
-        """
-        super().__init__(*args, **kwargs)
-        self.prompt_text = None
-        self.use_prompter = use_prompter
-        self.prompt_type = prompt_type
-        self.prompt_dict = prompt_dict
-        self.prompter = prompter
-        self.context = context
-        self.iinput = iinput
-        if self.use_prompter:
-            if self.prompter is not None:
-                assert self.prompter.prompt_type is not None
-            else:
-                self.prompter = Prompter(self.prompt_type, self.prompt_dict, debug=debug, chat=chat,
-                                         stream_output=stream_output)
-            self.human = self.prompter.humanstr
-            self.bot = self.prompter.botstr
-            self.can_stop = True
-        else:
-            self.prompter = None
-            self.human = None
-            self.bot = None
-            self.can_stop = False
-        self.sanitize_bot_response = sanitize_bot_response
-        self.max_input_tokens = max_input_tokens  # not for generate, so ok that not kwargs
-
-    @staticmethod
-    def limit_prompt(prompt_text, tokenizer, max_prompt_length=None):
-        verbose = bool(int(os.getenv('VERBOSE_PIPELINE', '0')))
-
-        if hasattr(tokenizer, 'model_max_length'):
-            # model_max_length only defined for generate.py, not raw use of h2oai_pipeline.py
-            model_max_length = tokenizer.model_max_length
-            if max_prompt_length is not None:
-                model_max_length = min(model_max_length, max_prompt_length)
-            # cut at some upper likely limit to avoid excessive tokenization etc
-            # upper bound of 10 chars/token, e.g. special chars sometimes are long
-            if len(prompt_text) > model_max_length * 10:
-                len0 = len(prompt_text)
-                prompt_text = prompt_text[-model_max_length * 10:]
-                if verbose:
-                    print("Cut of input: %s -> %s" % (len0, len(prompt_text)), flush=True)
-        else:
-            # unknown
-            model_max_length = None
-
-        num_prompt_tokens = None
-        if model_max_length is not None:
-            # can't wait for "hole" if not plain prompt_type, since would lose prefix like <human>:
-            # For https://github.com/h2oai/h2ogpt/issues/192
-            for trial in range(0, 3):
-                prompt_tokens = tokenizer(prompt_text)['input_ids']
-                num_prompt_tokens = len(prompt_tokens)
-                if num_prompt_tokens > model_max_length:
-                    # conservative by using int()
-                    chars_per_token = int(len(prompt_text) / num_prompt_tokens)
-                    # keep tail, where question is if using langchain
-                    prompt_text = prompt_text[-model_max_length * chars_per_token:]
-                    if verbose:
-                        print("reducing %s tokens, assuming average of %s chars/token for %s characters" % (
-                            num_prompt_tokens, chars_per_token, len(prompt_text)), flush=True)
-                else:
-                    if verbose:
-                        print("using %s tokens with %s chars" % (num_prompt_tokens, len(prompt_text)), flush=True)
-                    break
-
-            # Why Below False: don't limit max_new_tokens more, just rely upon stopping to reach limit of model
-            if False:
-                # if input prompt is some number of tokens, despite user request, can't have max_new_tokens more
-                #
-                assert num_prompt_tokens is not None
-                if self.prompt_type not in [PromptType.plain.name, PromptType.plain.value]:
-                    # then give room for prompt
-                    fudge = 20
-                else:
-                    fudge = 0
-                max_new_tokens = max(0, min(generate_kwargs['max_new_tokens'],
-                                            model_max_length - (num_prompt_tokens + fudge)))
-                if max_new_tokens < generate_kwargs['max_new_tokens']:
-                    if verbose:
-                        print("Reduced max_new_tokens from %s -> %s" % (
-                        generate_kwargs['max_new_tokens'], max_new_tokens))
-                    generate_kwargs['max_new_tokens'] = max_new_tokens
-        return prompt_text, num_prompt_tokens
-
-    def preprocess(self, prompt_text, prefix="", handle_long_generation=None, **generate_kwargs):
-        prompt_text, num_prompt_tokens = H2OTextGenerationPipeline.limit_prompt(prompt_text, self.tokenizer)
-
-        data_point = dict(context=self.context, instruction=prompt_text, input=self.iinput)
-        if self.prompter is not None:
-            prompt_text = self.prompter.generate_prompt(data_point)
-        self.prompt_text = prompt_text
-        if handle_long_generation is None:
-            # forces truncation of inputs to avoid critical failure
-            handle_long_generation = None  # disable with new approaches
-        return super().preprocess(prompt_text, prefix=prefix, handle_long_generation=handle_long_generation,
-                                  **generate_kwargs)
-
-    def postprocess(self, model_outputs, return_type=ReturnType.FULL_TEXT, clean_up_tokenization_spaces=True):
-        records = super().postprocess(model_outputs, return_type=return_type,
-                                      clean_up_tokenization_spaces=clean_up_tokenization_spaces)
-        for rec in records:
-            if self.use_prompter:
-                outputs = rec['generated_text']
-                outputs = self.prompter.get_response(outputs, prompt=self.prompt_text,
-                                                     sanitize_bot_response=self.sanitize_bot_response)
-            elif self.bot and self.human:
-                outputs = rec['generated_text'].split(self.bot)[1].split(self.human)[0]
-            else:
-                outputs = rec['generated_text']
-            rec['generated_text'] = outputs
-            print("prompt: %s\noutputs: %s\n\n" % (self.prompt_text, outputs), flush=True)
-        return records
-
-    def _forward(self, model_inputs, **generate_kwargs):
-        if self.can_stop:
-            stopping_criteria = get_stopping(self.prompt_type, self.prompt_dict,
-                                             self.tokenizer, self.device,
-                                             human=self.human, bot=self.bot,
-                                             model_max_length=self.tokenizer.model_max_length)
-            generate_kwargs['stopping_criteria'] = stopping_criteria
-        # return super()._forward(model_inputs, **generate_kwargs)
-        return self.__forward(model_inputs, **generate_kwargs)
-
-    # FIXME: Copy-paste of original _forward, but removed copy.deepcopy()
-    # FIXME: https://github.com/h2oai/h2ogpt/issues/172
-    def __forward(self, model_inputs, **generate_kwargs):
-        input_ids = model_inputs["input_ids"]
-        attention_mask = model_inputs.get("attention_mask", None)
-        # Allow empty prompts
-        if input_ids.shape[1] == 0:
-            input_ids = None
-            attention_mask = None
-            in_b = 1
-        else:
-            in_b = input_ids.shape[0]
-        prompt_text = model_inputs.pop("prompt_text")
-
-        ## If there is a prefix, we may need to adjust the generation length. Do so without permanently modifying
-        ## generate_kwargs, as some of the parameterization may come from the initialization of the pipeline.
-        # generate_kwargs = copy.deepcopy(generate_kwargs)
-        prefix_length = generate_kwargs.pop("prefix_length", 0)
-        if prefix_length > 0:
-            has_max_new_tokens = "max_new_tokens" in generate_kwargs or (
-                    "generation_config" in generate_kwargs
-                    and generate_kwargs["generation_config"].max_new_tokens is not None
-            )
-            if not has_max_new_tokens:
-                generate_kwargs["max_length"] = generate_kwargs.get("max_length") or self.model.config.max_length
-                generate_kwargs["max_length"] += prefix_length
-            has_min_new_tokens = "min_new_tokens" in generate_kwargs or (
-                    "generation_config" in generate_kwargs
-                    and generate_kwargs["generation_config"].min_new_tokens is not None
-            )
-            if not has_min_new_tokens and "min_length" in generate_kwargs:
-                generate_kwargs["min_length"] += prefix_length
-
-        # BS x SL
-        generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
-        out_b = generated_sequence.shape[0]
-        if self.framework == "pt":
-            generated_sequence = generated_sequence.reshape(in_b, out_b // in_b, *generated_sequence.shape[1:])
-        elif self.framework == "tf":
-            from transformers import is_tf_available
-            if is_tf_available():
-                import tensorflow as tf
-                generated_sequence = tf.reshape(generated_sequence,
-                                                (in_b, out_b // in_b, *generated_sequence.shape[1:]))
-            else:
-                raise ValueError("TF not avaialble.")
-        return {"generated_sequence": generated_sequence, "input_ids": input_ids, "prompt_text": prompt_text}

loaders.py

@@ -1,61 +0,0 @@
-import functools
-
-
-def get_loaders(model_name, reward_type, llama_type=None, load_gptq=''):
-    # NOTE: Some models need specific new prompt_type
-    # E.g. t5_xxl_true_nli_mixture has input format: "premise: PREMISE_TEXT hypothesis: HYPOTHESIS_TEXT".)
-    if load_gptq:
-        from transformers import AutoTokenizer
-        from auto_gptq import AutoGPTQForCausalLM
-        use_triton = False
-        functools.partial(AutoGPTQForCausalLM.from_quantized, quantize_config=None, use_triton=use_triton)
-        return AutoGPTQForCausalLM.from_quantized, AutoTokenizer
-    if llama_type is None:
-        llama_type = "llama" in model_name.lower()
-    if llama_type:
-        from transformers import LlamaForCausalLM, LlamaTokenizer
-        return LlamaForCausalLM.from_pretrained, LlamaTokenizer
-    elif 'distilgpt2' in model_name.lower():
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        return AutoModelForCausalLM.from_pretrained, AutoTokenizer
-    elif 'gpt2' in model_name.lower():
-        from transformers import GPT2LMHeadModel, GPT2Tokenizer
-        return GPT2LMHeadModel.from_pretrained, GPT2Tokenizer
-    elif 'mbart-' in model_name.lower():
-        from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-        return MBartForConditionalGeneration.from_pretrained, MBart50TokenizerFast
-    elif 't5' == model_name.lower() or \
-         't5-' in model_name.lower() or \
-         'flan-' in model_name.lower():
-        from transformers import AutoTokenizer, T5ForConditionalGeneration
-        return T5ForConditionalGeneration.from_pretrained, AutoTokenizer
-    elif 'bigbird' in model_name:
-        from transformers import BigBirdPegasusForConditionalGeneration, AutoTokenizer
-        return BigBirdPegasusForConditionalGeneration.from_pretrained, AutoTokenizer
-    elif 'bart-large-cnn-samsum' in model_name or 'flan-t5-base-samsum' in model_name:
-        from transformers import pipeline
-        return pipeline, "summarization"
-    elif reward_type or 'OpenAssistant/reward-model'.lower() in model_name.lower():
-        from transformers import AutoModelForSequenceClassification, AutoTokenizer
-        return AutoModelForSequenceClassification.from_pretrained, AutoTokenizer
-    else:
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-        model_loader = AutoModelForCausalLM
-        tokenizer_loader = AutoTokenizer
-        return model_loader.from_pretrained, tokenizer_loader
-
-
-def get_tokenizer(tokenizer_loader, tokenizer_base_model, local_files_only, resume_download, use_auth_token):
-    tokenizer = tokenizer_loader.from_pretrained(tokenizer_base_model,
-                                                 local_files_only=local_files_only,
-                                                 resume_download=resume_download,
-                                                 use_auth_token=use_auth_token,
-                                                 padding_side='left')
-
-    tokenizer.pad_token_id = 0  # different from the eos token
-    # when generating, we will use the logits of right-most token to predict the next token
-    # so the padding should be on the left,
-    # e.g. see: https://huggingface.co/transformers/v4.11.3/model_doc/t5.html#inference
-    tokenizer.padding_side = "left"  # Allow batched inference
-
-    return tokenizer

prompter.py

@@ -1,871 +0,0 @@
-import os
-import ast
-import time
-from enums import PromptType  # also supports imports from this file from other files
-
-non_hf_types = ['gpt4all_llama', 'llama', 'gptj']
-
-prompt_type_to_model_name = {
-    'plain': [
-        'EleutherAI/gpt-j-6B',
-        'EleutherAI/pythia-6.9b',
-        'EleutherAI/pythia-12b',
-        'EleutherAI/pythia-12b-deduped',
-        'EleutherAI/gpt-neox-20b',
-        'openlm-research/open_llama_7b_700bt_preview',
-        'decapoda-research/llama-7b-hf',
-        'decapoda-research/llama-13b-hf',
-        'decapoda-research/llama-30b-hf',
-        'decapoda-research/llama-65b-hf',
-        'facebook/mbart-large-50-many-to-many-mmt',
-        'philschmid/bart-large-cnn-samsum',
-        'philschmid/flan-t5-base-samsum',
-        'gpt2',
-        'distilgpt2',
-        'mosaicml/mpt-7b-storywriter',
-    ],
-    'gptj': ['gptj', 'gpt4all_llama'],
-    'prompt_answer': [
-        'h2oai/h2ogpt-gm-oasst1-en-1024-20b',
-        'h2oai/h2ogpt-gm-oasst1-en-1024-12b',
-        'h2oai/h2ogpt-gm-oasst1-multilang-1024-20b',
-        'h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b',
-        'h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b-v2',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v3',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v2',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-40b-v1',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-40b-v2',
-        'h2oai/h2ogpt-gm-oasst1-en-xgen-7b-8k',
-        'h2oai/h2ogpt-gm-oasst1-multilang-xgen-7b-8k',
-        'TheBloke/h2ogpt-gm-oasst1-en-2048-falcon-40b-v2-GPTQ',
-    ],
-    'prompt_answer_openllama': [
-        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b-preview-300bt',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b-preview-300bt-v2',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b-preview-700bt',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-13b',
-    ],
-    'instruct': ['TheBloke/llama-30b-supercot-SuperHOT-8K-fp16'],  # https://huggingface.co/TheBloke/llama-30b-supercot-SuperHOT-8K-fp16#prompting
-    'instruct_with_end': ['databricks/dolly-v2-12b'],
-    'quality': [],
-    'human_bot': [
-        'h2oai/h2ogpt-oasst1-512-12b',
-        'h2oai/h2ogpt-oasst1-512-20b',
-        'h2oai/h2ogpt-oig-oasst1-256-6_9b',
-        'h2oai/h2ogpt-oig-oasst1-512-6_9b',
-        'h2oai/h2ogpt-oig-oasst1-256-6.9b',  # legacy
-        'h2oai/h2ogpt-oig-oasst1-512-6.9b',  # legacy
-        'h2oai/h2ogpt-research-oasst1-512-30b',
-        'h2oai/h2ogpt-research-oasst1-llama-65b',
-        'h2oai/h2ogpt-oasst1-falcon-40b',
-        'h2oai/h2ogpt-oig-oasst1-falcon-40b',
-    ],
-    'dai_faq': [],
-    'summarize': [],
-    'simple_instruct': ['t5-small', 't5-large', 'google/flan-t5', 'google/flan-t5-xxl', 'google/flan-ul2'],
-    'instruct_vicuna': ['AlekseyKorshuk/vicuna-7b', 'TheBloke/stable-vicuna-13B-HF', 'junelee/wizard-vicuna-13b'],
-    'human_bot_orig': ['togethercomputer/GPT-NeoXT-Chat-Base-20B'],
-    "open_assistant": ['OpenAssistant/oasst-sft-7-llama-30b-xor', 'oasst-sft-7-llama-30b'],
-    "wizard_lm": ['ehartford/WizardLM-7B-Uncensored', 'ehartford/WizardLM-13B-Uncensored'],
-    "wizard_mega": ['openaccess-ai-collective/wizard-mega-13b'],
-    "instruct_simple": ['JosephusCheung/Guanaco'],
-    "wizard_vicuna": ['ehartford/Wizard-Vicuna-13B-Uncensored'],
-    "wizard2": ['llama'],
-    "mptinstruct": ['mosaicml/mpt-30b-instruct', 'mosaicml/mpt-7b-instruct', 'mosaicml/mpt-30b-instruct'],
-    "mptchat": ['mosaicml/mpt-7b-chat', 'mosaicml/mpt-30b-chat', 'TheBloke/mpt-30B-chat-GGML'],
-    "vicuna11": ['lmsys/vicuna-33b-v1.3'],
-    "falcon": ['tiiuae/falcon-40b-instruct', 'tiiuae/falcon-40b', 'tiiuae/falcon-7b-instruct', 'tiiuae/falcon-7b'],
-    "llama2": [
-        'meta-llama/Llama-2-7b-chat-hf',
-        'meta-llama/Llama-2-13b-chat-hf',
-        'meta-llama/Llama-2-34b-chat-hf',
-        'meta-llama/Llama-2-70b-chat-hf',
-    ],
-    # could be plain, but default is correct prompt_type for default TheBloke model ggml-wizardLM-7B.q4_2.bin
-}
-if os.getenv('OPENAI_API_KEY'):
-    prompt_type_to_model_name.update({
-        "openai": ["text-davinci-003", "text-curie-001", "text-babbage-001", "text-ada-001"],
-        "openai_chat": ["gpt-3.5-turbo", "gpt-3.5-turbo-16k"],
-    })
-
-inv_prompt_type_to_model_name = {v.strip(): k for k, l in prompt_type_to_model_name.items() for v in l}
-inv_prompt_type_to_model_lower = {v.strip().lower(): k for k, l in prompt_type_to_model_name.items() for v in l}
-
-prompt_types_strings = []
-for p in PromptType:
-    prompt_types_strings.extend([p.name])
-
-prompt_types = []
-for p in PromptType:
-    prompt_types.extend([p.name, p.value, str(p.value)])
-
-
-def get_prompt(prompt_type, prompt_dict, chat, context, reduced, making_context, return_dict=False):
-    prompt_dict_error = ''
-    generates_leading_space = False
-
-    if prompt_type == PromptType.custom.name and not isinstance(prompt_dict, dict):
-        try:
-            prompt_dict = ast.literal_eval(prompt_dict)
-        except BaseException as e:
-            prompt_dict_error = str(e)
-    if prompt_dict_error:
-        promptA = None
-        promptB = None
-        PreInstruct = None
-        PreInput = ''
-        PreResponse = ''
-        terminate_response = None
-        chat_sep = ''
-        chat_turn_sep = ''
-        humanstr = ''
-        botstr = ''
-        generates_leading_space = False
-    elif prompt_type in [PromptType.custom.value, str(PromptType.custom.value),
-                         PromptType.custom.name]:
-        promptA = prompt_dict.get('promptA', '')
-        promptB = prompt_dict.get('promptB', '')
-        PreInstruct = prompt_dict.get('PreInstruct', '')
-        PreInput = prompt_dict.get('PreInput', '')
-        PreResponse = prompt_dict.get('PreResponse', '')
-        terminate_response = prompt_dict.get('terminate_response', None)
-        chat_sep = prompt_dict.get('chat_sep', '\n')
-        chat_turn_sep = prompt_dict.get('chat_turn_sep', '\n')
-        humanstr = prompt_dict.get('humanstr', '')
-        botstr = prompt_dict.get('botstr', '')
-    elif prompt_type in [PromptType.plain.value, str(PromptType.plain.value),
-                         PromptType.plain.name]:
-        promptA = promptB = PreInstruct = PreInput = PreResponse = None
-        terminate_response = []
-        chat_turn_sep = chat_sep = ''
-        # plain should have None for human/bot, so nothing truncated out, not '' that would truncate after first token
-        humanstr = None
-        botstr = None
-    elif prompt_type == 'simple_instruct':
-        promptA = promptB = PreInstruct = PreInput = PreResponse = None
-        terminate_response = []
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = None
-        botstr = None
-    elif prompt_type in [PromptType.instruct.value, str(PromptType.instruct.value),
-                         PromptType.instruct.name] + [PromptType.instruct_with_end.value,
-                                                      str(PromptType.instruct_with_end.value),
-                                                      PromptType.instruct_with_end.name]:
-        promptA = 'Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n' if not (
-                chat and reduced) else ''
-        promptB = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.\n' if not (
-                chat and reduced) else ''
-
-        PreInstruct = """
-### Instruction:
-"""
-
-        PreInput = """
-### Input:
-"""
-
-        PreResponse = """
-### Response:
-"""
-        if prompt_type in [PromptType.instruct_with_end.value, str(PromptType.instruct_with_end.value),
-                           PromptType.instruct_with_end.name]:
-            terminate_response = ['### End']
-        else:
-            terminate_response = None
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.quality.value, str(PromptType.quality.value),
-                         PromptType.quality.name]:
-        promptA = 'Write a detailed high-quality, accurate, fair, Response with about 100 words by following the Instruction as applied on the Input.\n' if not (
-                chat and reduced) else ''
-        promptB = 'Write a detailed high-quality, accurate, fair, Response with about 100 words by following the Instruction.\n' if not (
-                chat and reduced) else ''
-
-        PreInstruct = """
-### Instruction:
-"""
-
-        PreInput = """
-### Input:
-"""
-
-        PreResponse = """
-### Response:
-"""
-        terminate_response = None
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct  # first thing human says
-        botstr = PreResponse  # first thing bot says
-    elif prompt_type in [PromptType.human_bot.value, str(PromptType.human_bot.value),
-                         PromptType.human_bot.name] + [PromptType.human_bot_orig.value,
-                                                       str(PromptType.human_bot_orig.value),
-                                                       PromptType.human_bot_orig.name]:
-        human = '<human>:'
-        bot = "<bot>:"
-        if reduced or context or prompt_type in [PromptType.human_bot.value, str(PromptType.human_bot.value),
-                                                 PromptType.human_bot.name]:
-            preprompt = ''
-        else:
-            cur_date = time.strftime('%Y-%m-%d')
-            cur_time = time.strftime('%H:%M:%S %p %Z')
-
-            PRE_PROMPT = """\
-Current Date: {}
-Current Time: {}
-
-"""
-            preprompt = PRE_PROMPT.format(cur_date, cur_time)
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-
-        PreInstruct = human + ' '
-
-        PreInput = None
-
-        if making_context:
-            # when making context, want it to appear as-if LLM generated, which starts with space after :
-            PreResponse = bot + ' '
-        else:
-            # normally LLM adds space after this, because was how trained.
-            # if add space here, non-unique tokenization will often make LLM produce wrong output
-            PreResponse = bot
-
-        terminate_response = ['\n' + human, '\n' + bot, human, bot, PreResponse]
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = human  # tag before human talks
-        botstr = bot  # tag before bot talks
-        generates_leading_space = True
-    elif prompt_type in [PromptType.dai_faq.value, str(PromptType.dai_faq.value),
-                         PromptType.dai_faq.name]:
-        promptA = ''
-        promptB = 'Answer the following Driverless AI question.\n'
-
-        PreInstruct = """
-### Driverless AI frequently asked question:
-"""
-
-        PreInput = None
-
-        PreResponse = """
-### Driverless AI documentation answer:
-"""
-        terminate_response = ['\n\n']
-        chat_turn_sep = chat_sep = terminate_response
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.summarize.value, str(PromptType.summarize.value),
-                         PromptType.summarize.name]:
-        promptA = promptB = PreInput = ''
-        PreInstruct = '## Main Text\n\n'
-        PreResponse = '\n\n## Summary\n\n'
-        terminate_response = None
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.instruct_vicuna.value, str(PromptType.instruct_vicuna.value),
-                         PromptType.instruct_vicuna.name]:
-        promptA = promptB = "A chat between a curious human and an artificial intelligence assistant. " \
-                            "The assistant gives helpful, detailed, and polite answers to the human's questions." if not (
-                chat and reduced) else ''
-
-        PreInstruct = """
-### Human:
-"""
-
-        PreInput = None
-
-        PreResponse = """
-### Assistant:
-"""
-        terminate_response = [
-            '### Human:']  # but only allow terminate after prompt is found correctly, else can't terminate
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.prompt_answer.value, str(PromptType.prompt_answer.value),
-                         PromptType.prompt_answer.name]:
-        preprompt = ''
-        prompt_tokens = "<|prompt|>"
-        answer_tokens = "<|answer|>"
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = prompt_tokens
-        PreInput = None
-        PreResponse = answer_tokens
-        eos = '<|endoftext|>'  # neox eos
-        humanstr = prompt_tokens
-        botstr = answer_tokens
-        terminate_response = [humanstr, PreResponse, eos]
-        chat_sep = eos
-        chat_turn_sep = eos
-    elif prompt_type in [PromptType.prompt_answer_openllama.value, str(PromptType.prompt_answer_openllama.value),
-                         PromptType.prompt_answer_openllama.name]:
-        preprompt = ''
-        prompt_tokens = "<|prompt|>"
-        answer_tokens = "<|answer|>"
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = prompt_tokens
-        PreInput = None
-        PreResponse = answer_tokens
-        eos = '</s>'  # llama eos
-        humanstr = prompt_tokens
-        botstr = answer_tokens
-        terminate_response = [humanstr, PreResponse, eos]
-        chat_sep = eos
-        chat_turn_sep = eos
-    elif prompt_type in [PromptType.open_assistant.value, str(PromptType.open_assistant.value),
-                         PromptType.open_assistant.name]:
-        # From added_tokens.json
-        preprompt = ''
-        prompt_tokens = "<|prompter|>"
-        answer_tokens = "<|assistant|>"
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = prompt_tokens
-        PreInput = None
-        PreResponse = answer_tokens
-        pend = "<|prefix_end|>"
-        eos = "</s>"
-        humanstr = prompt_tokens
-        botstr = answer_tokens
-        terminate_response = [humanstr, PreResponse, pend, eos]
-        chat_turn_sep = chat_sep = eos
-    elif prompt_type in [PromptType.wizard_lm.value, str(PromptType.wizard_lm.value),
-                         PromptType.wizard_lm.name]:
-        # https://github.com/ehartford/WizardLM/blob/main/src/train_freeform.py
-        preprompt = ''
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = ""
-        PreInput = None
-        PreResponse = "\n\n### Response\n"
-        eos = "</s>"
-        terminate_response = [PreResponse, eos]
-        chat_turn_sep = chat_sep = eos
-        humanstr = promptA
-        botstr = PreResponse
-    elif prompt_type in [PromptType.wizard_mega.value, str(PromptType.wizard_mega.value),
-                         PromptType.wizard_mega.name]:
-        preprompt = ''
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = """
-### Instruction:
-"""
-        PreInput = None
-        PreResponse = """
-### Assistant:
-"""
-        terminate_response = [PreResponse]
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.instruct_vicuna2.value, str(PromptType.instruct_vicuna2.value),
-                         PromptType.instruct_vicuna2.name]:
-        promptA = promptB = "" if not (chat and reduced) else ''
-
-        PreInstruct = """
-HUMAN:
-"""
-
-        PreInput = None
-
-        PreResponse = """
-ASSISTANT:
-"""
-        terminate_response = [
-            'HUMAN:']  # but only allow terminate after prompt is found correctly, else can't terminate
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.instruct_vicuna3.value, str(PromptType.instruct_vicuna3.value),
-                         PromptType.instruct_vicuna3.name]:
-        promptA = promptB = "" if not (chat and reduced) else ''
-
-        PreInstruct = """
-### User:
-"""
-
-        PreInput = None
-
-        PreResponse = """
-### Assistant:
-"""
-        terminate_response = [
-            '### User:']  # but only allow terminate after prompt is found correctly, else can't terminate
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.wizard2.value, str(PromptType.wizard2.value),
-                         PromptType.wizard2.name]:
-        # https://huggingface.co/TheBloke/WizardLM-7B-uncensored-GGML
-        preprompt = """Below is an instruction that describes a task. Write a response that appropriately completes the request.""" if not (
-                chat and reduced) else ''
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = """
-### Instruction:
-"""
-        PreInput = None
-        PreResponse = """
-### Response:
-"""
-        terminate_response = [PreResponse]
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.wizard3.value, str(PromptType.wizard3.value),
-                         PromptType.wizard3.name]:
-        # https://huggingface.co/TheBloke/wizardLM-13B-1.0-GGML
-        preprompt = """A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.""" if not (
-                chat and reduced) else ''
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = """USER: """
-        PreInput = None
-        PreResponse = """ASSISTANT: """
-        terminate_response = [PreResponse]
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.wizard_vicuna.value, str(PromptType.wizard_vicuna.value),
-                         PromptType.wizard_vicuna.name]:
-        preprompt = ''
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = """USER: """
-        PreInput = None
-        PreResponse = """ASSISTANT: """
-        terminate_response = [PreResponse]
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-
-    elif prompt_type in [PromptType.instruct_simple.value, str(PromptType.instruct_simple.value),
-                         PromptType.instruct_simple.name]:
-        promptB = promptA = '' if not (chat and reduced) else ''
-
-        PreInstruct = """
-### Instruction:
-"""
-
-        PreInput = """
-### Input:
-"""
-
-        PreResponse = """
-### Response:
-"""
-        terminate_response = None
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.openai.value, str(PromptType.openai.value),
-                         PromptType.openai.name]:
-        preprompt = """The following is a conversation with an AI assistant. The assistant is helpful, creative, clever, and very friendly.""" if not (
-                chat and reduced) else ''
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = "\nHuman: "
-        PreInput = None
-        PreResponse = "\nAI:"
-        terminate_response = [PreResponse] + [" Human:", " AI:"]
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.gptj.value, str(PromptType.gptj.value),
-                         PromptType.gptj.name]:
-        preprompt = "### Instruction:\n The prompt below is a question to answer, a task to complete, or a conversation to respond to; decide which and write an appropriate response." if not (
-                chat and reduced) else ''
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = "\n### Prompt: "
-        PreInput = None
-        PreResponse = "\n### Response: "
-        terminate_response = [PreResponse] + ["Prompt:", "Response:"]
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.openai_chat.value, str(PromptType.openai_chat.value),
-                         PromptType.openai_chat.name]:
-        # prompting and termination all handled by endpoint
-        preprompt = """"""
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = ""
-        PreInput = None
-        PreResponse = ""
-        terminate_response = []
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = None
-        botstr = None
-    elif prompt_type in [PromptType.vicuna11.value, str(PromptType.vicuna11.value),
-                         PromptType.vicuna11.name]:
-        preprompt = """A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. """ if not (
-                chat and reduced) else ''
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        eos = '</s>'
-        PreInstruct = """USER: """
-        PreInput = None
-        PreResponse = """ASSISTANT:"""
-        terminate_response = [PreResponse]
-        chat_sep = ' '
-        chat_turn_sep = eos
-        humanstr = PreInstruct
-        botstr = PreResponse
-
-        if making_context:
-            # when making context, want it to appear as-if LLM generated, which starts with space after :
-            PreResponse = PreResponse + ' '
-        else:
-            # normally LLM adds space after this, because was how trained.
-            # if add space here, non-unique tokenization will often make LLM produce wrong output
-            PreResponse = PreResponse
-    elif prompt_type in [PromptType.mptinstruct.value, str(PromptType.mptinstruct.value),
-                         PromptType.mptinstruct.name]:
-        # https://huggingface.co/mosaicml/mpt-30b-instruct#formatting
-        promptA = promptB = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.\n' if not (
-                chat and reduced) else ''
-
-        PreInstruct = """
-### Instruction
-"""
-
-        PreInput = """
-### Input
-"""
-
-        PreResponse = """
-### Response
-"""
-        terminate_response = None
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.mptchat.value, str(PromptType.mptchat.value),
-                         PromptType.mptchat.name]:
-        # https://huggingface.co/TheBloke/mpt-30B-chat-GGML#prompt-template
-        promptA = promptB = """<|im_start|>system\nA conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.\n<|im_end|>""" if not (
-                chat and reduced) else ''
-
-        PreInstruct = """<|im_start|>user
-"""
-
-        PreInput = None
-
-        PreResponse = """<|im_end|><|im_start|>assistant
-"""
-        terminate_response = ['<|im_end|>']
-        chat_sep = ''
-        chat_turn_sep = '<|im_end|>'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.falcon.value, str(PromptType.falcon.value),
-                         PromptType.falcon.name]:
-        promptA = promptB = "" if not (chat and reduced) else ''
-
-        PreInstruct = """User: """
-
-        PreInput = None
-
-        PreResponse = """Assistant:"""
-        terminate_response = ['\nUser', "<|endoftext|>"]
-        chat_sep = '\n\n'
-        chat_turn_sep = '\n\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-        if making_context:
-            # when making context, want it to appear as-if LLM generated, which starts with space after :
-            PreResponse = 'Assistant: '
-        else:
-            # normally LLM adds space after this, because was how trained.
-            # if add space here, non-unique tokenization will often make LLM produce wrong output
-            PreResponse = PreResponse
-        # generates_leading_space = True
-    elif prompt_type in [PromptType.guanaco.value, str(PromptType.guanaco.value),
-                         PromptType.guanaco.name]:
-        # https://huggingface.co/TheBloke/guanaco-65B-GPTQ
-        promptA = promptB = "" if not (chat and reduced) else ''
-
-        PreInstruct = """### Human: """
-
-        PreInput = None
-
-        PreResponse = """### Assistant:"""
-        terminate_response = ['### Human:']  # but only allow terminate after prompt is found correctly, else can't terminate
-        chat_turn_sep = chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-    elif prompt_type in [PromptType.llama2.value, str(PromptType.llama2.value),
-                         PromptType.llama2.name]:
-        PreInstruct = ""
-        llama2_sys = "<<SYS>>\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\n<</SYS>>\n\n"
-        prompt = "<s>[INST] "
-        enable_sys = False  # too much safety, hurts accuracy
-        if not (chat and reduced):
-            if enable_sys:
-                promptA = promptB = prompt + llama2_sys
-            else:
-                promptA = promptB = prompt
-        else:
-            promptA = promptB = ''
-        PreInput = None
-        PreResponse = ""
-        terminate_response = ["[INST]", "</s>"]
-        chat_sep = ' [/INST]'
-        chat_turn_sep = ' </s><s>[INST] '
-        humanstr = PreInstruct
-        botstr = PreResponse
-        if making_context:
-            PreResponse += " "
-    else:
-        raise RuntimeError("No such prompt_type=%s" % prompt_type)
-
-    if isinstance(terminate_response, (tuple, list)):
-        assert '' not in terminate_response, "Bad terminate_response"
-
-    ret_dict = dict(promptA=promptA, promptB=promptB, PreInstruct=PreInstruct, PreInput=PreInput,
-                    PreResponse=PreResponse, terminate_response=terminate_response, chat_sep=chat_sep,
-                    chat_turn_sep=chat_turn_sep,
-                    humanstr=humanstr, botstr=botstr,
-                    generates_leading_space=generates_leading_space)
-
-    if return_dict:
-        return ret_dict, prompt_dict_error
-    else:
-        return tuple(list(ret_dict.values()))
-
-
-def generate_prompt(data_point, prompt_type, prompt_dict, chat, reduced, making_context):
-    context = data_point.get('context')
-    if context is None:
-        context = ''
-    instruction = data_point.get('instruction')
-    input = data_point.get('input')
-    output = data_point.get('output')
-    prompt_type = data_point.get('prompt_type', prompt_type)
-    prompt_dict = data_point.get('prompt_dict', prompt_dict)
-    assert prompt_type in prompt_types, "Bad prompt type: %s" % prompt_type
-    promptA, promptB, PreInstruct, PreInput, PreResponse, \
-        terminate_response, chat_sep, chat_turn_sep, humanstr, botstr, \
-        generates_leading_space = get_prompt(prompt_type, prompt_dict, chat,
-                                             context, reduced, making_context)
-
-    # could avoid if reduce=True, but too complex for parent functions to handle
-    prompt = context
-
-    if input and promptA:
-        prompt += f"""{promptA}"""
-    elif promptB:
-        prompt += f"""{promptB}"""
-
-    if instruction and PreInstruct is not None and input and PreInput is not None:
-        prompt += f"""{PreInstruct}{instruction}{PreInput}{input}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-    elif instruction and input and PreInstruct is None and PreInput is not None:
-        prompt += f"""{PreInput}{instruction}
-{input}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-    elif input and instruction and PreInput is None and PreInstruct is not None:
-        prompt += f"""{PreInstruct}{instruction}
-{input}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-    elif instruction and PreInstruct is not None:
-        prompt += f"""{PreInstruct}{instruction}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-    elif input and PreInput is not None:
-        prompt += f"""{PreInput}{input}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-    elif input and instruction and PreInput is not None:
-        prompt += f"""{PreInput}{instruction}{input}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-    elif input and instruction and PreInstruct is not None:
-        prompt += f"""{PreInstruct}{instruction}{input}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-    elif input and instruction:
-        # i.e. for simple_instruct
-        prompt += f"""{instruction}: {input}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-    elif input:
-        prompt += f"""{input}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-    elif instruction:
-        prompt += f"""{instruction}"""
-        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
-
-    if PreResponse is not None:
-        prompt += f"""{PreResponse}"""
-        pre_response = PreResponse  # Don't use strip
-    else:
-        pre_response = ''
-
-    if output:
-        prompt += f"""{output}"""
-
-    return prompt, pre_response, terminate_response, chat_sep, chat_turn_sep
-
-
-def inject_chatsep(prompt_type, prompt, chat_sep=None):
-    if chat_sep:
-        # only add new line if structured prompt, while 'plain' is just generation of next tokens from input
-        prompt += chat_sep
-    return prompt
-
-
-class Prompter(object):
-    def __init__(self, prompt_type, prompt_dict, debug=False, chat=False, stream_output=False, repeat_penalty=True,
-                 allowed_repeat_line_length=10):
-        self.prompt_type = prompt_type
-        self.prompt_dict = prompt_dict
-        self.debug = debug
-        self.chat = chat
-        self.stream_output = stream_output
-        self.repeat_penalty = repeat_penalty
-        self.allowed_repeat_line_length = allowed_repeat_line_length
-        self.prompt = None
-        context = ""  # not for chat context
-        reduced = False  # not for chat context
-        making_context = False  # not for chat context
-        self.promptA, self.promptB, self.PreInstruct, self.PreInput, self.PreResponse, \
-            self.terminate_response, self.chat_sep, self.chat_turn_sep, self.humanstr, self.botstr, \
-            self.generates_leading_space = \
-            get_prompt(self.prompt_type, self.prompt_dict, chat, context, reduced, making_context)
-        self.pre_response = self.PreResponse
-
-    def generate_prompt(self, data_point, reduced=None):
-        """
-        data_point['context'] is assumed to be like a system prompt or pre-conversation, not inserted after user prompt
-        :param data_point:
-        :param reduced:
-        :return:
-        """
-        reduced = data_point.get('context') not in ['', None] if reduced is None else reduced
-        making_context = False  # whether really making final prompt or just generating context
-        prompt, _, _, _, _ = generate_prompt(data_point, self.prompt_type, self.prompt_dict, self.chat, reduced,
-                                             making_context)
-        if self.debug:
-            print("prompt: %s" % prompt, flush=True)
-        # if have context, should have always reduced and only preappend promptA/B here
-        if data_point.get('context'):
-            if data_point.get('input') and self.promptA:
-                prompt = self.promptA + prompt
-            elif self.promptB:
-                prompt = self.promptB + prompt
-
-        self.prompt = prompt
-        return prompt
-
-    def get_response(self, outputs, prompt=None, sanitize_bot_response=False):
-        if isinstance(outputs, str):
-            outputs = [outputs]
-        if self.debug:
-            print("output:\n%s" % '\n\n'.join(outputs), flush=True)
-        if prompt is not None:
-            self.prompt = prompt
-
-        def clean_response(response):
-            meaningless_words = ['<pad>', '</s>', '<|endoftext|>']
-            for word in meaningless_words:
-                response = response.replace(word, "")
-            if sanitize_bot_response:
-                from better_profanity import profanity
-                response = profanity.censor(response)
-            if self.generates_leading_space and isinstance(response, str) and len(response) > 0 and response[0] == ' ':
-                response = response[1:]
-            return response
-
-        def clean_repeats(response):
-            lines = response.split('\n')
-            new_lines = []
-            [new_lines.append(line) for line in lines if
-             line not in new_lines or len(line) < self.allowed_repeat_line_length]
-            if self.debug and len(lines) != len(new_lines):
-                print("cleaned repeats: %s %s" % (len(lines), len(new_lines)), flush=True)
-            response = '\n'.join(new_lines)
-            return response
-
-        multi_output = len(outputs) > 1
-
-        for oi, output in enumerate(outputs):
-            if self.prompt_type in [PromptType.plain.value, str(PromptType.plain.value), PromptType.plain.name]:
-                output = clean_response(output)
-            elif prompt is None:
-                # then use most basic parsing like pipeline
-                if not self.botstr:
-                    pass
-                elif self.botstr in output:
-                    if self.humanstr:
-                        output = clean_response(output.split(self.botstr)[1].split(self.humanstr)[0])
-                    else:
-                        # i.e. use after bot but only up to next bot
-                        output = clean_response(output.split(self.botstr)[1].split(self.botstr)[0])
-                else:
-                    # output = clean_response(output)
-                    # assume just not printed yet
-                    output = ""
-            else:
-                # find first instance of prereponse
-                # prompt sometimes has odd characters, that mutate length,
-                # so can't go by length alone
-                if self.pre_response:
-                    outputi = output.find(prompt)
-                    if outputi >= 0:
-                        output = output[outputi + len(prompt):]
-                        allow_terminate = True
-                    else:
-                        # subtraction is risky due to space offsets sometimes, so only do if necessary
-                        output = output[len(prompt) - len(self.pre_response):]
-                        # [1] to avoid repeated pre_response, just take first (after prompt - pre_response for chat)
-                        if self.pre_response in output:
-                            output = output.split(self.pre_response)[1]
-                            allow_terminate = True
-                        else:
-                            if output:
-                                print("Failure of parsing or not enough output yet: %s" % output, flush=True)
-                            allow_terminate = False
-                else:
-                    allow_terminate = True
-                    output = output[len(prompt):]
-                # clean after subtract prompt out, so correct removal of pre_response
-                output = clean_response(output)
-                if self.repeat_penalty:
-                    output = clean_repeats(output)
-                if self.terminate_response and allow_terminate:
-                    finds = []
-                    for term in self.terminate_response:
-                        finds.append(output.find(term))
-                    finds = [x for x in finds if x >= 0]
-                    if len(finds) > 0:
-                        termi = finds[0]
-                        output = output[:termi]
-                    else:
-                        output = output
-            if multi_output:
-                # prefix with output counter
-                output = "\n=========== Output %d\n\n" % (1 + oi) + output
-                if oi > 0:
-                    # post fix outputs with seperator
-                    output += '\n'
-            output = self.fix_text(self.prompt_type, output)
-            outputs[oi] = output
-        # join all outputs, only one extra new line between outputs
-        output = '\n'.join(outputs)
-        if self.debug:
-            print("outputclean:\n%s" % '\n\n'.join(outputs), flush=True)
-        return output
-
-    @staticmethod
-    def fix_text(prompt_type1, text1):
-        if prompt_type1 == 'human_bot':
-            # hack bug in vLLM with stopping, stops right, but doesn't return last token
-            hfix = '<human'
-            if text1.endswith(hfix):
-                text1 = text1[:-len(hfix)]
-        return text1
-

stopping.py

@@ -1,78 +0,0 @@
-import torch
-from transformers import StoppingCriteria, StoppingCriteriaList
-
-from enums import PromptType
-
-
-class StoppingCriteriaSub(StoppingCriteria):
-
-    def __init__(self, stops=[], encounters=[], device="cuda", model_max_length=None):
-        super().__init__()
-        assert len(stops) % len(encounters) == 0, "Number of stops and encounters must match"
-        self.encounters = encounters
-        self.stops = [stop.to(device) for stop in stops]
-        self.num_stops = [0] * len(stops)
-        self.model_max_length = model_max_length
-
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        for stopi, stop in enumerate(self.stops):
-            if torch.all((stop == input_ids[0][-len(stop):])).item():
-                self.num_stops[stopi] += 1
-                if self.num_stops[stopi] >= self.encounters[stopi % len(self.encounters)]:
-                    # print("Stopped", flush=True)
-                    return True
-        if self.model_max_length is not None and input_ids[0].shape[0] >= self.model_max_length:
-            # critical limit
-            return True
-        # print("Tokens: %s" % input_ids[0].cpu().numpy(), flush=True)
-        # print("Stop Tokens: %s" % [x.cpu().numpy() for x in self.stops], flush=True)
-        return False
-
-
-def get_stopping(prompt_type, prompt_dict, tokenizer, device, human='<human>:', bot="<bot>:", model_max_length=None):
-    # FIXME: prompt_dict unused currently
-    if prompt_type in [PromptType.human_bot.name, PromptType.instruct_vicuna.name, PromptType.instruct_with_end.name]:
-        if prompt_type == PromptType.human_bot.name:
-            # encounters = [prompt.count(human) + 1, prompt.count(bot) + 1]
-            # stopping only starts once output is beyond prompt
-            # 1 human is enough to trigger, but need 2 bots, because very first view back will be bot we added
-            stop_words = [human, bot, '\n' + human, '\n' + bot]
-            encounters = [1, 2]
-        elif prompt_type == PromptType.instruct_vicuna.name:
-            # even below is not enough, generic strings and many ways to encode
-            stop_words = [
-                '### Human:',
-                """
-### Human:""",
-                """
-### Human:
-""",
-                '### Assistant:',
-                """
-### Assistant:""",
-                """
-### Assistant:
-""",
-            ]
-            encounters = [1, 2]
-        else:
-            # some instruct prompts have this as end, doesn't hurt to stop on it since not common otherwise
-            stop_words = ['### End']
-            encounters = [1]
-        stop_words_ids = [
-            tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze() for stop_word in stop_words]
-        # handle single token case
-        stop_words_ids = [x if len(x.shape) > 0 else torch.tensor([x]) for x in stop_words_ids]
-        stop_words_ids = [x for x in stop_words_ids if x.shape[0] > 0]
-        # avoid padding in front of tokens
-        if tokenizer._pad_token:  # use hidden variable to avoid annoying properly logger bug
-            stop_words_ids = [x[1:] if x[0] == tokenizer.pad_token_id and len(x) > 1 else x for x in stop_words_ids]
-        # handle fake \n added
-        stop_words_ids = [x[1:] if y[0] == '\n' else x for x, y in zip(stop_words_ids, stop_words)]
-        # build stopper
-        stopping_criteria = StoppingCriteriaList(
-            [StoppingCriteriaSub(stops=stop_words_ids, encounters=encounters, device=device,
-                                 model_max_length=model_max_length)])
-    else:
-        stopping_criteria = StoppingCriteriaList()
-    return stopping_criteria

utils.py

@@ -1,1080 +0,0 @@
-import contextlib
-import functools
-import hashlib
-import inspect
-import os
-import gc
-import pathlib
-import pickle
-import random
-import shutil
-import subprocess
-import sys
-import threading
-import time
-import traceback
-import zipfile
-from datetime import datetime
-
-import filelock
-import requests, uuid
-from typing import Tuple, Callable, Dict
-from tqdm.auto import tqdm
-from joblib import Parallel
-from concurrent.futures import ProcessPoolExecutor
-import numpy as np
-import pandas as pd
-
-
-def set_seed(seed: int):
-    """
-    Sets the seed of the entire notebook so results are the same every time we run.
-    This is for REPRODUCIBILITY.
-    """
-    import torch
-    np.random.seed(seed)
-    random_state = np.random.RandomState(seed)
-    random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-    os.environ['PYTHONHASHSEED'] = str(seed)
-    return random_state
-
-
-def flatten_list(lis):
-    """Given a list, possibly nested to any level, return it flattened."""
-    new_lis = []
-    for item in lis:
-        if type(item) == type([]):
-            new_lis.extend(flatten_list(item))
-        else:
-            new_lis.append(item)
-    return new_lis
-
-
-def clear_torch_cache():
-    import torch
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
-        gc.collect()
-
-
-def ping():
-    try:
-        print('Ping: %s' % str(datetime.now()), flush=True)
-    except AttributeError:
-        # some programs wrap print and will fail with flush passed
-        pass
-
-
-def ping_gpu():
-    try:
-        print('Ping_GPU: %s %s' % (str(datetime.now()), system_info()), flush=True)
-    except AttributeError:
-        # some programs wrap print and will fail with flush passed
-        pass
-    try:
-        ping_gpu_memory()
-    except Exception as e:
-        print('Ping_GPU memory failure: %s' % str(e), flush=True)
-
-
-def ping_gpu_memory():
-    from models.gpu_mem_track import MemTracker
-    gpu_tracker = MemTracker()  # define a GPU tracker
-    from torch.cuda import memory_summary
-    gpu_tracker.track()
-
-
-def get_torch_allocated():
-    import torch
-    return torch.cuda.memory_allocated()
-
-
-def get_device():
-    import torch
-    if torch.cuda.is_available():
-        device = "cuda"
-    elif torch.backends.mps.is_built():
-        device = "mps"
-    else:
-        device = "cpu"
-
-    return device
-
-
-def system_info():
-    import psutil
-
-    system = {}
-    # https://stackoverflow.com/questions/48951136/plot-multiple-graphs-in-one-plot-using-tensorboard
-    # https://arshren.medium.com/monitoring-your-devices-in-python-5191d672f749
-    try:
-        temps = psutil.sensors_temperatures(fahrenheit=False)
-        if 'coretemp' in temps:
-            coretemp = temps['coretemp']
-            temp_dict = {k.label: k.current for k in coretemp}
-            for k, v in temp_dict.items():
-                system['CPU_C/%s' % k] = v
-    except AttributeError:
-        pass
-
-    # https://github.com/gpuopenanalytics/pynvml/blob/master/help_query_gpu.txt
-    try:
-        from pynvml.smi import nvidia_smi
-        nvsmi = nvidia_smi.getInstance()
-
-        gpu_power_dict = {'W_gpu%d' % i: x['power_readings']['power_draw'] for i, x in
-                          enumerate(nvsmi.DeviceQuery('power.draw')['gpu'])}
-        for k, v in gpu_power_dict.items():
-            system['GPU_W/%s' % k] = v
-
-        gpu_temp_dict = {'C_gpu%d' % i: x['temperature']['gpu_temp'] for i, x in
-                         enumerate(nvsmi.DeviceQuery('temperature.gpu')['gpu'])}
-        for k, v in gpu_temp_dict.items():
-            system['GPU_C/%s' % k] = v
-
-        gpu_memory_free_dict = {'MiB_gpu%d' % i: x['fb_memory_usage']['free'] for i, x in
-                                enumerate(nvsmi.DeviceQuery('memory.free')['gpu'])}
-        gpu_memory_total_dict = {'MiB_gpu%d' % i: x['fb_memory_usage']['total'] for i, x in
-                                 enumerate(nvsmi.DeviceQuery('memory.total')['gpu'])}
-        gpu_memory_frac_dict = {k: gpu_memory_free_dict[k] / gpu_memory_total_dict[k] for k in gpu_memory_total_dict}
-        for k, v in gpu_memory_frac_dict.items():
-            system[f'GPU_M/%s' % k] = v
-    except (KeyError, ModuleNotFoundError):
-        pass
-    system['hash'] = get_githash()
-
-    return system
-
-
-def system_info_print():
-    try:
-        df = pd.DataFrame.from_dict(system_info(), orient='index')
-        # avoid slamming GPUs
-        time.sleep(1)
-        return df.to_markdown()
-    except Exception as e:
-        return "Error: %s" % str(e)
-
-
-def zip_data(root_dirs=None, zip_file=None, base_dir='./', fail_any_exception=False):
-    try:
-        return _zip_data(zip_file=zip_file, base_dir=base_dir, root_dirs=root_dirs)
-    except Exception as e:
-        traceback.print_exc()
-        print('Exception in zipping: %s' % str(e))
-        if not fail_any_exception:
-            raise
-
-
-def _zip_data(root_dirs=None, zip_file=None, base_dir='./'):
-    if isinstance(root_dirs, str):
-        root_dirs = [root_dirs]
-    if zip_file is None:
-        datetime_str = str(datetime.now()).replace(" ", "_").replace(":", "_")
-        host_name = os.getenv('HF_HOSTNAME', 'emptyhost')
-        zip_file = "data_%s_%s.zip" % (datetime_str, host_name)
-    assert root_dirs is not None
-    if not os.path.isdir(os.path.dirname(zip_file)) and os.path.dirname(zip_file):
-        os.makedirs(os.path.dirname(zip_file), exist_ok=True)
-    with zipfile.ZipFile(zip_file, "w") as expt_zip:
-        for root_dir in root_dirs:
-            if root_dir is None:
-                continue
-            for root, d, files in os.walk(root_dir):
-                for file in files:
-                    file_to_archive = os.path.join(root, file)
-                    assert os.path.exists(file_to_archive)
-                    path_to_archive = os.path.relpath(file_to_archive, base_dir)
-                    expt_zip.write(filename=file_to_archive, arcname=path_to_archive)
-    return zip_file, zip_file
-
-
-def save_generate_output(prompt=None, output=None, base_model=None, save_dir=None, where_from='unknown where from',
-                         extra_dict={}):
-    try:
-        return _save_generate_output(prompt=prompt, output=output, base_model=base_model, save_dir=save_dir,
-                                     where_from=where_from, extra_dict=extra_dict)
-    except Exception as e:
-        traceback.print_exc()
-        print('Exception in saving: %s' % str(e))
-
-
-def _save_generate_output(prompt=None, output=None, base_model=None, save_dir=None, where_from='unknown where from',
-                          extra_dict={}):
-    """
-    Save conversation to .json, row by row.
-    json_file_path is path to final JSON file. If not in ., then will attempt to make directories.
-    Appends if file exists
-    """
-    prompt = '<not set>' if prompt is None else prompt
-    output = '<not set>' if output is None else output
-    assert save_dir, "save_dir must be provided"
-    if os.path.exists(save_dir) and not os.path.isdir(save_dir):
-        raise RuntimeError("save_dir already exists and is not a directory!")
-    os.makedirs(save_dir, exist_ok=True)
-    import json
-    dict_to_save = dict(prompt=prompt, text=output, time=time.ctime(), base_model=base_model, where_from=where_from)
-    dict_to_save.update(extra_dict)
-    with filelock.FileLock("save_dir.lock"):
-        # lock logging in case have concurrency
-        with open(os.path.join(save_dir, "history.json"), "a") as f:
-            # just add [ at start, and ] at end, and have proper JSON dataset
-            f.write(
-                "  " + json.dumps(
-                    dict_to_save
-                ) + ",\n"
-            )
-
-
-def s3up(filename):
-    try:
-        return _s3up(filename)
-    except Exception as e:
-        traceback.print_exc()
-        print('Exception for file %s in s3up: %s' % (filename, str(e)))
-        return "Failed to upload %s: Error: %s" % (filename, str(e))
-
-
-def _s3up(filename):
-    import boto3
-
-    aws_access_key_id = os.getenv('AWS_SERVER_PUBLIC_KEY')
-    aws_secret_access_key = os.getenv('AWS_SERVER_SECRET_KEY')
-    bucket = os.getenv('AWS_BUCKET')
-    assert aws_access_key_id, "Set AWS key"
-    assert aws_secret_access_key, "Set AWS secret"
-    assert bucket, "Set AWS Bucket"
-
-    s3 = boto3.client('s3',
-                      aws_access_key_id=os.getenv('AWS_SERVER_PUBLIC_KEY'),
-                      aws_secret_access_key=os.getenv('AWS_SERVER_SECRET_KEY'),
-                      )
-    ret = s3.upload_file(
-        Filename=filename,
-        Bucket=os.getenv('AWS_BUCKET'),
-        Key=filename,
-    )
-    if ret in [None, '']:
-        return "Successfully uploaded %s" % filename
-
-
-def get_githash():
-    try:
-        githash = subprocess.run(['git', 'rev-parse', 'HEAD'], stdout=subprocess.PIPE).stdout.decode('utf-8')[0:-1]
-    except:
-        githash = ''
-    return githash
-
-
-def copy_code(run_id):
-    """
-    copy code to track changes
-    :param run_id:
-    :return:
-    """
-    rnd_num = str(random.randint(0, 2 ** 31))
-    run_id = 'run_' + str(run_id)
-    os.makedirs(run_id, exist_ok=True)
-    me_full = os.path.join(pathlib.Path(__file__).parent.resolve(), __file__)
-    me_file = os.path.basename(__file__)
-    new_me = os.path.join(run_id, me_file + '_' + get_githash())
-    if os.path.isfile(new_me):
-        new_me = os.path.join(run_id, me_file + '_' + get_githash() + '_' + rnd_num)
-        shutil.copy(me_full, new_me)
-    else:
-        shutil.copy(me_full, new_me)
-
-
-class NullContext(threading.local):
-    """No-op context manager, executes block without doing any additional processing.
-
-    Used as a stand-in if a particular block of code is only sometimes
-    used with a normal context manager:
-    """
-
-    def __init__(self, *args, **kwargs):
-        pass
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, exc_traceback):
-        self.finally_act()
-
-    def finally_act(self):
-        pass
-
-
-def wrapped_partial(func, *args, **kwargs):
-    """
-    Give partial properties of normal function, like __name__ attribute etc.
-    :param func:
-    :param args:
-    :param kwargs:
-    :return:
-    """
-    partial_func = functools.partial(func, *args, **kwargs)
-    functools.update_wrapper(partial_func, func)
-    return partial_func
-
-
-class ThreadException(Exception):
-    pass
-
-
-class EThread(threading.Thread):
-    # Function that raises the custom exception
-    def __init__(self, group=None, target=None, name=None,
-                 args=(), kwargs=None, *, daemon=None, streamer=None, bucket=None):
-        self.bucket = bucket
-        self.streamer = streamer
-        self.exc = None
-        self._return = None
-        super().__init__(group=group, target=target, name=name, args=args, kwargs=kwargs, daemon=daemon)
-
-    def run(self):
-        # Variable that stores the exception, if raised by someFunction
-        try:
-            if self._target is not None:
-                self._return = self._target(*self._args, **self._kwargs)
-        except BaseException as e:
-            print("thread exception: %s" % str(sys.exc_info()))
-            self.bucket.put(sys.exc_info())
-            self.exc = e
-            if self.streamer:
-                print("make stop: %s" % str(sys.exc_info()), flush=True)
-                self.streamer.do_stop = True
-        finally:
-            # Avoid a refcycle if the thread is running a function with
-            # an argument that has a member that points to the thread.
-            del self._target, self._args, self._kwargs
-
-    def join(self, timeout=None):
-        threading.Thread.join(self)
-        # Since join() returns in caller thread
-        # we re-raise the caught exception
-        # if any was caught
-        if self.exc:
-            raise self.exc
-        return self._return
-
-
-def import_matplotlib():
-    import matplotlib
-    matplotlib.use('agg')
-    # KEEP THESE HERE! START
-    import matplotlib.pyplot as plt
-    import pandas as pd
-    # to avoid dlopen deadlock in fork
-    import pandas.core.computation.expressions as pd_expressions
-    import pandas._libs.groupby as pd_libgroupby
-    import pandas._libs.reduction as pd_libreduction
-    import pandas.core.algorithms as pd_algorithms
-    import pandas.core.common as pd_com
-    import numpy as np
-    # KEEP THESE HERE! END
-
-
-def get_sha(value):
-    return hashlib.md5(str(value).encode('utf-8')).hexdigest()
-
-
-def sanitize_filename(name):
-    """
-    Sanitize file *base* names.
-    :param name: name to sanitize
-    :return:
-    """
-    bad_chars = ['[', ']', ',', '/', '\\', '\\w', '\\s', '-', '+', '\"', '\'', '>', '<', ' ', '=', ')', '(', ':', '^']
-    for char in bad_chars:
-        name = name.replace(char, "_")
-
-    length = len(name)
-    file_length_limit = 250  # bit smaller than 256 for safety
-    sha_length = 32
-    real_length_limit = file_length_limit - (sha_length + 2)
-    if length > file_length_limit:
-        sha = get_sha(name)
-        half_real_length_limit = max(1, int(real_length_limit / 2))
-        name = name[0:half_real_length_limit] + "_" + sha + "_" + name[length - half_real_length_limit:length]
-
-    return name
-
-
-def shutil_rmtree(*args, **kwargs):
-    return shutil.rmtree(*args, **kwargs)
-
-
-def remove(path: str):
-    try:
-        if path is not None and os.path.exists(path):
-            if os.path.isdir(path):
-                shutil_rmtree(path, ignore_errors=True)
-            else:
-                with contextlib.suppress(FileNotFoundError):
-                    os.remove(path)
-    except:
-        pass
-
-
-def makedirs(path, exist_ok=True):
-    """
-    Avoid some inefficiency in os.makedirs()
-    :param path:
-    :param exist_ok:
-    :return:
-    """
-    if os.path.isdir(path) and os.path.exists(path):
-        assert exist_ok, "Path already exists"
-        return path
-    os.makedirs(path, exist_ok=exist_ok)
-
-
-def atomic_move_simple(src, dst):
-    try:
-        shutil.move(src, dst)
-    except (shutil.Error, FileExistsError):
-        pass
-    remove(src)
-
-
-def download_simple(url, dest=None, print_func=None):
-    if print_func is not None:
-        print_func("BEGIN get url %s" % str(url))
-    if url.startswith("file://"):
-        from requests_file import FileAdapter
-        s = requests.Session()
-        s.mount('file://', FileAdapter())
-        url_data = s.get(url, stream=True)
-    else:
-        url_data = requests.get(url, stream=True)
-    if dest is None:
-        dest = os.path.basename(url)
-    if url_data.status_code != requests.codes.ok:
-        msg = "Cannot get url %s, code: %s, reason: %s" % (
-            str(url),
-            str(url_data.status_code),
-            str(url_data.reason),
-        )
-        raise requests.exceptions.RequestException(msg)
-    url_data.raw.decode_content = True
-    makedirs(os.path.dirname(dest), exist_ok=True)
-    uuid_tmp = str(uuid.uuid4())[:6]
-    dest_tmp = dest + "_dl_" + uuid_tmp + ".tmp"
-    with open(dest_tmp, "wb") as f:
-        shutil.copyfileobj(url_data.raw, f)
-    atomic_move_simple(dest_tmp, dest)
-    if print_func is not None:
-        print_func("END get url %s" % str(url))
-
-
-def download(url, dest=None, dest_path=None):
-    if dest_path is not None:
-        dest = os.path.join(dest_path, os.path.basename(url))
-        if os.path.isfile(dest):
-            print("already downloaded %s -> %s" % (url, dest))
-            return dest
-    elif dest is not None:
-        if os.path.exists(dest):
-            print("already downloaded %s -> %s" % (url, dest))
-            return dest
-    else:
-        uuid_tmp = "dl2_" + str(uuid.uuid4())[:6]
-        dest = uuid_tmp + os.path.basename(url)
-
-    print("downloading %s to %s" % (url, dest))
-
-    if url.startswith("file://"):
-        from requests_file import FileAdapter
-        s = requests.Session()
-        s.mount('file://', FileAdapter())
-        url_data = s.get(url, stream=True)
-    else:
-        url_data = requests.get(url, stream=True)
-
-    if url_data.status_code != requests.codes.ok:
-        msg = "Cannot get url %s, code: %s, reason: %s" % (
-            str(url), str(url_data.status_code), str(url_data.reason))
-        raise requests.exceptions.RequestException(msg)
-    url_data.raw.decode_content = True
-    dirname = os.path.dirname(dest)
-    if dirname != "" and not os.path.isdir(dirname):
-        makedirs(os.path.dirname(dest), exist_ok=True)
-    uuid_tmp = "dl3_" + str(uuid.uuid4())[:6]
-    dest_tmp = dest + "_" + uuid_tmp + ".tmp"
-    with open(dest_tmp, 'wb') as f:
-        shutil.copyfileobj(url_data.raw, f)
-    try:
-        shutil.move(dest_tmp, dest)
-    except FileExistsError:
-        pass
-    remove(dest_tmp)
-    return dest
-
-
-def get_url(x, from_str=False, short_name=False):
-    if not from_str:
-        source = x.metadata['source']
-    else:
-        source = x
-    if short_name:
-        source_name = get_short_name(source)
-    else:
-        source_name = source
-    if source.startswith('http://') or source.startswith('https://'):
-        return """<a href="%s" target="_blank"  rel="noopener noreferrer">%s</a>""" % (
-            source, source_name)
-    else:
-        return """<a href="file/%s" target="_blank"  rel="noopener noreferrer">%s</a>""" % (
-            source, source_name)
-
-
-def get_short_name(name, maxl=50):
-    if name is None:
-        return ''
-    length = len(name)
-    if length > maxl:
-        allow_length = maxl - 3
-        half_allowed = max(1, int(allow_length / 2))
-        name = name[0:half_allowed] + "..." + name[length - half_allowed:length]
-    return name
-
-
-def cuda_vis_check(total_gpus):
-    """Helper function to count GPUs by environment variable
-    Stolen from Jon's h2o4gpu utils
-    """
-    cudavis = os.getenv("CUDA_VISIBLE_DEVICES")
-    which_gpus = []
-    if cudavis is not None:
-        # prune away white-space, non-numerics,
-        # except commas for simple checking
-        cudavis = "".join(cudavis.split())
-        import re
-        cudavis = re.sub("[^0-9,]", "", cudavis)
-
-        lencudavis = len(cudavis)
-        if lencudavis == 0:
-            total_gpus = 0
-        else:
-            total_gpus = min(
-                total_gpus,
-                os.getenv("CUDA_VISIBLE_DEVICES").count(",") + 1)
-            which_gpus = os.getenv("CUDA_VISIBLE_DEVICES").split(",")
-            which_gpus = [int(x) for x in which_gpus]
-    else:
-        which_gpus = list(range(0, total_gpus))
-
-    return total_gpus, which_gpus
-
-
-def get_ngpus_vis(raise_if_exception=True):
-    ngpus_vis1 = 0
-
-    shell = False
-    if shell:
-        cmd = "nvidia-smi -L 2> /dev/null"
-    else:
-        cmd = ["nvidia-smi", "-L"]
-
-    try:
-        timeout = 5 * 3
-        o = subprocess.check_output(cmd, shell=shell, timeout=timeout)
-        lines = o.decode("utf-8").splitlines()
-        ngpus_vis1 = 0
-        for line in lines:
-            if 'Failed to initialize NVML' not in line:
-                ngpus_vis1 += 1
-    except (FileNotFoundError, subprocess.CalledProcessError, OSError):
-        # GPU systems might not have nvidia-smi, so can't fail
-        pass
-    except subprocess.TimeoutExpired as e:
-        print('Failed get_ngpus_vis: %s' % str(e))
-        if raise_if_exception:
-            raise
-
-    ngpus_vis1, which_gpus = cuda_vis_check(ngpus_vis1)
-    return ngpus_vis1
-
-
-def get_mem_gpus(raise_if_exception=True, ngpus=None):
-    totalmem_gpus1 = 0
-    usedmem_gpus1 = 0
-    freemem_gpus1 = 0
-
-    if ngpus == 0:
-        return totalmem_gpus1, usedmem_gpus1, freemem_gpus1
-
-    try:
-        cmd = "nvidia-smi -q 2> /dev/null | grep -A 3 'FB Memory Usage'"
-        o = subprocess.check_output(cmd, shell=True, timeout=15)
-        lines = o.decode("utf-8").splitlines()
-        for line in lines:
-            if 'Total' in line:
-                totalmem_gpus1 += int(line.split()[2]) * 1024 ** 2
-            if 'Used' in line:
-                usedmem_gpus1 += int(line.split()[2]) * 1024 ** 2
-            if 'Free' in line:
-                freemem_gpus1 += int(line.split()[2]) * 1024 ** 2
-    except (FileNotFoundError, subprocess.CalledProcessError, OSError):
-        # GPU systems might not have nvidia-smi, so can't fail
-        pass
-    except subprocess.TimeoutExpired as e:
-        print('Failed get_mem_gpus: %s' % str(e))
-        if raise_if_exception:
-            raise
-
-    return totalmem_gpus1, usedmem_gpus1, freemem_gpus1
-
-
-class ForkContext(threading.local):
-    """
-        Set context for forking
-        Ensures state is returned once done
-    """
-
-    def __init__(self, args=None, kwargs=None, forkdata_capable=True):
-        """
-        :param args:
-        :param kwargs:
-        :param forkdata_capable: whether fork is forkdata capable and will use copy-on-write forking of args/kwargs
-        """
-        self.forkdata_capable = forkdata_capable
-        if self.forkdata_capable:
-            self.has_args = args is not None
-            self.has_kwargs = kwargs is not None
-            forkdatacontext.args = args
-            forkdatacontext.kwargs = kwargs
-        else:
-            self.has_args = False
-            self.has_kwargs = False
-
-    def __enter__(self):
-        try:
-            # flush all outputs so doesn't happen during fork -- don't print/log inside ForkContext contexts!
-            sys.stdout.flush()
-            sys.stderr.flush()
-        except BaseException as e:
-            # exit not called if exception, and don't want to leave forkdatacontext filled in that case
-            print("ForkContext failure on enter: %s" % str(e))
-            self.finally_act()
-            raise
-        return self
-
-    def __exit__(self, exc_type, exc_value, exc_traceback):
-        self.finally_act()
-
-    def finally_act(self):
-        """
-            Done when exception hit or exit is reached in context
-            first reset forkdatacontext as crucial to have reset even if later 2 calls fail
-        :return: None
-        """
-        if self.forkdata_capable and (self.has_args or self.has_kwargs):
-            forkdatacontext._reset()
-
-
-class _ForkDataContext(threading.local):
-    def __init__(
-            self,
-            args=None,
-            kwargs=None,
-    ):
-        """
-        Global context for fork to carry data to subprocess instead of relying upon copy/pickle/serialization
-
-        :param args: args
-        :param kwargs: kwargs
-        """
-        assert isinstance(args, (tuple, type(None)))
-        assert isinstance(kwargs, (dict, type(None)))
-        self.__args = args
-        self.__kwargs = kwargs
-
-    @property
-    def args(self) -> Tuple:
-        """returns args"""
-        return self.__args
-
-    @args.setter
-    def args(self, args):
-        if self.__args is not None:
-            raise AttributeError(
-                "args cannot be overwritten: %s %s" % (str(self.__args), str(self.__kwargs))
-            )
-
-        self.__args = args
-
-    @property
-    def kwargs(self) -> Dict:
-        """returns kwargs"""
-        return self.__kwargs
-
-    @kwargs.setter
-    def kwargs(self, kwargs):
-        if self.__kwargs is not None:
-            raise AttributeError(
-                "kwargs cannot be overwritten: %s %s" % (str(self.__args), str(self.__kwargs))
-            )
-
-        self.__kwargs = kwargs
-
-    def _reset(self):
-        """Reset fork arg-kwarg context to default values"""
-        self.__args = None
-        self.__kwargs = None
-
-    def get_args_kwargs(self, func, args, kwargs) -> Tuple[Callable, Tuple, Dict]:
-        if self.__args:
-            args = self.__args[1:]
-            if not func:
-                assert len(self.__args) > 0, "if have no func, must have in args"
-                func = self.__args[0]  # should always be there
-        if self.__kwargs:
-            kwargs = self.__kwargs
-        try:
-            return func, args, kwargs
-        finally:
-            forkdatacontext._reset()
-
-    @staticmethod
-    def get_args_kwargs_for_traced_func(func, args, kwargs):
-        """
-        Return args/kwargs out of forkdatacontext when using copy-on-write way of passing args/kwargs
-        :param func: actual function ran by _traced_func, which itself is directly what mppool treats as function
-        :param args:
-        :param kwargs:
-        :return: func, args, kwargs from forkdatacontext if used, else originals
-        """
-        # first 3 lines are debug
-        func_was_None = func is None
-        args_was_None_or_empty = args is None or len(args) == 0
-        kwargs_was_None_or_empty = kwargs is None or len(kwargs) == 0
-
-        forkdatacontext_args_was_None = forkdatacontext.args is None
-        forkdatacontext_kwargs_was_None = forkdatacontext.kwargs is None
-        func, args, kwargs = forkdatacontext.get_args_kwargs(func, args, kwargs)
-        using_forkdatacontext = func_was_None and func is not None  # pulled func out of forkdatacontext.__args[0]
-        assert forkdatacontext.args is None, "forkdatacontext.args should be None after get_args_kwargs"
-        assert forkdatacontext.kwargs is None, "forkdatacontext.kwargs should be None after get_args_kwargs"
-
-        proc_type = kwargs.get('proc_type', 'SUBPROCESS')
-        if using_forkdatacontext:
-            assert proc_type == "SUBPROCESS" or proc_type == "SUBPROCESS"
-        if proc_type == "NORMAL":
-            assert forkdatacontext_args_was_None, "if no fork, expect forkdatacontext.args None entering _traced_func"
-            assert forkdatacontext_kwargs_was_None, "if no fork, expect forkdatacontext.kwargs None entering _traced_func"
-        assert func is not None, "function should not be None, indicates original args[0] was None or args was None"
-
-        return func, args, kwargs
-
-
-forkdatacontext = _ForkDataContext()
-
-
-def _traced_func(func, *args, **kwargs):
-    func, args, kwargs = forkdatacontext.get_args_kwargs_for_traced_func(func, args, kwargs)
-    return func(*args, **kwargs)
-
-
-def call_subprocess_onetask(func, args=None, kwargs=None):
-    import platform
-    if platform.system() in ['Darwin', 'Windows']:
-        return func(*args, **kwargs)
-    if isinstance(args, list):
-        args = tuple(args)
-    if args is None:
-        args = ()
-    if kwargs is None:
-        kwargs = {}
-    args = list(args)
-    args = [func] + args
-    args = tuple(args)
-    with ForkContext(args=args, kwargs=kwargs):
-        args = (None,)
-        kwargs = {}
-        with ProcessPoolExecutor(max_workers=1) as executor:
-            future = executor.submit(_traced_func, *args, **kwargs)
-            return future.result()
-
-
-class ProgressParallel(Parallel):
-    def __init__(self, use_tqdm=True, total=None, *args, **kwargs):
-        self._use_tqdm = use_tqdm
-        self._total = total
-        super().__init__(*args, **kwargs)
-
-    def __call__(self, *args, **kwargs):
-        with tqdm(disable=not self._use_tqdm, total=self._total) as self._pbar:
-            return Parallel.__call__(self, *args, **kwargs)
-
-    def print_progress(self):
-        if self._total is None:
-            self._pbar.total = self.n_dispatched_tasks
-        self._pbar.n = self.n_completed_tasks
-        self._pbar.refresh()
-
-
-def get_kwargs(func, exclude_names=None, **kwargs):
-    func_names = list(inspect.signature(func).parameters)
-    missing_kwargs = [x for x in func_names if x not in kwargs]
-    if exclude_names:
-        for k in exclude_names:
-            if k in missing_kwargs:
-                missing_kwargs.remove(k)
-            if k in func_names:
-                func_names.remove(k)
-    assert not missing_kwargs, "Missing %s" % missing_kwargs
-    kwargs = {k: v for k, v in kwargs.items() if k in func_names}
-    return kwargs
-
-
-import pkg_resources
-
-have_faiss = False
-
-try:
-    assert pkg_resources.get_distribution('faiss') is not None
-    have_faiss = True
-except (pkg_resources.DistributionNotFound, AssertionError):
-    pass
-try:
-    assert pkg_resources.get_distribution('faiss_gpu') is not None
-    have_faiss = True
-except (pkg_resources.DistributionNotFound, AssertionError):
-    pass
-try:
-    assert pkg_resources.get_distribution('faiss_cpu') is not None
-    have_faiss = True
-except (pkg_resources.DistributionNotFound, AssertionError):
-    pass
-
-
-def hash_file(file):
-    try:
-        import hashlib
-
-        # BUF_SIZE is totally arbitrary, change for your app!
-        BUF_SIZE = 65536  # lets read stuff in 64kb chunks!
-
-        md5 = hashlib.md5()
-        # sha1 = hashlib.sha1()
-
-        with open(file, 'rb') as f:
-            while True:
-                data = f.read(BUF_SIZE)
-                if not data:
-                    break
-                md5.update(data)
-                # sha1.update(data)
-    except BaseException as e:
-        print("Cannot hash %s due to %s" % (file, str(e)))
-        traceback.print_exc()
-        md5 = None
-    return md5.hexdigest()
-
-
-def start_faulthandler():
-    # If hit server or any subprocess with signal SIGUSR1, it'll print out all threads stack trace, but wont't quit or coredump
-    # If more than one fork tries to write at same time, then looks corrupted.
-    import faulthandler
-
-    # SIGUSR1 in h2oai/__init__.py as well
-    faulthandler.enable()
-    if hasattr(faulthandler, 'register'):
-        # windows/mac
-        import signal
-        faulthandler.register(signal.SIGUSR1)
-
-
-def get_hf_server(inference_server):
-    inf_split = inference_server.split("    ")
-    assert len(inf_split) == 1 or len(inf_split) == 3
-    inference_server = inf_split[0]
-    if len(inf_split) == 3:
-        headers = {"authorization": "%s %s" % (inf_split[1], inf_split[2])}
-    else:
-        headers = None
-    return inference_server, headers
-
-
-class FakeTokenizer:
-    """
-    1) For keeping track of model_max_length
-    2) For when model doesn't directly expose tokenizer but need to count tokens
-    """
-
-    def __init__(self, model_max_length=2048, encoding_name="cl100k_base"):
-        # dont' push limit, since if using fake tokenizer, only estimate, and seen underestimates by order 250
-        self.model_max_length = model_max_length - 250
-        self.encoding_name = encoding_name
-        # The first time this runs, it will require an internet connection to download. Later runs won't need an internet connection.
-        import tiktoken
-        self.encoding = tiktoken.get_encoding(self.encoding_name)
-
-    def encode(self, x, *args, return_tensors="pt", **kwargs):
-        input_ids = self.encoding.encode(x, disallowed_special=())
-        if return_tensors == 'pt' and isinstance(input_ids, list):
-            import torch
-            input_ids = torch.tensor(input_ids)
-        return dict(input_ids=input_ids)
-
-    def decode(self, x, *args, **kwargs):
-        # input is input_ids[0] form
-        return self.encoding.decode(x)
-
-    def num_tokens_from_string(self, prompt: str) -> int:
-        """Returns the number of tokens in a text string."""
-        num_tokens = len(self.encoding.encode(prompt))
-        return num_tokens
-
-    def __call__(self, x, *args, **kwargs):
-        return self.encode(x, *args, **kwargs)
-
-
-def get_local_ip():
-    import socket
-    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-    try:
-        # doesn't even have to be reachable
-        s.connect(('10.255.255.255', 1))
-        IP = s.getsockname()[0]
-    except Exception:
-        IP = '127.0.0.1'
-    finally:
-        s.close()
-    return IP
-
-
-try:
-    assert pkg_resources.get_distribution('langchain') is not None
-    have_langchain = True
-except (pkg_resources.DistributionNotFound, AssertionError):
-    have_langchain = False
-
-import distutils.spawn
-
-have_tesseract = distutils.spawn.find_executable("tesseract")
-have_libreoffice = distutils.spawn.find_executable("libreoffice")
-
-import pkg_resources
-
-try:
-    assert pkg_resources.get_distribution('arxiv') is not None
-    assert pkg_resources.get_distribution('pymupdf') is not None
-    have_arxiv = True
-except (pkg_resources.DistributionNotFound, AssertionError):
-    have_arxiv = False
-
-try:
-    assert pkg_resources.get_distribution('pymupdf') is not None
-    have_pymupdf = True
-except (pkg_resources.DistributionNotFound, AssertionError):
-    have_pymupdf = False
-
-try:
-    assert pkg_resources.get_distribution('selenium') is not None
-    have_selenium = True
-except (pkg_resources.DistributionNotFound, AssertionError):
-    have_selenium = False
-
-try:
-    assert pkg_resources.get_distribution('playwright') is not None
-    have_playwright = True
-except (pkg_resources.DistributionNotFound, AssertionError):
-    have_playwright = False
-
-# disable, hangs too often
-have_playwright = False
-
-
-def set_openai(inference_server):
-    if inference_server.startswith('vllm'):
-        import openai_vllm
-        openai_vllm.api_key = "EMPTY"
-        inf_type = inference_server.split(':')[0]
-        ip_vllm = inference_server.split(':')[1]
-        port_vllm = inference_server.split(':')[2]
-        openai_vllm.api_base = f"http://{ip_vllm}:{port_vllm}/v1"
-        return openai_vllm, inf_type
-    else:
-        import openai
-        openai.api_key = os.getenv("OPENAI_API_KEY")
-        openai.api_base = os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1")
-        inf_type = inference_server
-        return openai, inf_type
-
-
-visible_langchain_modes_file = 'visible_langchain_modes.pkl'
-
-
-def save_collection_names(langchain_modes, visible_langchain_modes, langchain_mode_paths, LangChainMode, db1s):
-    """
-    extra controls if UserData type of MyData type
-    """
-
-    # use first default MyData hash as general user hash to maintain file
-    # if user moves MyData from langchain modes, db will still survive, so can still use hash
-    scratch_collection_names = list(db1s.keys())
-    user_hash = db1s.get(LangChainMode.MY_DATA.value, '')[1]
-
-    llms = ['LLM', 'Disabled']
-
-    scratch_langchain_modes = [x for x in langchain_modes if x in scratch_collection_names]
-    scratch_visible_langchain_modes = [x for x in visible_langchain_modes if x in scratch_collection_names]
-    scratch_langchain_mode_paths = {k: v for k, v in langchain_mode_paths.items() if
-                                    k in scratch_collection_names and k not in llms}
-
-    user_langchain_modes = [x for x in langchain_modes if x not in scratch_collection_names]
-    user_visible_langchain_modes = [x for x in visible_langchain_modes if x not in scratch_collection_names]
-    user_langchain_mode_paths = {k: v for k, v in langchain_mode_paths.items() if
-                                 k not in scratch_collection_names and k not in llms}
-
-    base_path = 'locks'
-    makedirs(base_path)
-
-    # user
-    extra = ''
-    file = "%s%s" % (visible_langchain_modes_file, extra)
-    with filelock.FileLock(os.path.join(base_path, "%s.lock" % file)):
-        with open(file, 'wb') as f:
-            pickle.dump((user_langchain_modes, user_visible_langchain_modes, user_langchain_mode_paths), f)
-
-    # scratch
-    extra = user_hash
-    file = "%s%s" % (visible_langchain_modes_file, extra)
-    with filelock.FileLock(os.path.join(base_path, "%s.lock" % file)):
-        with open(file, 'wb') as f:
-            pickle.dump((scratch_langchain_modes, scratch_visible_langchain_modes, scratch_langchain_mode_paths), f)
-
-
-def load_collection_enum(extra):
-    """
-    extra controls if UserData type of MyData type
-    """
-    file = "%s%s" % (visible_langchain_modes_file, extra)
-    langchain_modes_from_file = []
-    visible_langchain_modes_from_file = []
-    langchain_mode_paths_from_file = {}
-    if os.path.isfile(visible_langchain_modes_file):
-        try:
-            with filelock.FileLock("%s.lock" % file):
-                with open(file, 'rb') as f:
-                    langchain_modes_from_file, visible_langchain_modes_from_file, langchain_mode_paths_from_file = pickle.load(
-                        f)
-        except BaseException as e:
-            print("Cannot load %s, ignoring error: %s" % (file, str(e)), flush=True)
-    for k, v in langchain_mode_paths_from_file.items():
-        if v is not None and not os.path.isdir(v) and isinstance(v, str):
-            # assume was deleted, but need to make again to avoid extra code elsewhere
-            makedirs(v)
-    return langchain_modes_from_file, visible_langchain_modes_from_file, langchain_mode_paths_from_file
-
-
-def remove_collection_enum():
-    remove(visible_langchain_modes_file)

utils_langchain.py

@@ -1,64 +0,0 @@
-from typing import Any, Dict, List, Union, Optional
-import time
-import queue
-
-from langchain.callbacks.base import BaseCallbackHandler
-from langchain.schema import LLMResult
-
-
-class StreamingGradioCallbackHandler(BaseCallbackHandler):
-    """
-    Similar to H2OTextIteratorStreamer that is for HF backend, but here LangChain backend
-    """
-    def __init__(self, timeout: Optional[float] = None, block=True):
-        super().__init__()
-        self.text_queue = queue.SimpleQueue()
-        self.stop_signal = None
-        self.do_stop = False
-        self.timeout = timeout
-        self.block = block
-
-    def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        """Run when LLM starts running. Clean the queue."""
-        while not self.text_queue.empty():
-            try:
-                self.text_queue.get(block=False)
-            except queue.Empty:
-                continue
-
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Run on new LLM token. Only available when streaming is enabled."""
-        self.text_queue.put(token)
-
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        """Run when LLM ends running."""
-        self.text_queue.put(self.stop_signal)
-
-    def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when LLM errors."""
-        self.text_queue.put(self.stop_signal)
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        while True:
-            try:
-                value = self.stop_signal  # value looks unused in pycharm, not true
-                if self.do_stop:
-                    print("hit stop", flush=True)
-                    # could raise or break, maybe best to raise and make parent see if any exception in thread
-                    raise StopIteration()
-                    # break
-                value = self.text_queue.get(block=self.block, timeout=self.timeout)
-                break
-            except queue.Empty:
-                time.sleep(0.01)
-        if value == self.stop_signal:
-            raise StopIteration()
-        else:
-            return value