Update with h2oGPT hash 13a8343d2a96885985bda8c4480bbb23cf55bb9b

LICENSE

@@ -1,201 +0,0 @@
-                                Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
-
-   APPENDIX: How to apply the Apache License to your work.
-
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "[]"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.

LICENSE

@@ -0,0 +1 @@
+../../LICENSE

client_test.py

@@ -1,179 +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: --infer_devices=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 client_test.py
-
-
-
-For HF spaces:
-
-HOST="https://h2oai-h2ogpt-chatbot.hf.space" python 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 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
-
-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, langchain_mode='Disabled'):
-    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,
-                         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,
-                         top_k_docs=4,
-                         document_choice=['All'],
-                         )
-    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():
-    return run_client_nochat(prompt='Who are you?', prompt_type='human_bot', 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(ast.literal_eval(res)['response']),
-                    sources=ast.literal_eval(res)['sources'])
-    print(res_dict)
-    return res_dict
-
-
-@pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_chat():
-    return run_client_chat(prompt='Who are you?', prompt_type='human_bot', stream_output=False, max_new_tokens=50,
-                           langchain_mode='Disabled')
-
-
-def run_client_chat(prompt, prompt_type, stream_output, max_new_tokens, langchain_mode):
-    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)
-    return run_client(client, prompt, args, kwargs)
-
-
-def run_client(client, prompt, args, kwargs, do_md_to_text=True, verbose=False):
-    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
-
-
-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()
-
-
-if __name__ == '__main__':
-    test_client_basic()

client_test.py

@@ -0,0 +1 @@
+../../client_test.py

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)[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(llama_type=llama_type, model_name=base_model, reward_type=False)
-    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))

create_data.py

@@ -0,0 +1 @@
+../../create_data.py

enums.py

@@ -0,0 +1 @@
+../../enums.py

finetune.py

@@ -1,670 +0,0 @@
-import os
-import sys
-from functools import partial
-from typing import List, Union
-import fire
-import numpy as np
-
-from loaders import get_loaders, get_tokenizer
-from prompter import generate_prompt, prompt_types
-from utils import get_githash, copy_code
-import torch
-
-
-def log(*args, **kwargs):
-    if int(os.environ.get("LOCAL_RANK", 0)) == 0:
-        if 'flush' not in kwargs:
-            kwargs['flush'] = True
-        print(*args, **kwargs)
-
-
-# supported by huggingface evaluate
-supported_metrics = ['bleu', 'rouge', 'sacrebleu', 'meteor']
-
-
-def train(
-        save_code: bool = False,
-        run_id: int = None,
-
-        base_model: str = 'h2oai/h2ogpt-oig-oasst1-512-6_9b',
-        # base_model: str = 'h2oai/h2ogpt-oasst1-512-12b',
-        # base_model: str = 'h2oai/h2ogpt-oasst1-512-20b',
-        # base_model: str = 'EleutherAI/gpt-neox-20b',
-        # base_model: str = 'EleutherAI/pythia-12b-deduped',
-        # base_model: str = 'togethercomputer/GPT-NeoXT-Chat-Base-20B',
-        # base_model: str = 'decapoda-research/llama-7b-hf',
-        # base_model: str = 'decapoda-research/llama-13b-hf',
-        # base_model: str = 'decapoda-research/llama-30b-hf',
-        # base_model: str = 'EleutherAI/gpt-j-6B',
-
-        # only needed if base_model is self-exported HF state without tokenizer
-        tokenizer_base_model: str = None,
-        # tokenizer_base_model: str = 'EleutherAI/gpt-neox-20b',
-
-        data_path: str = "h2oai/openassistant_oasst1_h2ogpt",
-        data_col_dict: dict = None,
-        # data_path: str = "./dai_docs.train.json",
-        prompt_type: Union[str, int] = "plain",  # "plain", "instruct", "quality", "human_bot", "dai_faq"
-
-        valid_path: str = None,
-        # valid_path: str = "./dai_docs.valid.json",
-
-        # data_mix_in_path: str = "laion/OIG",  # way too big, medium quality
-        data_mix_in_path: str = "0-hero/OIG-small-chip2",  # high quality, 50 MB, good enough for now
-        data_mix_in_factor: float = 0.0,  # >1: more mix-in data, <1: more of data_path data
-        data_mix_in_col_dict: dict = {'user': 'instruction', 'chip2': 'output'},
-        data_mix_in_prompt_type: str = "instruct",  # just instruction->output, same as instruct
-
-        output_dir: str = None,
-
-        # LoRA checkpoint continuation
-        lora_weights: str = "",
-
-        # batching training hyperparams
-        batch_size: int = 128,
-        micro_batch_size: int = 4,
-        gradient_checkpointing=False,  # unnecessary with gradient accumulation enabled
-        fp16=True,
-        train_8bit=False,
-        train_4bit=False,
-
-        # general training hyperparams
-        num_epochs: float = 1,
-        learning_rate: float = 3e-4,
-
-        # validation settings
-        val_set_size: int = None,
-        val_metrics: List[str] = [],
-        eval_steps: int = None,  # to control eval steps via steps
-        eval_epochs: float = None,  # to control eval steps via epochs
-
-        # lora hyperparams
-        lora_r: int = 8,
-        lora_alpha: int = 16,
-        lora_dropout: float = 0.05,
-        lora_target_modules: List[str] = None,
-        llama_type: bool = None,
-        llama_flash_attn: bool = False,
-
-        # llm hyperparams
-        train_on_inputs: bool = True,  # if False, masks out inputs in loss
-        group_by_length: bool = False,  # if True, faster, but produces an odd training loss curve
-        resume_from_checkpoint: str = None,  # either training checkpoint or final adapter
-        cutoff_len: int = 512,  # larger values use more memory
-        drop_truncations: bool = False,  # if True, drop any truncated long sequences
-
-        # torch training params
-        ddp: bool = True,  # set to False if OOM with True, for multi-GPU model parallelism
-        local_files_only: bool = False,  # else will download new versions, normally unwanted
-        resume_download: bool = True,
-        use_auth_token: Union[str, bool] = False,  # True requires CLI did huggingface-cli login before running
-        warmup_steps: int = 100,
-        logging_steps: int = 1,
-        save_steps: int = None,  # must be round multiple of eval_steps
-        save_total_limit: int = 3,
-        add_eos_token: bool = False,
-):
-
-    if llama_flash_attn:
-        # Need to call this before importing transformers.
-        from llama_flash_attn_monkey_patch import replace_llama_attn_with_flash_attn
-        replace_llama_attn_with_flash_attn()
-
-    # allow set token directly
-    use_auth_token = os.environ.get("HUGGINGFACE_API_TOKEN", use_auth_token)
-
-    prompt_type = str(prompt_type)  # migration from integers
-    assert prompt_type in prompt_types
-
-    world_size = int(os.getenv("WORLD_SIZE", 1))
-    local_rank = int(os.getenv("LOCAL_RANK", 0))
-    rank = int(os.getenv("RANK", 0))
-    print(f"local_rank: {local_rank}")
-    print(f"global rank: {rank}")
-
-    gpus = max(world_size, torch.cuda.device_count())
-    run_id = run_id or 0
-    if not data_path:
-        raise ValueError("No data_path provided")
-    if not output_dir:
-        output_dir = f"{base_model.split('/')[-1]}.{data_path.replace('/', '')}.{num_epochs}_epochs.{get_githash() or 'nogit'}.{run_id}"
-        if os.path.exists(output_dir) and not resume_from_checkpoint:
-            raise FileExistsError(f"output_dir {output_dir} based on run_id {run_id} already exists. Please pick a different run_id.")
-    else:
-        if os.path.exists(output_dir) and not resume_from_checkpoint:
-            raise FileExistsError(f"output_dir {output_dir} already exists. Please pick a different output_dir, or specify a run_id instead.")
-    device_map = "auto"
-
-    if save_code:
-        copy_code(run_id)
-    if tokenizer_base_model is None:
-        tokenizer_base_model = base_model
-    if llama_type is None:
-        llama_type = "llama" in base_model.lower()
-    if llama_type and llama_flash_attn:
-        import pkg_resources
-        try:
-            pkg_resources.get_distribution('flash_attn')
-            can_do_flash_attn = True
-        except (pkg_resources.DistributionNotFound, pkg_resources.ContextualVersionConflict):
-            can_do_flash_attn = False
-
-        if not can_do_flash_attn:
-            raise RuntimeError("""Flash attention not installed.
-            NOTE: for current pytorch 2.0, flash attention requires installing cuda 11.7 via https://developer.nvidia.com/cuda-11-7-0-download-archive?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=20.04&target_type=runfile_local and then when running, to avoid installing driver, docs, samples, just install toolkit.  Then when pip installing flash attention do:
-
-            CUDA_HOME=/usr/local/cuda-11.7 pip install flash-attn""")
-    assert (
-        base_model
-    ), "Please specify a --base_model, e.g. --base_model='decapoda-research/llama-7b-hf'"
-    gradient_accumulation_steps = batch_size // micro_batch_size
-    assert gradient_accumulation_steps >= world_size, "must increase batch_size for multi-GPU"
-
-    device_map = "auto"
-
-    locals_dict = locals()
-    locals_print = '\n'.join(['%s: %s' % (k, v) for k, v in locals_dict.items()])
-    log(f"Training model with params:\n{locals_print}")
-    log("Command: %s\nHash: %s" % (str(' '.join(sys.argv)), get_githash()))
-
-    max_memory = None
-    if gpus > 1:
-        if ddp:
-            log("Distributed: data parallel")
-            device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
-            gradient_accumulation_steps = gradient_accumulation_steps // world_size
-        else:
-            free_in_GB = int(min(torch.cuda.mem_get_info()) / 1024 ** 3)
-            max_memory = f"{free_in_GB - 2}GB"
-            max_memory = {i: max_memory for i in range(gpus)}
-            log("world_size: %d" % world_size)
-            log("num_gpus: %d" % gpus)
-            log("max mem: %s" % max_memory)
-
-    model_loader, tokenizer_loader = get_loaders(llama_type=llama_type, model_name=base_model, reward_type=False)
-
-    model = model_loader.from_pretrained(
-        base_model,
-        load_in_8bit=train_8bit,
-        load_in_4bit=train_4bit,
-        device_map=device_map,
-        torch_dtype=torch.float16,
-        max_memory=max_memory,
-        local_files_only=local_files_only,
-        trust_remote_code=True,
-        resume_download=resume_download,
-        use_auth_token=use_auth_token,
-    )
-    if gpus > 1:
-        if not ddp:
-            log("model parallel")
-            model.is_parallelizable = True
-            model.model_parallel = True
-
-    tokenizer = get_tokenizer(tokenizer_loader, tokenizer_base_model, local_files_only, resume_download, use_auth_token)
-
-    if train_8bit or train_4bit:
-        from peft import (
-            prepare_model_for_kbit_training,
-        )
-
-        model = prepare_model_for_kbit_training(model)
-
-    from peft import LoraConfig, get_peft_model, set_peft_model_state_dict
-    try:
-        from peft import utils
-        lora_mappings = utils.TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING.copy()
-    except AttributeError:
-        from peft import mapping
-        lora_mappings = mapping.TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING.copy()
-    lora_mappings['distilgpt2'] = ["c_attn"]
-
-    if lora_weights:
-
-        from peft import PeftModel
-        model = PeftModel.from_pretrained(
-            model,
-            lora_weights,
-            torch_dtype=torch.float16,
-            device_map=device_map,
-            local_files_only=local_files_only,
-            resume_download=resume_download,
-            use_auth_token=use_auth_token,
-        )
-    elif lora_r > 0:
-        if lora_target_modules is None:
-            base_model_lower = base_model.lower()
-            if base_model_lower in lora_mappings:
-                lora_target_modules_cand = [lora_mappings[base_model_lower]]
-            else:
-                lora_target_modules_cand = [["query_key_value"], ["q_proj", "v_proj"]]
-        else:
-            lora_target_modules_cand = [lora_target_modules]
-
-        for lora_target_modules in lora_target_modules_cand:
-            try:
-                config = LoraConfig(
-                    r=lora_r,
-                    lora_alpha=lora_alpha,
-                    target_modules=lora_target_modules,
-                    lora_dropout=lora_dropout,
-                    bias="none",
-                    task_type="CAUSAL_LM",
-                )
-                model = get_peft_model(model, config)
-                break
-            except ValueError as e:
-                if "Target modules" in str(e) and "not found" in str(e):
-                    continue
-                else:
-                    raise
-        from peft import PeftModel
-        assert isinstance(model, PeftModel), "LoRA failed. Please provide --lora_target_modules explicitly."
-    if resume_from_checkpoint:
-        # Check the available weights and load them
-        checkpoint_name = os.path.join(
-            resume_from_checkpoint, "pytorch_model.bin"
-        )  # Full checkpoint
-        if not os.path.exists(checkpoint_name):
-            checkpoint_name = os.path.join(
-                resume_from_checkpoint, "adapter_model.bin"
-            )  # only LoRA model - LoRA config above has to fit
-            resume_from_checkpoint = False  # So the trainer won't try loading its state
-        # The two files above have a different name depending on how they were saved, but are actually the same.
-        if os.path.exists(checkpoint_name):
-            log(f"Restarting from {checkpoint_name}")
-            adapters_weights = torch.load(checkpoint_name)
-            set_peft_model_state_dict(model, adapters_weights)
-        else:
-            log(f"Checkpoint {checkpoint_name} not found")
-
-    print(model)
-    try:
-        # only for PeftModel
-        model.print_trainable_parameters()  # Be more transparent about the % of trainable params.
-    except:
-        pass
-
-    metrics = {}
-    for name in supported_metrics:
-        if name in val_metrics:
-            import evaluate  # Causes hang for 'python generate.py' on dual 4090 if imported early, 100% reproducible
-            metrics[name] = evaluate.load(name)
-    log("Using Validation Metrics: %s" % str(list(metrics.keys())))
-    log("Supported Metrics: %s" % supported_metrics)
-
-    if val_set_size is None:
-        if len(metrics) == 0:
-            val_set_size = 1000
-        else:
-            val_set_size = 100
-        log("Auto set val_set_size %s" % val_set_size)
-    elif val_set_size < 1.0 and val_set_size != 0:
-        raise RuntimeError("Fractional validation size not supported.")
-
-    from datasets import load_dataset, concatenate_datasets
-    if valid_path:
-        data = load_dataset("json", data_files={"train": data_path, "valid": valid_path})
-    else:
-        if "json" in data_path:
-            data = load_dataset("json", data_files={"train": data_path})
-        else:
-            data = load_dataset(data_path)
-            data = data.rename_columns(data_col_dict or {})
-
-    valid_data = None
-    train_data_mix_in = None
-    valid_data_mix_in = None
-
-    if data_mix_in_path and data_mix_in_factor > 0:
-        # get mix-in training/validation data - to keep model "sane"
-        num_rows = data["train"].num_rows
-        log("Loading mix-in dataset: %s" % data_mix_in_path)
-        if "json" in data_mix_in_path:
-            data_mix_in = load_dataset("json", data_files={"train": data_mix_in_path})["train"]
-        else:
-            data_mix_in = load_dataset(data_mix_in_path)["train"]  # can be large
-        data_mix_in = data_mix_in.rename_columns(data_mix_in_col_dict or {})
-        mix_in_rows = int(num_rows * data_mix_in_factor)
-
-        if mix_in_rows > data_mix_in.num_rows:
-            # duplicate rows if mix-in is smaller than required
-            log("Duplicating mixin to compensate for its size for training size and mixin fraction")
-            data_mix_in = concatenate_datasets([data_mix_in] * int(np.ceil(mix_in_rows / data_mix_in.num_rows)))
-
-        # only get as much as we need to balance
-        valid_size = min(data_mix_in.num_rows // 2, val_set_size or 0)
-        train_size = max(1, min(data_mix_in.num_rows - valid_size, mix_in_rows))
-        mixin_small = data_mix_in.train_test_split(
-            test_size=train_size + valid_size,
-            shuffle=True, seed=np.random.randint(10000),
-        )["test"]
-        if valid_size:
-            mixin_train_test = mixin_small.train_test_split(
-                test_size=valid_size, shuffle=False,
-            )
-            train_data_mix_in = mixin_train_test["train"]
-            valid_data_mix_in = mixin_train_test["test"]
-        else:
-            train_data_mix_in = mixin_small
-
-        if "prompt_type" not in train_data_mix_in.column_names:
-            train_data_mix_in = train_data_mix_in.add_column(
-                "prompt_type",
-                [data_mix_in_prompt_type] * train_data_mix_in.num_rows,
-            )
-            log("Added prompt type %s to mix-in training data" % data_mix_in_prompt_type)
-        if valid_data_mix_in and "prompt_type" not in valid_data_mix_in.column_names:
-            valid_data_mix_in = valid_data_mix_in.add_column(
-                "prompt_type",
-                [data_mix_in_prompt_type] * valid_data_mix_in.num_rows,
-            )
-            log("Added prompt type %s to mix-in validation data" % data_mix_in_prompt_type)
-        log("Created mix-in data:\nTrain %s\nValid %s" % (train_data_mix_in, valid_data_mix_in))
-
-    # get our own training/validation data - for fine-tuning
-    if val_set_size > 0 and not valid_path and not data_mix_in_path:
-        # create valid split from train
-        train_val = data["train"].train_test_split(
-            test_size=val_set_size, shuffle=True, seed=42
-        )
-        train_data = train_val["train"]
-        valid_data = train_val["test"]
-    else:
-        train_data = data["train"]
-        if valid_path:
-            # use given valid split, has priority over data_mix_in_path
-            valid_data = data["valid"]
-    if "prompt_type" not in train_data.column_names:
-        train_data = train_data.add_column(
-            "prompt_type",
-            [prompt_type] * train_data.num_rows,
-        )
-        log("Added prompt type %s to training data" % prompt_type)
-    if valid_data and "prompt_type" not in valid_data.column_names:
-        valid_data = valid_data.add_column(
-            "prompt_type",
-            [prompt_type] * valid_data.num_rows,
-        )
-        log("Added prompt type %s to validation data" % prompt_type)
-
-    assert train_data is not None
-
-    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)
-
-    # shuffle and tokenize data
-    if train_data_mix_in:
-        train_data = concatenate_datasets([train_data, train_data_mix_in])
-    log("Tokenizing %s training rows" % train_data.num_rows)
-    train_data = train_data.shuffle().map(generate_and_tokenize_prompt_fun, num_proc=os.cpu_count() // torch.cuda.device_count())
-    if drop_truncations:
-        log("avoid keeping truncated cases to avoid contaminating model with truncation cases.  Original size: %s" % train_data.num_rows)
-        prune_long_sequences_func = partial(prune_long_sequences, cutoff_len=cutoff_len)
-        train_data = train_data.filter(prune_long_sequences_func, num_proc=os.cpu_count() // torch.cuda.device_count())
-        log("avoid keeping truncated cases to avoid contaminating model with truncation cases.  New size: %s" % train_data.num_rows)
-    train_set_size = len(train_data)
-
-    if valid_data and valid_data_mix_in:
-        valid_data = concatenate_datasets([valid_data, valid_data_mix_in])
-    elif valid_data_mix_in:
-        valid_data = valid_data_mix_in
-
-    if valid_data:
-        log("Tokenizing %s validation rows" % valid_data.num_rows)
-        valid_data = valid_data.shuffle().map(generate_and_tokenize_prompt_fun, num_proc=os.cpu_count() // torch.cuda.device_count())
-        val_set_size = len(valid_data)
-    else:
-        val_set_size = 0
-    log("Final fine-tuning data:\nTrain %s\nValid %s" % (train_data, valid_data))
-    sample_row_dict = train_data[:1]
-    del sample_row_dict['input_ids']
-    del sample_row_dict['attention_mask']
-    del sample_row_dict['labels']
-    log("Sample input: %s" % sample_row_dict)
-
-    try:
-        import neptune
-        from transformers.integrations import NeptuneCallback
-
-        neptune_run = neptune.init_run(
-            source_files=[],
-        )
-        log("Connected to Neptune.")
-    except ImportError:
-        neptune_run = None
-        log("Please pip install neptune for tracking.")
-    except neptune.exceptions.NeptuneMissingApiTokenException:
-        neptune_run = None
-        os.environ["NEPTUNE_MODE"] = 'debug'
-        log("No neptune configured, set NEPTUNE_API_TOKEN env var.")
-
-    if neptune_run:
-        neptune_callback = NeptuneCallback(run=neptune_run)
-        callbacks = [neptune_callback]
-    else:
-        from transformers.integrations import TensorBoardCallback, is_tensorboard_available
-        if is_tensorboard_available:
-            # tensorboard --logdir=runs/
-            from torch.utils.tensorboard import SummaryWriter
-            tb_writer = SummaryWriter()
-            callbacks = [TensorBoardCallback(tb_writer=tb_writer)]
-        else:
-            callbacks = []
-
-    expected_steps = (train_set_size * num_epochs) // batch_size
-    if eval_steps is None and eval_epochs is None:
-        # 20 evaluations for a run
-        eval_steps = max(1, int(expected_steps / 20))
-        log("Auto set eval_steps to %s out of %s total training steps" % (eval_steps, expected_steps))
-    elif eval_steps is None and eval_epochs is not None:
-        eval_steps = max(1, int(expected_steps * eval_epochs / num_epochs))
-        log("Auto converted eval_epochs=%s to eval_steps %s"
-            " out of %s total training steps" % (eval_epochs, eval_steps, expected_steps))
-    if save_steps is None:
-        save_steps = eval_steps
-        log("Auto step save_steps to %s" % save_steps)
-    elif save_steps > eval_steps:
-        # save steps must be round multiple of eval_steps
-        save_steps0 = save_steps
-        save_steps = max(1, (save_steps//eval_steps)) * eval_steps
-        if save_steps0 != save_steps:
-            log("Auto converted save_steps from %s to %s" % (save_steps0, save_steps))
-
-    def compute_metrics(eval_preds):
-        # e.g. see: https://huggingface.co/docs/transformers/v4.25.1/en/tasks/translation#evaluate
-        inputs = eval_preds.inputs
-        label_ids = eval_preds.label_ids
-        predictions = eval_preds.predictions
-
-        #inputs = np.where(inputs != -100, inputs, tokenizer.pad_token_id)
-        #decoded_inputs = tokenizer.batch_decode(inputs, skip_special_tokens=True)
-        #decoded_inputs = [pred.strip() for pred in decoded_inputs]
-
-        label_ids = np.where(label_ids != -100, label_ids, tokenizer.pad_token_id)
-        # tokenizer behavior like generate time
-        decoded_labels = tokenizer.batch_decode(label_ids, skip_special_tokens=True,
-                                                           clean_up_tokenization_spaces=True)
-        decoded_labels = [pred.strip() for pred in decoded_labels]
-
-        predictions = np.argmax(predictions, -1)
-        predictions = np.where(predictions != -100, predictions, tokenizer.pad_token_id)
-        # tokenizer behavior like generate time
-        decoded_predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True,
-                                                                  clean_up_tokenization_spaces=True)
-        decoded_predictions = [pred.strip() for pred in decoded_predictions]
-
-        result = {}
-        for metric in metrics.values():
-            result1 = metric.compute(predictions=decoded_predictions, references=decoded_labels)
-            # get rid of lists, for precision etc., for now
-            numeric_results = {k: v for k, v in result1.items() if isinstance(v, (int, float))}
-            result.update(numeric_results)
-        return result
-
-    # the callback that computes metrics of interest
-    if val_metrics:
-        trainer_kwargs = dict(compute_metrics=compute_metrics)
-    else:
-        trainer_kwargs = dict()
-
-    import transformers
-    trainer = transformers.Trainer(
-        model=model,
-        tokenizer=tokenizer,
-        train_dataset=train_data,
-        eval_dataset=valid_data,
-        # FIXME: might need Seq2SeqTrainingArguments for some models
-        args=transformers.TrainingArguments(
-            per_device_train_batch_size=micro_batch_size,
-            per_device_eval_batch_size=1,
-            eval_accumulation_steps=10,
-            # predict_with_generate=True,  # SEQ2SEQ only
-            include_inputs_for_metrics=True,
-            gradient_accumulation_steps=gradient_accumulation_steps,
-            warmup_steps=warmup_steps,
-            num_train_epochs=num_epochs,
-            learning_rate=learning_rate,
-            gradient_checkpointing=gradient_checkpointing,
-            fp16=fp16,
-            # cosnider 8-bit adam: https://huggingface.co/docs/transformers/v4.18.0/en/performance#8bit-adam
-            optim="adamw_torch",  # consider "adafactor" to save memory
-            logging_steps=logging_steps,
-            logging_strategy="steps",
-            evaluation_strategy="steps" if val_set_size > 0 else "no",
-            save_strategy="steps",
-            eval_steps=eval_steps if val_set_size > 0 else None,
-            save_steps=save_steps,
-            output_dir=output_dir,
-            save_total_limit=save_total_limit,
-            load_best_model_at_end=True if val_set_size > 0 else False,
-            ddp_find_unused_parameters=False if ddp else None,
-            group_by_length=group_by_length,
-            #fsdp="shard_grad_op auto_wrap" if gpus > 1 and not ddp else None,
-            #fsdp_min_num_params=20000 if gpus > 1 and not ddp else None,
-            report_to='tensorboard' if not neptune_run else 'neptune',
-        ),
-        data_collator=transformers.DataCollatorForSeq2Seq(
-            tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
-        ),
-        callbacks=callbacks,
-        **trainer_kwargs,
-    )
-    model.config.use_cache = False
-
-    old_state_dict = model.state_dict
-    from peft import get_peft_model_state_dict
-
-    model.state_dict = (
-        lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())
-    ).__get__(model, type(model))
-
-    if torch.__version__ >= "2" and sys.platform != "win32":
-        model = torch.compile(model)
-        # WIP (not generally replacing layers until pytorch 2.1)
-        if not llama_flash_attn:
-            torch.backends.cuda.enable_flash_sdp(True)
-
-    if gpus > 1 and not ddp:
-        assert trainer.is_model_parallel
-    else:
-        assert not trainer.is_model_parallel
-    trainer.train(resume_from_checkpoint=resume_from_checkpoint)
-
-    model.save_pretrained(output_dir)
-
-    log("\n If there's a warning about missing keys above, please disregard :)")
-
-
-def tokenize(prompt, tokenizer, cutoff_len, add_eos_token=False):
-    # there's probably a way to do this with the tokenizer settings
-    # but again, gotta move fast
-    result = tokenizer(
-        prompt,
-        truncation=True,
-        max_length=cutoff_len,
-        padding=False,
-        return_tensors=None,
-    )
-    if (
-            result["input_ids"][-1] != tokenizer.eos_token_id
-            and len(result["input_ids"]) < cutoff_len
-            and add_eos_token
-    ):
-        result["input_ids"].append(tokenizer.eos_token_id)
-        result["attention_mask"].append(1)
-
-    result["labels"] = result["input_ids"].copy()
-
-    return result
-
-
-def prune_long_sequences(data_point, cutoff_len=None):
-    """
-    Prune if too long for tokenizer, so truncation doesn't lead training to learn from truncated language
-    :param data_point:
-    :param cutoff_len:
-    :return:
-    """
-    assert cutoff_len is not None
-    return len(data_point['input_ids']) < cutoff_len
-
-
-def generate_and_tokenize_prompt(data_point, prompt_type=None, train_on_inputs=False, add_eos_token=False,
-                                 cutoff_len=None, tokenizer=None):
-    assert prompt_type is not None
-    assert cutoff_len is not None
-    assert tokenizer is not None
-    full_prompt, _, _, _ = generate_prompt(data_point, prompt_type, False, False)
-    tokenized_full_prompt = tokenize(full_prompt, tokenizer, cutoff_len, add_eos_token=add_eos_token)
-    if not train_on_inputs:
-        user_prompt, _, _, _ = generate_prompt({**data_point, "output": ""}, prompt_type, False, False)
-        tokenized_user_prompt = tokenize(user_prompt, tokenizer, cutoff_len, add_eos_token=add_eos_token)
-        user_prompt_len = len(tokenized_user_prompt["input_ids"])
-        if add_eos_token:
-            user_prompt_len -= 1
-
-        # ignore_index=-100 ensures torch/tf don't include padding token id in CrossEntropyLoss
-        tokenized_full_prompt["labels"] = [
-                                              -100
-                                          ] * user_prompt_len + tokenized_full_prompt["labels"][
-                                                                user_prompt_len:
-                                                                ]  # could be sped up, probably
-    return tokenized_full_prompt
-
-
-def test_debug():
-    fire.Fire(train)
-
-
-if __name__ == "__main__":
-    CONFIG = "NCCL_P2P_LEVEL=LOC WORLD_SIZE=5 torchrun --nnodes=5 --master_addr=10.10.10.2 --master_port=1111 --nproc_per_node=1"
-    CMD = "finetune.py --data_path=config.json --num_epochs=1 --base_model=decapoda-research/llama-13b-hf"
-    log(f"""
-    Example runs on 4 GPUs:
-    WORLD_SIZE=4 CUDA_VISIBLE_DEVICES="0,1,2,3" torchrun --nproc_per_node=4 finetune.py --base_model='decapoda-research/llama-7b-hf' --data_path=data/config.json --run_id=0 &> 0.log
-    WORLD_SIZE=4 CUDA_VISIBLE_DEVICES="0,1,2,3" torchrun --nproc_per_node=4 finetune.py --base_model='decapoda-research/llama-30b-hf' --data_path=data/config.json --batch_size=16 --micro_batch_size=1 --run_id=1 --save_code=True &> 1.log
-    WORLD_SIZE=4 CUDA_VISIBLE_DEVICES="0,1,2,3" torchrun --nproc_per_node=4 finetune.py --base_model='EleutherAI/gpt-j-6B' --data_path=data/config.json --run_id=2 &> 2.log
-    WORLD_SIZE=4 CUDA_VISIBLE_DEVICES="0,1,2,3" torchrun --nproc_per_node=4 finetune.py --base_model='EleutherAI/gpt-neox-20b' --data_path=data/config.json --run_id=8 --batch_size=16 --micro_batch_size=4 &> 8.log
-    WORLD_SIZE=4 CUDA_VISIBLE_DEVICES="0,1,2,3" torchrun --nproc_per_node=4 finetune.py --base_model='togethercomputer/GPT-NeoXT-Chat-Base-20B' --data_path=data/config.json --prompt_type='dai_faq' --run_id=13 --batch_size=16 --micro_batch_size=4 --num_epochs=100 --val_set_size=0 data_mix_in_path='' &> 13.log
-    WORLD_SIZE=4 CUDA_VISIBLE_DEVICES="0,1,2,3" torchrun --nproc_per_node=4 finetune.py --base_model='togethercomputer/GPT-NeoXT-Chat-Base-20B' --data_path=data/config.json --run_id=28 --batch_size=16 --micro_batch_size=4 --num_epochs=8 --val_set_size=0 --data_mix_in_factor=0.1 --data_mix_in_prompt_type='human_bot' --save_code=True --cutoff_len=512  &> 28.log
-
-    All metrics:
-    CUDA_VISIBLE_DEVICES= finetune.py --data_mix_in_factor=0 --eval_steps=100 --warmup_steps=2 --val_set_size=100 --val_metrics="['bleu', 'rouge', 'sacrebleu', 'meteor']"
-
-    # Fine-tune 20B on 24GB GPUs across 3 nodes with 3+2+2 GPUs
-    rippa>
-NCCL_P2P_LEVEL=LOC WORLD_SIZE=7 CUDA_VISIBLE_DEVICES="0,1,2" torchrun --node_rank 0 --nproc_per_node=3 --master_port=1234 --nnodes=3 --master_addr=10.10.10.2 finetune.py --data_path=merged_shuffled_OIG_87f6a1e788.json --micro_batch_size=1 --batch_size=7 --cutoff_len=512 --run_id=17 &>log.17.rank0
-    ova>
-NCCL_P2P_LEVEL=LOC WORLD_SIZE=7 CUDA_VISIBLE_DEVICES="0,1" torchrun --node_rank 1 --nproc_per_node=2 --master_port=1234 --nnodes=3 --master_addr=10.10.10.2 finetune.py --data_path=merged_shuffled_OIG_87f6a1e788.json --micro_batch_size=1 --batch_size=7 --cutoff_len=512 --run_id=17 &>log.17.rank1
-    timemachine>
-NCCL_P2P_LEVEL=LOC WORLD_SIZE=7 CUDA_VISIBLE_DEVICES="0,1" torchrun --node_rank 2 --nproc_per_node=2 --master_port=1234 --nnodes=3 --master_addr=10.10.10.2 finetune.py --data_path=merged_shuffled_OIG_87f6a1e788.json --micro_batch_size=1 --batch_size=7 --cutoff_len=512 --run_id=17 &>log.17.rank2
-
-    """, flush=True)
-
-    if os.environ.get("LOCAL_RANK") is None:
-        # then not using torchrun, so can't do distributed, ensure CVD set
-        assert os.environ.get("CUDA_VISIBLE_DEVICES") is not None, "Run python script using: torchrun finetune.py OR set CUDA_VISIBLE_DEVICES to single GPU"
-
-    fire.Fire(train)

finetune.py

@@ -0,0 +1 @@
+../../finetune.py

generate.py

@@ -1,1548 +0,0 @@
-import ast
-import functools
-import glob
-import inspect
-import queue
-import shutil
-import sys
-import os
-import time
-import traceback
-import typing
-import warnings
-from datetime import datetime
-import filelock
-import psutil
-
-os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
-os.environ['BITSANDBYTES_NOWELCOME'] = '1'
-warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
-
-from loaders import get_loaders
-from utils import set_seed, clear_torch_cache, save_generate_output, NullContext, wrapped_partial, EThread, get_githash, \
-    import_matplotlib, get_device, makedirs, get_kwargs
-
-import_matplotlib()
-
-SEED = 1236
-set_seed(SEED)
-
-from typing import Union
-
-import fire
-import torch
-from peft import PeftModel
-from transformers import GenerationConfig, AutoModel, TextIteratorStreamer
-from accelerate import init_empty_weights, infer_auto_device_map
-
-from prompter import Prompter, inv_prompt_type_to_model_lower, non_hf_types
-from stopping import get_stopping
-
-eval_extra_columns = ['prompt', 'response', 'score']
-
-langchain_modes = ['Disabled', 'ChatLLM', 'LLM', 'All', 'wiki', 'wiki_full', 'UserData', 'MyData', 'github h2oGPT',
-                   'DriverlessAI docs']
-
-scratch_base_dir = '/tmp/'
-
-
-def main(
-        load_8bit: bool = False,
-        load_4bit: bool = False,
-        load_half: bool = True,
-        infer_devices: bool = True,
-        base_model: str = '',
-        tokenizer_base_model: str = '',
-        lora_weights: str = "",
-        gpu_id: int = 0,
-        compile_model: bool = True,
-
-        prompt_type: Union[int, str] = None,
-        # input to generation
-        temperature: float = None,
-        top_p: float = None,
-        top_k: int = None,
-        num_beams: int = None,
-        repetition_penalty: float = None,
-        num_return_sequences: int = None,
-        do_sample: bool = None,
-        max_new_tokens: int = None,
-        min_new_tokens: int = None,
-        early_stopping: Union[bool, str] = None,
-        max_time: float = None,
-
-        memory_restriction_level: int = None,
-        debug: bool = False,
-        save_dir: str = None,
-        share: bool = True,
-        local_files_only: bool = False,
-        resume_download: bool = True,
-        use_auth_token: Union[str, bool] = False,
-        trust_remote_code: Union[str, bool] = True,
-        offload_folder: str = "offline_folder",
-
-        src_lang: str = "English",
-        tgt_lang: str = "Russian",
-
-        cli: bool = False,
-        cli_loop: bool = True,
-        gradio: bool = True,
-        gradio_avoid_processing_markdown: bool = False,
-        gradio_offline_level: int = 0,
-        chat: bool = True,
-        chat_context: bool = False,
-        stream_output: bool = True,
-        show_examples: bool = None,
-        verbose: bool = False,
-        h2ocolors: bool = False,
-        height: int = 600,
-        show_lora: bool = True,
-        login_mode_if_model0: bool = False,
-        block_gradio_exit: bool = True,
-        concurrency_count: int = 1,
-        api_open: bool = False,
-        allow_api: bool = True,
-        input_lines: int = 1,
-        auth: typing.List[typing.Tuple[str, str]] = None,
-
-        sanitize_user_prompt: bool = True,
-        sanitize_bot_response: bool = True,
-
-        extra_model_options: typing.List[str] = [],
-        extra_lora_options: typing.List[str] = [],
-
-        score_model: str = 'OpenAssistant/reward-model-deberta-v3-large-v2',
-        auto_score: bool = True,
-
-        eval_filename: str = None,
-        eval_prompts_only_num: int = 0,
-        eval_prompts_only_seed: int = 1234,
-        eval_as_output: bool = False,
-
-        langchain_mode: str = 'Disabled',
-        visible_langchain_modes: list = ['UserData', 'MyData'],
-        document_choice: list = ['All'],
-        user_path: str = None,
-        detect_user_path_changes_every_query: bool = False,
-        load_db_if_exists: bool = True,
-        keep_sources_in_context: bool = False,
-        db_type: str = 'chroma',
-        use_openai_embedding: bool = False,
-        use_openai_model: bool = False,
-        hf_embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
-        allow_upload_to_user_data: bool = True,
-        allow_upload_to_my_data: bool = True,
-        enable_url_upload: bool = True,
-        enable_text_upload: bool = True,
-        enable_sources_list: bool = True,
-        chunk: bool = True,
-        chunk_size: int = 512,
-        top_k_docs: int = 3,  # FIXME: Can go back to 4 once https://github.com/h2oai/h2ogpt/issues/192 fixed
-        n_jobs: int = -1,
-        enable_captions: bool = True,
-        captions_model: str = "Salesforce/blip-image-captioning-base",
-        pre_load_caption_model: bool = False,
-        caption_gpu: bool = True,
-        enable_ocr: bool = False,
-):
-    """
-
-    :param load_8bit: load model in 8-bit using bitsandbytes
-    :param load_4bit: load model in 4-bit using bitsandbytes
-    :param load_half: load model in float16
-    :param infer_devices: whether to control devices with gpu_id.  If False, then spread across GPUs
-    :param base_model: model HF-type name.  If use --base_model to preload model, cannot unload in gradio in models tab
-    :param tokenizer_base_model: tokenizer HF-type name.  Usually not required, inferred from base_model.
-    :param lora_weights: LORA weights path/HF link
-    :param gpu_id: if infer_devices, then use gpu_id for cuda device ID, or auto mode if gpu_id != -1
-    :param compile_model Whether to compile the model
-    :param prompt_type: type of prompt, usually matched to fine-tuned model or plain for foundational model
-    :param temperature: generation temperature
-    :param top_p: generation top_p
-    :param top_k: generation top_k
-    :param num_beams: generation number of beams
-    :param repetition_penalty: generation repetition penalty
-    :param num_return_sequences: generation number of sequences (1 forced for chat)
-    :param do_sample: generation sample
-    :param max_new_tokens: generation max new tokens
-    :param min_new_tokens: generation min tokens
-    :param early_stopping: generation early stopping
-    :param max_time: maximum time to allow for generation
-    :param memory_restriction_level: 0 = no restriction to tokens or model, 1 = some restrictions on token 2 = HF like restriction 3 = very low memory case
-    :param debug: enable debug mode
-    :param save_dir: directory chat data is saved to
-    :param share: whether to share the gradio app with sharable URL
-    :param local_files_only: whether to only use local files instead of doing to HF for models
-    :param resume_download: whether to resume downloads from HF for models
-    :param use_auth_token: whether to use HF auth token (requires CLI did huggingface-cli login before)
-    :param trust_remote_code: whether to use trust any code needed for HF model
-    :param offload_folder: path for spilling model onto disk
-    :param src_lang: source languages to include if doing translation (None = all)
-    :param tgt_lang: target languages to include if doing translation (None = all)
-    :param cli: whether to use CLI (non-gradio) interface.
-    :param cli_loop: whether to loop for CLI (False usually only for testing)
-    :param gradio: whether to enable gradio, or to enable benchmark mode
-    :param gradio_avoid_processing_markdown:
-    :param gradio_offline_level: > 0, then change fonts so full offline
-           == 1 means backend won't need internet for fonts, but front-end UI might if font not cached
-           == 2 means backend and frontend don't need internet to download any fonts.
-           Note: Some things always disabled include HF telemetry, gradio telemetry, chromadb posthog that involve uploading.
-           This option further disables google fonts for downloading, which is less intrusive than uploading,
-           but still required in air-gapped case.  The fonts don't look as nice as google fonts, but ensure full offline behavior.
-    :param chat: whether to enable chat mode with chat history
-    :param chat_context: whether to use extra helpful context if human_bot
-    :param stream_output: whether to stream output from generate
-    :param show_examples: whether to show clickable examples in gradio
-    :param verbose: whether to show verbose prints
-    :param h2ocolors: whether to use H2O.ai theme
-    :param height: height of chat window
-    :param show_lora: whether to show LORA options in UI (expert so can be hard to understand)
-    :param login_mode_if_model0: set to True to load --base_model after client logs in, to be able to free GPU memory when model is swapped
-    :param block_gradio_exit: whether to block gradio exit (used for testing)
-    :param concurrency_count: gradio concurrency count (1 is optimal for LLMs)
-    :param api_open: If False, don't let API calls skip gradio queue
-    :param allow_api: whether to allow API calls at all to gradio server
-    :param input_lines: how many input lines to show for chat box (>1 forces shift-enter for submit, else enter is submit)
-    :param auth: gradio auth for launcher in form [(user1, pass1), (user2, pass2), ...]
-                 e.g. --auth=[('jon','password')] with no spaces
-    :param sanitize_user_prompt: whether to remove profanity from user input
-    :param sanitize_bot_response: whether to remove profanity and repeat lines from bot output
-    :param extra_model_options: extra models to show in list in gradio
-    :param extra_lora_options: extra LORA to show in list in gradio
-    :param score_model: which model to score responses (None means no scoring)
-    :param auto_score: whether to automatically score responses
-    :param eval_filename: json file to use for evaluation, if None is sharegpt
-    :param eval_prompts_only_num: for no gradio benchmark, if using eval_filename prompts for eval instead of examples
-    :param eval_prompts_only_seed: for no gradio benchmark, seed for eval_filename sampling
-    :param eval_as_output: for no gradio benchmark, whether to test eval_filename output itself
-    :param langchain_mode: Data source to include.  Choose "UserData" to only consume files from make_db.py.
-           WARNING: wiki_full requires extra data processing via read_wiki_full.py and requires really good workstation to generate db, unless already present.
-    :param user_path: user path to glob from to generate db for vector search, for 'UserData' langchain mode.
-           If already have db, any new/changed files are added automatically if path set, does not have to be same path used for prior db sources
-    :param detect_user_path_changes_every_query: whether to detect if any files changed or added every similarity search (by file hashes).
-           Expensive for large number of files, so not done by default.  By default only detect changes during db loading.
-    :param visible_langchain_modes: dbs to generate at launch to be ready for LLM
-           Can be up to ['wiki', 'wiki_full', 'UserData', 'MyData', 'github h2oGPT', 'DriverlessAI docs']
-           But wiki_full is expensive and requires preparation
-           To allow scratch space only live in session, add 'MyData' to list
-           Default: If only want to consume local files, e.g. prepared by make_db.py, only include ['UserData']
-           FIXME: Avoid 'All' for now, not implemented
-    :param document_choice: Default document choice when taking subset of collection
-    :param load_db_if_exists: Whether to load chroma db if exists or re-generate db
-    :param keep_sources_in_context: Whether to keep url sources in context, not helpful usually
-    :param db_type: 'faiss' for in-memory or 'chroma' or 'weaviate' for persisted on disk
-    :param use_openai_embedding: Whether to use OpenAI embeddings for vector db
-    :param use_openai_model: Whether to use OpenAI model for use with vector db
-    :param hf_embedding_model: Which HF embedding model to use for vector db
-    :param allow_upload_to_user_data: Whether to allow file uploads to update shared vector db
-    :param allow_upload_to_my_data: Whether to allow file uploads to update scratch vector db
-    :param enable_url_upload: Whether to allow upload from URL
-    :param enable_text_upload: Whether to allow upload of text
-    :param enable_sources_list: Whether to allow list (or download for non-shared db) of list of sources for chosen db
-    :param chunk: Whether to chunk data (True unless know data is already optimally chunked)
-    :param chunk_size: Size of chunks, with typically top-4 passed to LLM, so neesd to be in context length
-    :param top_k_docs: number of chunks to give LLM
-    :param n_jobs: Number of processors to use when consuming documents (-1 = all, is default)
-    :param enable_captions: Whether to support captions using BLIP for image files as documents, then preloads that model
-    :param captions_model: Which model to use for captions.
-           captions_model: int = "Salesforce/blip-image-captioning-base",  # continue capable
-           captions_model: str = "Salesforce/blip2-flan-t5-xl",   # question/answer capable, 16GB state
-           captions_model: int = "Salesforce/blip2-flan-t5-xxl",  # question/answer capable, 60GB state
-           Note: opt-based blip2 are not permissive license due to opt and Meta license restrictions
-    :param pre_load_caption_model: Whether to preload caption model, or load after forking parallel doc loader
-           parallel loading disabled if preload and have images, to prevent deadlocking on cuda context
-           Recommended if using larger caption model
-    :param caption_gpu: If support caption, then use GPU if exists
-    :param enable_ocr: Whether to support OCR on images
-    :return:
-    """
-    is_hf = bool(os.getenv("HUGGINGFACE_SPACES"))
-    is_gpth2oai = bool(os.getenv("GPT_H2O_AI"))
-    is_public = is_hf or is_gpth2oai  # multi-user case with fixed model and disclaimer
-    if memory_restriction_level is None:
-        memory_restriction_level = 2 if is_hf else 0  # 2 assumes run on 24GB consumer GPU
-    else:
-        assert 0 <= memory_restriction_level <= 3, "Bad memory_restriction_level=%s" % memory_restriction_level
-    admin_pass = os.getenv("ADMIN_PASS")
-    # will sometimes appear in UI or sometimes actual generation, but maybe better than empty result
-    # but becomes unrecoverable sometimes if raise, so just be silent for now
-    raise_generate_gpu_exceptions = True
-
-    # allow set token directly
-    use_auth_token = os.environ.get("HUGGINGFACE_API_TOKEN", use_auth_token)
-    allow_upload_to_user_data = bool(os.environ.get("allow_upload_to_user_data", allow_upload_to_user_data))
-    allow_upload_to_my_data = bool(os.environ.get("allow_upload_to_my_data", allow_upload_to_my_data))
-    height = os.environ.get("HEIGHT", height)
-
-    # allow enabling langchain via ENV
-    # FIRST PLACE where LangChain referenced, but no imports related to it
-    langchain_mode = os.environ.get("LANGCHAIN_MODE", langchain_mode)
-    assert langchain_mode in langchain_modes, "Invalid langchain_mode %s" % langchain_mode
-    visible_langchain_modes = ast.literal_eval(os.environ.get("visible_langchain_modes", str(visible_langchain_modes)))
-    if langchain_mode not in visible_langchain_modes and langchain_mode in langchain_modes:
-        visible_langchain_modes += [langchain_mode]
-
-    if is_public:
-        allow_upload_to_user_data = False
-        input_lines = 1  # ensure set, for ease of use
-        temperature = 0.2 if temperature is None else temperature
-        top_p = 0.85 if top_p is None else top_p
-        top_k = 70 if top_k is None else top_k
-        if is_hf:
-            do_sample = True if do_sample is None else do_sample
-        else:
-            # by default don't sample, too chatty
-            do_sample = False if do_sample is None else do_sample
-
-        if memory_restriction_level == 2:
-            if not base_model:
-                base_model = 'h2oai/h2ogpt-oasst1-512-12b'
-                # don't set load_8bit if passed base_model, doesn't always work so can't just override
-                load_8bit = True
-                load_4bit = False  # FIXME - consider using 4-bit instead of 8-bit
-        else:
-            base_model = 'h2oai/h2ogpt-oasst1-512-20b' if not base_model else base_model
-    if memory_restriction_level >= 2:
-        load_8bit = True
-        load_4bit = False  # FIXME - consider using 4-bit instead of 8-bit
-    if is_hf:
-        # must override share if in spaces
-        share = False
-    save_dir = os.getenv('SAVE_DIR', save_dir)
-    score_model = os.getenv('SCORE_MODEL', score_model)
-    if score_model == 'None' or score_model is None:
-        score_model = ''
-    concurrency_count = int(os.getenv('CONCURRENCY_COUNT', concurrency_count))
-    api_open = bool(int(os.getenv('API_OPEN', api_open)))
-    allow_api = bool(int(os.getenv('ALLOW_API', allow_api)))
-
-    n_gpus = torch.cuda.device_count() if torch.cuda.is_available else 0
-    if n_gpus == 0:
-        gpu_id = None
-        load_8bit = False
-        load_4bit = False
-        load_half = False
-        infer_devices = False
-        torch.backends.cudnn.benchmark = True
-        torch.backends.cudnn.enabled = False
-        torch.set_default_dtype(torch.float32)
-        if psutil.virtual_memory().available < 94 * 1024 ** 3:
-            # 12B uses ~94GB
-            # 6.9B uses ~47GB
-            base_model = 'h2oai/h2ogpt-oig-oasst1-512-6_9b' if not base_model else base_model
-
-    # get defaults
-    model_lower = base_model.lower()
-    if not gradio:
-        # force, else not single response like want to look at
-        stream_output = False
-        # else prompt removal can mess up output
-        chat = False
-    # hard-coded defaults
-    first_para = False
-    text_limit = None
-
-    if offload_folder:
-        makedirs(offload_folder)
-
-    placeholder_instruction, placeholder_input, \
-        stream_output, show_examples, \
-        prompt_type, temperature, top_p, top_k, num_beams, \
-        max_new_tokens, min_new_tokens, early_stopping, max_time, \
-        repetition_penalty, num_return_sequences, \
-        do_sample, \
-        src_lang, tgt_lang, \
-        examples, \
-        task_info = \
-        get_generate_params(model_lower, chat,
-                            stream_output, show_examples,
-                            prompt_type, temperature, top_p, top_k, num_beams,
-                            max_new_tokens, min_new_tokens, early_stopping, max_time,
-                            repetition_penalty, num_return_sequences,
-                            do_sample,
-                            top_k_docs,
-                            verbose,
-                            )
-
-    locals_dict = locals()
-    locals_print = '\n'.join(['%s: %s' % (k, v) for k, v in locals_dict.items()])
-    if verbose:
-        print(f"Generating model with params:\n{locals_print}", flush=True)
-        print("Command: %s\nHash: %s" % (str(' '.join(sys.argv)), get_githash()), flush=True)
-
-    if langchain_mode != "Disabled":
-        # SECOND PLACE where LangChain referenced, but all imports are kept local so not required
-        from gpt_langchain import prep_langchain, get_some_dbs_from_hf
-        if is_hf:
-            get_some_dbs_from_hf()
-        dbs = {}
-        for langchain_mode1 in visible_langchain_modes:
-            if langchain_mode1 in ['MyData']:
-                # don't use what is on disk, remove it instead
-                for gpath1 in glob.glob(os.path.join(scratch_base_dir, 'db_dir_%s*' % langchain_mode1)):
-                    if os.path.isdir(gpath1):
-                        print("Removing old MyData: %s" % gpath1, flush=True)
-                        shutil.rmtree(gpath1)
-                continue
-            if langchain_mode1 in ['All']:
-                # FIXME: All should be avoided until scans over each db, shouldn't be separate db
-                continue
-            persist_directory1 = 'db_dir_%s' % langchain_mode1  # single place, no special names for each case
-            db = prep_langchain(persist_directory1,
-                                load_db_if_exists,
-                                db_type, use_openai_embedding,
-                                langchain_mode1, user_path,
-                                hf_embedding_model,
-                                kwargs_make_db=locals())
-            dbs[langchain_mode1] = db
-        # remove None db's so can just rely upon k in dbs for if hav db
-        dbs = {k: v for k, v in dbs.items() if v is not None}
-    else:
-        dbs = {}
-        # import control
-        if os.environ.get("TEST_LANGCHAIN_IMPORT"):
-            assert 'gpt_langchain' not in sys.modules, "Dev bug, import of langchain when should not have"
-            assert 'langchain' not in sys.modules, "Dev bug, import of langchain when should not have"
-
-    if cli:
-        from cli import run_cli
-        return run_cli(**get_kwargs(run_cli, exclude_names=['model_state0'], **locals()))
-    elif not gradio:
-        from eval import run_eval
-        return run_eval(**get_kwargs(run_eval, exclude_names=['model_state0'], **locals()))
-    elif gradio:
-        # imported here so don't require gradio to run generate
-        from gradio_runner import go_gradio
-
-        # get default model
-        all_kwargs = locals().copy()
-        if all_kwargs.get('base_model') and not all_kwargs['login_mode_if_model0']:
-            model0, tokenizer0, device = get_model(reward_type=False,
-                                                   **get_kwargs(get_model, exclude_names=['reward_type'], **all_kwargs))
-        else:
-            # if empty model, then don't load anything, just get gradio up
-            model0, tokenizer0, device = None, None, None
-        model_state0 = [model0, tokenizer0, device, all_kwargs['base_model']]
-
-        # get score model
-        smodel, stokenizer, sdevice = get_score_model(reward_type=True,
-                                                      **get_kwargs(get_score_model, exclude_names=['reward_type'],
-                                                                   **all_kwargs))
-        score_model_state0 = [smodel, stokenizer, sdevice, score_model]
-
-        if enable_captions:
-            if pre_load_caption_model:
-                from image_captions import H2OImageCaptionLoader
-                caption_loader = H2OImageCaptionLoader(caption_gpu=caption_gpu).load_model()
-            else:
-                caption_loader = 'gpu' if caption_gpu else 'cpu'
-        else:
-            caption_loader = False
-
-        # assume gradio needs everything
-        go_gradio(**locals())
-
-
-def get_non_lora_model(base_model, model_loader, load_half, model_kwargs, reward_type,
-                       gpu_id=0,
-                       use_auth_token=False,
-                       trust_remote_code=True,
-                       offload_folder=None,
-                       triton_attn=False,
-                       long_sequence=True,
-                       ):
-    """
-    Ensure model gets on correct device
-    :param base_model:
-    :param model_loader:
-    :param load_half:
-    :param model_kwargs:
-    :param reward_type:
-    :param gpu_id:
-    :param use_auth_token:
-    :param trust_remote_code:
-    :param offload_folder:
-    :param triton_attn:
-    :param long_sequence:
-    :return:
-    """
-    with init_empty_weights():
-        from transformers import AutoConfig
-        config = AutoConfig.from_pretrained(base_model, use_auth_token=use_auth_token,
-                                            trust_remote_code=trust_remote_code,
-                                            offload_folder=offload_folder)
-        if triton_attn and 'mpt-' in base_model.lower():
-            config.attn_config['attn_impl'] = 'triton'
-        if long_sequence:
-            if 'mpt-7b-storywriter' in base_model.lower():
-                config.update({"max_seq_len": 83968})
-            if 'mosaicml/mpt-7b-chat' in base_model.lower():
-                config.update({"max_seq_len": 4096})
-        if issubclass(config.__class__, tuple(AutoModel._model_mapping.keys())):
-            model = AutoModel.from_config(
-                config,
-            )
-        else:
-            # can't infer
-            model = None
-
-    if model is not None:
-        # NOTE: Can specify max_memory={0: max_mem, 1: max_mem}, to shard model
-        # NOTE: Some models require avoiding sharding some layers,
-        # then would pass no_split_module_classes and give list of those layers.
-        device_map = infer_auto_device_map(
-            model,
-            dtype=torch.float16 if load_half else torch.float32,
-        )
-        if hasattr(model, 'model'):
-            device_map_model = infer_auto_device_map(
-                model.model,
-                dtype=torch.float16 if load_half else torch.float32,
-            )
-            device_map.update(device_map_model)
-    else:
-        device_map = "auto"
-
-    n_gpus = torch.cuda.device_count() if torch.cuda.is_available else 0
-
-    if n_gpus > 0:
-        if gpu_id >= 0:
-            # FIXME: If really distributes model, tend to get things like: ValueError: gpt_neox.embed_in.weight doesn't have any device set.
-            # So avoid for now, just put on first GPU, unless score_model, put on last
-            if reward_type:
-                device_map = {'': n_gpus - 1}
-            else:
-                device_map = {'': min(n_gpus - 1, gpu_id)}
-        if gpu_id == -1:
-            device_map = {'': 'cuda'}
-    else:
-        device_map = {'': 'cpu'}
-        model_kwargs['load_in_8bit'] = False
-        model_kwargs['load_in_4bit'] = False
-    print('device_map: %s' % device_map, flush=True)
-
-    load_in_8bit = model_kwargs.get('load_in_8bit', False)
-    load_in_4bit = model_kwargs.get('load_in_4bit', False)
-    model_kwargs['device_map'] = device_map
-    pop_unused_model_kwargs(model_kwargs)
-
-    if load_in_8bit or load_in_4bit or not load_half:
-        model = model_loader.from_pretrained(
-            base_model,
-            config=config,
-            **model_kwargs,
-        )
-    else:
-        model = model_loader.from_pretrained(
-            base_model,
-            config=config,
-            **model_kwargs,
-        ).half()
-    return model
-
-
-def get_model(
-        load_8bit: bool = False,
-        load_4bit: bool = False,
-        load_half: bool = True,
-        infer_devices: bool = True,
-        base_model: str = '',
-        tokenizer_base_model: str = '',
-        lora_weights: str = "",
-        gpu_id: int = 0,
-
-        reward_type: bool = None,
-        local_files_only: bool = False,
-        resume_download: bool = True,
-        use_auth_token: Union[str, bool] = False,
-        trust_remote_code: bool = True,
-        offload_folder: str = None,
-        compile_model: bool = True,
-
-        verbose: bool = False,
-):
-    """
-
-    :param load_8bit: load model in 8-bit, not supported by all models
-    :param load_4bit: load model in 4-bit, not supported by all models
-    :param load_half: load model in 16-bit
-    :param infer_devices: Use torch infer of optimal placement of layers on devices (for non-lora case)
-           For non-LORA case, False will spread shards across multiple GPUs, but this can lead to cuda:x cuda:y mismatches
-           So it is not the default
-    :param base_model: name/path of base model
-    :param tokenizer_base_model: name/path of tokenizer
-    :param lora_weights: name/path
-    :param gpu_id: which GPU (0..n_gpus-1) or allow all GPUs if relevant (-1)
-    :param reward_type: reward type model for sequence classification
-    :param local_files_only: use local files instead of from HF
-    :param resume_download: resume downloads from HF
-    :param use_auth_token: assumes user did on CLI `huggingface-cli login` to access private repo
-    :param trust_remote_code: trust code needed by model
-    :param offload_folder: offload folder
-    :param compile_model: whether to compile torch model
-    :param verbose:
-    :return:
-    """
-    if verbose:
-        print("Get %s model" % base_model, flush=True)
-    if base_model in non_hf_types:
-        from gpt4all_llm import get_model_tokenizer_gpt4all
-        model, tokenizer, device = get_model_tokenizer_gpt4all(base_model)
-        return model, tokenizer, device
-
-    if lora_weights is not None and lora_weights.strip():
-        if verbose:
-            print("Get %s lora weights" % lora_weights, flush=True)
-    device = get_device()
-
-    if 'gpt2' in base_model.lower():
-        # RuntimeError: where expected condition to be a boolean tensor, but got a tensor with dtype Half
-        load_8bit = False
-        load_4bit = False
-
-    assert base_model.strip(), (
-        "Please choose a base model with --base_model (CLI) or load one from Models Tab (gradio)"
-    )
-
-    from transformers import AutoConfig
-    config = AutoConfig.from_pretrained(base_model, use_auth_token=use_auth_token,
-                                        trust_remote_code=trust_remote_code,
-                                        offload_folder=offload_folder)
-    llama_type_from_config = 'llama' in str(config).lower()
-    llama_type_from_name = "llama" in base_model.lower()
-    llama_type = llama_type_from_config or llama_type_from_name
-    if llama_type:
-        if verbose:
-            print("Detected as llama type from"
-                  " config (%s) or name (%s)" % (llama_type_from_config, llama_type_from_name), flush=True)
-
-    model_loader, tokenizer_loader = get_loaders(llama_type=llama_type, model_name=base_model, reward_type=reward_type)
-    if not tokenizer_base_model:
-        tokenizer_base_model = base_model
-
-    if tokenizer_loader is not None and not isinstance(tokenizer_loader, str):
-        tokenizer = tokenizer_loader.from_pretrained(tokenizer_base_model,
-                                                     local_files_only=local_files_only,
-                                                     resume_download=resume_download,
-                                                     use_auth_token=use_auth_token,
-                                                     trust_remote_code=trust_remote_code,
-                                                     offload_folder=offload_folder,
-                                                     )
-    else:
-        tokenizer = tokenizer_loader
-
-    if isinstance(tokenizer, str):
-        # already a pipeline, tokenizer_loader is string for task
-        model = model_loader(tokenizer,
-                             model=base_model,
-                             device=0 if device == "cuda" else -1,
-                             torch_dtype=torch.float16 if device == 'cuda' else torch.float32)
-    else:
-        assert device in ["cuda", "cpu"], "Unsupported device %s" % device
-        model_kwargs = dict(local_files_only=local_files_only,
-                            torch_dtype=torch.float16 if device == 'cuda' else torch.float32,
-                            resume_download=resume_download,
-                            use_auth_token=use_auth_token,
-                            trust_remote_code=trust_remote_code,
-                            offload_folder=offload_folder,
-                            )
-        if 'mbart-' not in base_model.lower() and 'mpt-' not in base_model.lower():
-            model_kwargs.update(dict(load_in_8bit=load_8bit,
-                                     load_in_4bit=load_4bit,
-                                     device_map={"": 0} if (load_8bit or load_4bit) and device == 'cuda' else "auto",
-                                     ))
-        if 'mpt-' in base_model.lower() and gpu_id >= 0:
-            model_kwargs.update(dict(device_map={"": gpu_id} if device == 'cuda' else "cpu"))
-
-        if 'OpenAssistant/reward-model'.lower() in base_model.lower():
-            # FIXME: could put on other GPUs
-            model_kwargs['device_map'] = {"": 0} if device == 'cuda' else {"": 'cpu'}
-            model_kwargs.pop('torch_dtype', None)
-        pop_unused_model_kwargs(model_kwargs)
-
-        if not lora_weights:
-            with torch.device(device):
-                if infer_devices:
-                    model = get_non_lora_model(base_model, model_loader, load_half, model_kwargs, reward_type,
-                                               gpu_id=gpu_id,
-                                               use_auth_token=use_auth_token,
-                                               trust_remote_code=trust_remote_code,
-                                               offload_folder=offload_folder,
-                                               )
-                else:
-                    if load_half and not (load_8bit or load_4bit):
-                        model = model_loader.from_pretrained(
-                            base_model,
-                            **model_kwargs).half()
-                    else:
-                        model = model_loader.from_pretrained(
-                            base_model,
-                            **model_kwargs)
-        elif load_8bit or load_4bit:
-            model = model_loader.from_pretrained(
-                base_model,
-                **model_kwargs
-            )
-            model = PeftModel.from_pretrained(
-                model,
-                lora_weights,
-                torch_dtype=torch.float16 if device == 'cuda' else torch.float32,
-                local_files_only=local_files_only,
-                resume_download=resume_download,
-                use_auth_token=use_auth_token,
-                trust_remote_code=trust_remote_code,
-                offload_folder=offload_folder,
-                device_map={"": 0} if device == 'cuda' else {"": 'cpu'},  # seems to be required
-            )
-        else:
-            with torch.device(device):
-                model = model_loader.from_pretrained(
-                    base_model,
-                    **model_kwargs
-                )
-                model = PeftModel.from_pretrained(
-                    model,
-                    lora_weights,
-                    torch_dtype=torch.float16 if device == 'cuda' else torch.float32,
-                    local_files_only=local_files_only,
-                    resume_download=resume_download,
-                    use_auth_token=use_auth_token,
-                    trust_remote_code=trust_remote_code,
-                    offload_folder=offload_folder,
-                    device_map="auto",
-                )
-                if load_half:
-                    model.half()
-
-    # unwind broken decapoda-research config
-    if llama_type:
-        model.config.pad_token_id = tokenizer.pad_token_id = 0  # unk
-        model.config.bos_token_id = 1
-        model.config.eos_token_id = 2
-    if 'gpt2' in base_model.lower():
-        # add special tokens that otherwise all share the same id
-        tokenizer.add_special_tokens({'bos_token': '<bos>',
-                                      'eos_token': '<eos>',
-                                      'pad_token': '<pad>'})
-
-    if not isinstance(tokenizer, str):
-        model.eval()
-        if torch.__version__ >= "2" and sys.platform != "win32" and compile_model:
-            model = torch.compile(model)
-
-    if hasattr(config, 'max_position_embeddings') and isinstance(config.max_position_embeddings, int):
-        # help automatically limit inputs to generate
-        tokenizer.model_max_length = config.max_position_embeddings
-    else:
-        tokenizer.model_max_length = 2048
-
-    return model, tokenizer, device
-
-
-def pop_unused_model_kwargs(model_kwargs):
-    """
-    in-place pop unused kwargs that are not dependency-upgrade friendly
-    no point passing in False, is default, and helps avoid needing to update requirements for new deps
-    :param model_kwargs:
-    :return:
-    """
-    check_list = ['load_in_8bit', 'load_in_4bit']
-    for k in check_list:
-        if k in model_kwargs and not model_kwargs[k]:
-            model_kwargs.pop(k)
-
-
-def get_score_model(score_model: str = None,
-                    load_8bit: bool = False,
-                    load_4bit: bool = False,
-                    load_half: bool = True,
-                    infer_devices: bool = True,
-                    base_model: str = '',
-                    tokenizer_base_model: str = '',
-                    lora_weights: str = "",
-                    gpu_id: int = 0,
-
-                    reward_type: bool = None,
-                    local_files_only: bool = False,
-                    resume_download: bool = True,
-                    use_auth_token: Union[str, bool] = False,
-                    trust_remote_code: bool = True,
-                    offload_folder: str = None,
-                    compile_model: bool = True,
-
-                    verbose: bool = False,
-                    ):
-    if score_model is not None and score_model.strip():
-        load_8bit = False
-        load_4bit = False
-        load_half = False
-        base_model = score_model.strip()
-        tokenizer_base_model = ''
-        lora_weights = ''
-        llama_type = False
-        compile_model = False
-        smodel, stokenizer, sdevice = get_model(reward_type=True,
-                                                **get_kwargs(get_model, exclude_names=['reward_type'], **locals()))
-    else:
-        smodel, stokenizer, sdevice = None, None, None
-    return smodel, stokenizer, sdevice
-
-
-eval_func_param_names = ['instruction',
-                         'iinput',
-                         'context',
-                         'stream_output',
-                         'prompt_type',
-                         'temperature',
-                         'top_p',
-                         'top_k',
-                         'num_beams',
-                         'max_new_tokens',
-                         'min_new_tokens',
-                         'early_stopping',
-                         'max_time',
-                         'repetition_penalty',
-                         'num_return_sequences',
-                         'do_sample',
-                         'chat',
-                         'instruction_nochat',
-                         'iinput_nochat',
-                         'langchain_mode',
-                         'top_k_docs',
-                         'document_choice',
-                         ]
-
-
-def evaluate(
-        model_state,
-        my_db_state,
-        # START NOTE: Examples must have same order of parameters
-        instruction,
-        iinput,
-        context,
-        stream_output,
-        prompt_type,
-        temperature,
-        top_p,
-        top_k,
-        num_beams,
-        max_new_tokens,
-        min_new_tokens,
-        early_stopping,
-        max_time,
-        repetition_penalty,
-        num_return_sequences,
-        do_sample,
-        chat,
-        instruction_nochat,
-        iinput_nochat,
-        langchain_mode,
-        top_k_docs,
-        document_choice,
-        # END NOTE: Examples must have same order of parameters
-        src_lang=None,
-        tgt_lang=None,
-        debug=False,
-        concurrency_count=None,
-        save_dir=None,
-        sanitize_bot_response=True,
-        model_state0=None,
-        memory_restriction_level=None,
-        raise_generate_gpu_exceptions=None,
-        chat_context=None,
-        lora_weights=None,
-        load_db_if_exists=True,
-        dbs=None,
-        user_path=None,
-        detect_user_path_changes_every_query=None,
-        use_openai_embedding=None,
-        use_openai_model=None,
-        hf_embedding_model=None,
-        chunk=None,
-        chunk_size=None,
-        db_type=None,
-        n_jobs=None,
-        first_para=None,
-        text_limit=None,
-        verbose=False,
-        cli=False,
-):
-    # ensure passed these
-    assert concurrency_count is not None
-    assert memory_restriction_level is not None
-    assert raise_generate_gpu_exceptions is not None
-    assert chat_context is not None
-    assert use_openai_embedding is not None
-    assert use_openai_model is not None
-    assert hf_embedding_model is not None
-    assert chunk is not None
-    assert chunk_size is not None
-    assert db_type is not None
-    assert top_k_docs is not None and isinstance(top_k_docs, int)
-    assert n_jobs is not None
-    assert first_para is not None
-
-    if debug:
-        locals_dict = locals().copy()
-        locals_dict.pop('model_state', None)
-        locals_dict.pop('model_state0', None)
-        print(locals_dict)
-
-    no_model_msg = "Please choose a base model with --base_model (CLI) or load in Models Tab (gradio).\nThen start New Conversation"
-
-    if model_state0 is None:
-        # e.g. for no gradio case, set dummy value, else should be set
-        model_state0 = [None, None, None, None]
-
-    if model_state is not None and len(model_state) == 4 and not isinstance(model_state[0], str):
-        # try to free-up original model (i.e. list was passed as reference)
-        if model_state0 is not None and model_state0[0] is not None:
-            model_state0[0].cpu()
-            model_state0[0] = None
-        # try to free-up original tokenizer (i.e. list was passed as reference)
-        if model_state0 is not None and model_state0[1] is not None:
-            model_state0[1] = None
-        clear_torch_cache()
-        model, tokenizer, device, base_model = model_state
-    elif model_state0 is not None and len(model_state0) == 4 and model_state0[0] is not None:
-        assert isinstance(model_state[0], str)
-        model, tokenizer, device, base_model = model_state0
-    else:
-        raise AssertionError(no_model_msg)
-
-    if base_model is None:
-        raise AssertionError(no_model_msg)
-
-    assert base_model.strip(), no_model_msg
-    assert model, "Model is missing"
-    assert tokenizer, "Tokenizer is missing"
-
-    # choose chat or non-chat mode
-    if not chat:
-        instruction = instruction_nochat
-        iinput = iinput_nochat
-
-    if not context:
-        # get hidden context if have one
-        context = get_context(chat_context, prompt_type)
-
-    prompter = Prompter(prompt_type, debug=debug, chat=chat, stream_output=stream_output)
-    data_point = dict(context=context, instruction=instruction, input=iinput)
-    prompt = prompter.generate_prompt(data_point)
-
-    # THIRD PLACE where LangChain referenced, but imports only occur if enabled and have db to use
-    assert langchain_mode in langchain_modes, "Invalid langchain_mode %s" % langchain_mode
-    if langchain_mode in ['MyData'] and my_db_state is not None and len(my_db_state) > 0 and my_db_state[0] is not None:
-        db1 = my_db_state[0]
-    elif dbs is not None and langchain_mode in dbs:
-        db1 = dbs[langchain_mode]
-    else:
-        db1 = None
-    if langchain_mode not in [False, 'Disabled', 'ChatLLM', 'LLM'] and db1 is not None or base_model in non_hf_types:
-        query = instruction if not iinput else "%s\n%s" % (instruction, iinput)
-        outr = ""
-        # use smaller cut_distanct for wiki_full since so many matches could be obtained, and often irrelevant unless close
-        from gpt_langchain import run_qa_db
-        for r in run_qa_db(query=query,
-                           model_name=base_model, model=model, tokenizer=tokenizer,
-                           stream_output=stream_output,
-                           prompter=prompter,
-                           load_db_if_exists=load_db_if_exists,
-                           db=db1,
-                           user_path=user_path,
-                           detect_user_path_changes_every_query=detect_user_path_changes_every_query,
-                           max_new_tokens=max_new_tokens,
-                           cut_distanct=1.1 if langchain_mode in ['wiki_full'] else 1.64,  # FIXME, too arbitrary
-                           use_openai_embedding=use_openai_embedding,
-                           use_openai_model=use_openai_model,
-                           hf_embedding_model=hf_embedding_model,
-                           first_para=first_para,
-                           text_limit=text_limit,
-                           chunk=chunk,
-                           chunk_size=chunk_size,
-                           langchain_mode=langchain_mode,
-                           document_choice=document_choice,
-                           db_type=db_type,
-                           k=top_k_docs,
-                           temperature=temperature,
-                           repetition_penalty=repetition_penalty,
-                           top_k=top_k,
-                           top_p=top_p,
-                           prompt_type=prompt_type,
-                           n_jobs=n_jobs,
-                           verbose=verbose,
-                           cli=cli,
-                           ):
-            outr, extra = r  # doesn't accumulate, new answer every yield, so only save that full answer
-            yield dict(response=outr, sources=extra)
-        if save_dir:
-            save_generate_output(output=outr, base_model=base_model, save_dir=save_dir)
-            if verbose:
-                print(
-                    'Post-Generate Langchain: %s decoded_output: %s' % (str(datetime.now()), len(outr) if outr else -1),
-                    flush=True)
-        if outr or base_model in non_hf_types:
-            # if got no response (e.g. not showing sources and got no sources,
-            # so nothing to give to LLM), then slip through and ask LLM
-            # Or if llama/gptj, then just return since they had no response and can't go down below code path
-            return
-
-    if isinstance(tokenizer, str):
-        # pipeline
-        if tokenizer == "summarization":
-            key = 'summary_text'
-        else:
-            raise RuntimeError("No such task type %s" % tokenizer)
-        # NOTE: uses max_length only
-        yield dict(response=model(prompt, max_length=max_new_tokens)[0][key], sources='')
-
-    if 'mbart-' in base_model.lower():
-        assert src_lang is not None
-        tokenizer.src_lang = languages_covered()[src_lang]
-
-    if chat:
-        # override, ignore user change
-        num_return_sequences = 1
-    stopping_criteria = get_stopping(prompt_type, tokenizer, device)
-    _, _, max_length_tokenize, max_prompt_length = get_cutoffs(memory_restriction_level, model_max_length=tokenizer.model_max_length)
-    prompt = prompt[-max_prompt_length:]
-    inputs = tokenizer(prompt,
-                       return_tensors="pt",
-                       truncation=True,
-                       max_length=max_length_tokenize)
-    if inputs['input_ids'].shape[1] >= max_length_tokenize - 1:
-        print("Cutting off input: %s %s" % (inputs['input_ids'].shape[1], max_length_tokenize), flush=True)
-    if debug and len(inputs["input_ids"]) > 0:
-        print('input_ids length', len(inputs["input_ids"][0]), flush=True)
-    input_ids = inputs["input_ids"].to(device)
-    # CRITICAL LIMIT else will fail
-    max_max_tokens = tokenizer.model_max_length
-    max_input_tokens = max_max_tokens - max_new_tokens
-    input_ids = input_ids[:, -max_input_tokens:]
-    generation_config = GenerationConfig(
-        temperature=float(temperature),
-        top_p=float(top_p),
-        top_k=top_k,
-        num_beams=num_beams,
-        do_sample=do_sample,
-        repetition_penalty=float(repetition_penalty),
-        num_return_sequences=num_return_sequences,
-        renormalize_logits=True,
-        remove_invalid_values=True,
-    )
-
-    gen_kwargs = dict(input_ids=input_ids,
-                      generation_config=generation_config,
-                      return_dict_in_generate=True,
-                      output_scores=True,
-                      max_new_tokens=max_new_tokens,  # prompt + new
-                      min_new_tokens=min_new_tokens,  # prompt + new
-                      early_stopping=early_stopping,  # False, True, "never"
-                      max_time=max_time,
-                      stopping_criteria=stopping_criteria,
-                      )
-    if 'gpt2' in base_model.lower():
-        gen_kwargs.update(dict(bos_token_id=tokenizer.bos_token_id, pad_token_id=tokenizer.eos_token_id))
-    elif 'mbart-' in base_model.lower():
-        assert tgt_lang is not None
-        tgt_lang = languages_covered()[tgt_lang]
-        gen_kwargs.update(dict(forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang]))
-    else:
-        gen_kwargs.update(dict(pad_token_id=tokenizer.eos_token_id))
-
-    decoder_kwargs = dict(skip_special_tokens=True,
-                          clean_up_tokenization_spaces=True)
-
-    decoder = functools.partial(tokenizer.decode,
-                                **decoder_kwargs
-                                )
-    decoder_raw_kwargs = dict(skip_special_tokens=False,
-                              clean_up_tokenization_spaces=True)
-
-    decoder_raw = functools.partial(tokenizer.decode,
-                                    **decoder_raw_kwargs
-                                    )
-
-    with torch.no_grad():
-        context_class_cast = NullContext if device == 'cpu' or lora_weights else torch.autocast
-        with context_class_cast(device):
-            # protection for gradio not keeping track of closed users,
-            # else hit bitsandbytes lack of thread safety:
-            # https://github.com/h2oai/h2ogpt/issues/104
-            # but only makes sense if concurrency_count == 1
-            context_class = NullContext  # if concurrency_count > 1 else filelock.FileLock
-            if verbose:
-                print('Pre-Generate: %s' % str(datetime.now()), flush=True)
-            decoded_output = None
-            with context_class("generate.lock"):
-                if verbose:
-                    print('Generate: %s' % str(datetime.now()), flush=True)
-                # decoded tokenized prompt can deviate from prompt due to special characters
-                inputs_decoded = decoder(input_ids[0])
-                inputs_decoded_raw = decoder_raw(input_ids[0])
-                if inputs_decoded == prompt:
-                    # normal
-                    pass
-                elif inputs_decoded.lstrip() == prompt.lstrip():
-                    # sometimes extra space in front, make prompt same for prompt removal
-                    prompt = inputs_decoded
-                elif inputs_decoded_raw == prompt:
-                    # some models specify special tokens that are part of normal prompt, so can't skip them
-                    inputs_decoded = prompt = inputs_decoded_raw
-                    decoder = decoder_raw
-                    decoder_kwargs = decoder_raw_kwargs
-                elif inputs_decoded_raw.replace("<unk> ", "").replace("<unk>", "").replace('\n', ' ').replace(' ',
-                                                                                                              '') == prompt.replace(
-                    '\n', ' ').replace(' ', ''):
-                    inputs_decoded = prompt = inputs_decoded_raw
-                    decoder = decoder_raw
-                    decoder_kwargs = decoder_raw_kwargs
-                else:
-                    if verbose:
-                        print("WARNING: Special characters in prompt", flush=True)
-                if stream_output:
-                    skip_prompt = False
-                    streamer = H2OTextIteratorStreamer(tokenizer, skip_prompt=skip_prompt, block=False,
-                                                       **decoder_kwargs)
-                    gen_kwargs.update(dict(streamer=streamer))
-                    target = wrapped_partial(generate_with_exceptions, model.generate,
-                                             prompt=prompt, inputs_decoded=inputs_decoded,
-                                             raise_generate_gpu_exceptions=raise_generate_gpu_exceptions,
-                                             **gen_kwargs)
-                    bucket = queue.Queue()
-                    thread = EThread(target=target, streamer=streamer, bucket=bucket)
-                    thread.start()
-                    outputs = ""
-                    try:
-                        for new_text in streamer:
-                            if bucket.qsize() > 0 or thread.exc:
-                                thread.join()
-                            outputs += new_text
-                            yield dict(response=prompter.get_response(outputs, prompt=inputs_decoded,
-                                                                      sanitize_bot_response=sanitize_bot_response),
-                                       sources='')
-                    except BaseException:
-                        # if any exception, raise that exception if was from thread, first
-                        if thread.exc:
-                            raise thread.exc
-                        raise
-                    finally:
-                        # in case no exception and didn't join with thread yet, then join
-                        if not thread.exc:
-                            thread.join()
-                    # in case raise StopIteration or broke queue loop in streamer, but still have exception
-                    if thread.exc:
-                        raise thread.exc
-                    decoded_output = outputs
-                else:
-                    outputs = model.generate(**gen_kwargs)
-                    outputs = [decoder(s) for s in outputs.sequences]
-                    yield dict(response=prompter.get_response(outputs, prompt=inputs_decoded,
-                                                              sanitize_bot_response=sanitize_bot_response), sources='')
-                    if outputs and len(outputs) >= 1:
-                        decoded_output = prompt + outputs[0]
-                if save_dir and decoded_output:
-                    save_generate_output(output=decoded_output, base_model=base_model, save_dir=save_dir)
-            if verbose:
-                print('Post-Generate: %s decoded_output: %s' % (
-                    str(datetime.now()), len(decoded_output) if decoded_output else -1), flush=True)
-
-
-inputs_list_names = list(inspect.signature(evaluate).parameters)
-state_names = ['model_state', 'my_db_state']
-inputs_kwargs_list = [x for x in inputs_list_names if x not in eval_func_param_names + state_names]
-
-
-def get_cutoffs(memory_restriction_level, for_context=False, model_max_length=2048):
-    # help to avoid errors like:
-    # RuntimeError: The size of tensor a (2048) must match the size of tensor b (2049) at non-singleton dimension 3
-    # RuntimeError: expected scalar type Half but found Float
-    # with - 256
-    if memory_restriction_level > 0:
-        max_length_tokenize = 768 - 256 if memory_restriction_level <= 2 else 512 - 256
-    else:
-        max_length_tokenize = model_max_length - 256
-    cutoff_len = max_length_tokenize * 4  # if reaches limit, then can't generate new tokens
-    output_smallest = 30 * 4
-    max_prompt_length = cutoff_len - output_smallest
-
-    if for_context:
-        # then lower even more to avoid later chop, since just estimate tokens in context bot
-        max_prompt_length = max(64, int(max_prompt_length * 0.8))
-
-    return cutoff_len, output_smallest, max_length_tokenize, max_prompt_length
-
-
-class H2OTextIteratorStreamer(TextIteratorStreamer):
-    """
-    normally, timeout required for now to handle exceptions, else get()
-    but with H2O version of TextIteratorStreamer, loop over block to handle
-    """
-
-    def __init__(self, tokenizer, skip_prompt: bool = False, timeout: typing.Optional[float] = None,
-                 block=True, **decode_kwargs):
-        super().__init__(tokenizer, skip_prompt, **decode_kwargs)
-        self.text_queue = queue.Queue()
-        self.stop_signal = None
-        self.do_stop = False
-        self.timeout = timeout
-        self.block = block
-
-    def on_finalized_text(self, text: str, stream_end: bool = False):
-        """Put the new text in the queue. If the stream is ending, also put a stop signal in the queue."""
-        self.text_queue.put(text, timeout=self.timeout)
-        if stream_end:
-            self.text_queue.put(self.stop_signal, timeout=self.timeout)
-
-    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
-
-
-def generate_with_exceptions(func, *args, prompt='', inputs_decoded='', raise_generate_gpu_exceptions=True, **kwargs):
-    try:
-        func(*args, **kwargs)
-    except torch.cuda.OutOfMemoryError as e:
-        print("GPU OOM 2: prompt: %s inputs_decoded: %s exception: %s" % (prompt, inputs_decoded, str(e)),
-              flush=True)
-        if 'input_ids' in kwargs:
-            if kwargs['input_ids'] is not None:
-                kwargs['input_ids'].cpu()
-            kwargs['input_ids'] = None
-        traceback.print_exc()
-        clear_torch_cache()
-        return
-    except (Exception, RuntimeError) as e:
-        if 'Expected all tensors to be on the same device' in str(e) or \
-                'expected scalar type Half but found Float' in str(e) or \
-                'probability tensor contains either' in str(e) or \
-                'cublasLt ran into an error!' in str(e) or \
-                'mat1 and mat2 shapes cannot be multiplied' in str(e):
-            print(
-                "GPU Error: prompt: %s inputs_decoded: %s exception: %s" % (prompt, inputs_decoded, str(e)),
-                flush=True)
-            traceback.print_exc()
-            clear_torch_cache()
-            if raise_generate_gpu_exceptions:
-                raise
-            return
-        else:
-            clear_torch_cache()
-            if raise_generate_gpu_exceptions:
-                raise
-
-
-def get_generate_params(model_lower, chat,
-                        stream_output, show_examples,
-                        prompt_type, temperature, top_p, top_k, num_beams,
-                        max_new_tokens, min_new_tokens, early_stopping, max_time,
-                        repetition_penalty, num_return_sequences,
-                        do_sample, k, verbose):
-    use_defaults = False
-    use_default_examples = True
-    examples = []
-    task_info = f"{prompt_type}"
-    if model_lower:
-        print(f"Using Model {model_lower}", flush=True)
-    else:
-        print("No model defined yet", flush=True)
-
-    min_new_tokens = min_new_tokens if min_new_tokens is not None else 0
-    early_stopping = early_stopping if early_stopping is not None else False
-    max_time_defaults = 60 * 3
-    max_time = max_time if max_time is not None else max_time_defaults
-
-    if not prompt_type and model_lower in inv_prompt_type_to_model_lower:
-        prompt_type = inv_prompt_type_to_model_lower[model_lower]
-        if verbose:
-            print("Auto-selecting prompt_type=%s for %s" % (prompt_type, model_lower), flush=True)
-
-    # examples at first don't include chat, instruction_nochat, iinput_nochat, added at end
-    if show_examples is None:
-        if chat:
-            show_examples = False
-        else:
-            show_examples = True
-
-    summarize_example1 = """Jeff: Can I train a ? Transformers model on Amazon SageMaker? 
-Philipp: Sure you can use the new Hugging Face Deep Learning Container. 
-Jeff: ok.
-Jeff: and how can I get started? 
-Jeff: where can I find documentation? 
-Philipp: ok, ok you can find everything here. https://huggingface.co/blog/the-partnership-amazon-sagemaker-and-hugging-face"""
-
-    if 'bart-large-cnn-samsum' in model_lower or 'flan-t5-base-samsum' in model_lower:
-        placeholder_instruction = summarize_example1
-        placeholder_input = ""
-        use_defaults = True
-        use_default_examples = False
-        examples += [
-            [placeholder_instruction, "", "", stream_output, 'plain', 1.0, 1.0, 50, 1, 128, 0, False, max_time_defaults,
-             1.0, 1,
-             False]]
-        task_info = "Summarization"
-    elif 't5-' in model_lower or 't5' == model_lower or 'flan-' in model_lower:
-        placeholder_instruction = "The square root of x is the cube root of y. What is y to the power of 2, if x = 4?"
-        placeholder_input = ""
-        use_defaults = True
-        use_default_examples = True
-        task_info = "Multi-Task: Q/A, translation, Chain-of-Thought, Logical Reasoning, Summarization, etc.  Best to use task prefix as trained on, e.g. `translate English to German: ` (space after colon)"
-    elif 'mbart-' in model_lower:
-        placeholder_instruction = "The girl has long hair."
-        placeholder_input = ""
-        use_defaults = True
-        use_default_examples = False
-        examples += [
-            [placeholder_instruction, "", "", stream_output, 'plain', 1.0, 1.0, 50, 1, 128, 0, False, max_time_defaults,
-             1.0, 1,
-             False]]
-    elif 'gpt2' in model_lower:
-        placeholder_instruction = "The sky is"
-        placeholder_input = ""
-        prompt_type = prompt_type or 'plain'
-        use_default_examples = True  # some will be odd "continuations" but can be ok
-        examples += [
-            [placeholder_instruction, "", "", stream_output, 'plain', 1.0, 1.0, 50, 1, 128, 0, False, max_time_defaults,
-             1.0, 1,
-             False]]
-        task_info = "Auto-complete phrase, code, etc."
-        use_defaults = True
-    else:
-        if chat:
-            placeholder_instruction = "Enter a question or imperative."
-        else:
-            placeholder_instruction = "Give detailed answer for whether Einstein or Newton is smarter."
-        placeholder_input = ""
-        if model_lower:
-            # default is plain, because might relly upon trust_remote_code to handle prompting
-            prompt_type = prompt_type or 'plain'
-        else:
-            prompt_type = ''
-        task_info = "No task"
-        if prompt_type == 'instruct':
-            task_info = "Answer question or follow imperative as instruction with optionally input."
-        elif prompt_type == 'plain':
-            task_info = "Auto-complete phrase, code, etc."
-        elif prompt_type == 'human_bot':
-            if chat:
-                task_info = "Chat (Shift-Enter to give question/imperative, input concatenated with instruction)"
-            else:
-                task_info = "Ask question/imperative (input concatenated with instruction)"
-
-    # revert to plain if still nothing
-    prompt_type = prompt_type or 'plain'
-    if use_defaults:
-        temperature = 1.0 if temperature is None else temperature
-        top_p = 1.0 if top_p is None else top_p
-        top_k = 40 if top_k is None else top_k
-        num_beams = num_beams or 1
-        max_new_tokens = max_new_tokens or 128
-        repetition_penalty = repetition_penalty or 1.07
-        num_return_sequences = min(num_beams, num_return_sequences or 1)
-        do_sample = False if do_sample is None else do_sample
-    else:
-        temperature = 0.1 if temperature is None else temperature
-        top_p = 0.75 if top_p is None else top_p
-        top_k = 40 if top_k is None else top_k
-        if chat:
-            num_beams = num_beams or 1
-        else:
-            num_beams = num_beams or 4
-        max_new_tokens = max_new_tokens or 256
-        repetition_penalty = repetition_penalty or 1.07
-        num_return_sequences = min(num_beams, num_return_sequences or 1)
-        do_sample = False if do_sample is None else do_sample
-    # doesn't include chat, instruction_nochat, iinput_nochat, added later
-    params_list = ["", stream_output, prompt_type, temperature, top_p, top_k, num_beams, max_new_tokens, min_new_tokens,
-                   early_stopping, max_time, repetition_penalty, num_return_sequences, do_sample]
-
-    if use_default_examples:
-        examples += [
-            ["Translate English to French", "Good morning"] + params_list,
-            ["Give detailed answer for whether Einstein or Newton is smarter.", ''] + params_list,
-            ["Explain in detailed list, all the best practices for coding in python.", ''] + params_list,
-            [
-                "Create a markdown table with 3 rows for the primary colors, and 2 columns, with color name and hex codes.",
-                ''] + params_list,
-            ['Translate to German:  My name is Arthur', ''] + params_list,
-            ["Please answer to the following question. Who is going to be the next Ballon d'or?", ''] + params_list,
-            ['Can Geoffrey Hinton have a conversation with George Washington? Give the rationale before answering.',
-             ''] + params_list,
-            ['Please answer the following question. What is the boiling point of Nitrogen?', ''] + params_list,
-            ['Answer the following yes/no question. Can you write a whole Haiku in a single tweet?', ''] + params_list,
-            ["Simplify the following expression: (False or False and True). Explain your answer.", ''] + params_list,
-            [
-                "Premise: At my age you will probably have learnt one lesson. Hypothesis:  It's not certain how many lessons you'll learn by your thirties. Does the premise entail the hypothesis?",
-                ''] + params_list,
-            ['The square root of x is the cube root of y. What is y to the power of 2, if x = 4?', ''] + params_list,
-            [
-                'Answer the following question by reasoning step by step.  The cafeteria had 23 apples. If they used 20 for lunch, and bought 6 more, how many apple do they have?',
-                ''] + params_list,
-            ["""def area_of_rectangle(a: float, b: float):
-    \"\"\"Return the area of the rectangle.\"\"\"""", ''] + params_list,
-            ["""# a function in native python:
-def mean(a):
-    return sum(a)/len(a)
-
-# the same function using numpy:
-import numpy as np
-def mean(a):""", ''] + params_list,
-            ["""X = np.random.randn(100, 100)
-y = np.random.randint(0, 1, 100)
-
-# fit random forest classifier with 20 estimators""", ''] + params_list,
-        ]
-    # add summary example
-    examples += [[summarize_example1, 'Summarize' if prompt_type not in ['plain', 'instruct_simple'] else ''] + params_list]
-
-    src_lang = "English"
-    tgt_lang = "Russian"
-
-    # move to correct position
-    for example in examples:
-        example += [chat, '', '', 'Disabled', k, ['All']]
-        # adjust examples if non-chat mode
-        if not chat:
-            example[eval_func_param_names.index('instruction_nochat')] = example[
-                eval_func_param_names.index('instruction')]
-            example[eval_func_param_names.index('instruction')] = ''
-
-            example[eval_func_param_names.index('iinput_nochat')] = example[eval_func_param_names.index('iinput')]
-            example[eval_func_param_names.index('iinput')] = ''
-        assert len(example) == len(eval_func_param_names), "Wrong example: %s %s" % (
-            len(example), len(eval_func_param_names))
-
-    return placeholder_instruction, placeholder_input, \
-        stream_output, show_examples, \
-        prompt_type, temperature, top_p, top_k, num_beams, \
-        max_new_tokens, min_new_tokens, early_stopping, max_time, \
-        repetition_penalty, num_return_sequences, \
-        do_sample, \
-        src_lang, tgt_lang, \
-        examples, \
-        task_info
-
-
-def languages_covered():
-    # https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt#languages-covered
-    covered = """Arabic (ar_AR), Czech (cs_CZ), German (de_DE), English (en_XX), Spanish (es_XX), Estonian (et_EE), Finnish (fi_FI), French (fr_XX), Gujarati (gu_IN), Hindi (hi_IN), Italian (it_IT), Japanese (ja_XX), Kazakh (kk_KZ), Korean (ko_KR), Lithuanian (lt_LT), Latvian (lv_LV), Burmese (my_MM), Nepali (ne_NP), Dutch (nl_XX), Romanian (ro_RO), Russian (ru_RU), Sinhala (si_LK), Turkish (tr_TR), Vietnamese (vi_VN), Chinese (zh_CN), Afrikaans (af_ZA), Azerbaijani (az_AZ), Bengali (bn_IN), Persian (fa_IR), Hebrew (he_IL), Croatian (hr_HR), Indonesian (id_ID), Georgian (ka_GE), Khmer (km_KH), Macedonian (mk_MK), Malayalam (ml_IN), Mongolian (mn_MN), Marathi (mr_IN), Polish (pl_PL), Pashto (ps_AF), Portuguese (pt_XX), Swedish (sv_SE), Swahili (sw_KE), Tamil (ta_IN), Telugu (te_IN), Thai (th_TH), Tagalog (tl_XX), Ukrainian (uk_UA), Urdu (ur_PK), Xhosa (xh_ZA), Galician (gl_ES), Slovene (sl_SI)"""
-    covered = covered.split(', ')
-    covered = {x.split(' ')[0]: x.split(' ')[1].replace(')', '').replace('(', '') for x in covered}
-    return covered
-
-
-def get_context(chat_context, prompt_type):
-    if chat_context and prompt_type == 'human_bot':
-        context0 = """<bot>: I am an intelligent, helpful, truthful, and fair assistant named h2oGPT, who will give accurate, balanced, and reliable responses.  I will not respond with I don't know or I don't understand.
-<human>: I am a human person seeking useful assistance and request all questions be answered completely, and typically expect detailed responses.  Give answers in numbered list format if several distinct but related items are being listed."""
-    else:
-        context0 = ''
-    return context0
-
-
-def score_qa(smodel, stokenizer, max_length_tokenize, question, answer, cutoff_len):
-    question = question[-cutoff_len:]
-    answer = answer[-cutoff_len:]
-
-    inputs = stokenizer(question, answer,
-                        return_tensors="pt",
-                        truncation=True,
-                        max_length=max_length_tokenize).to(smodel.device)
-    try:
-        score = torch.sigmoid(smodel(**inputs).logits[0]).cpu().detach().numpy()[0]
-    except torch.cuda.OutOfMemoryError as e:
-        print("GPU OOM 3: question: %s answer: %s exception: %s" % (question, answer, str(e)), flush=True)
-        del inputs
-        traceback.print_exc()
-        clear_torch_cache()
-        return 'Response Score: GPU OOM'
-    except (Exception, RuntimeError) as e:
-        if 'Expected all tensors to be on the same device' in str(e) or \
-                'expected scalar type Half but found Float' in str(e) or \
-                'probability tensor contains either' in str(e) or \
-                'cublasLt ran into an error!' in str(e):
-            print("GPU Error: question: %s answer: %s exception: %s" % (question, answer, str(e)),
-                  flush=True)
-            traceback.print_exc()
-            clear_torch_cache()
-            return 'Response Score: GPU Error'
-        else:
-            raise
-    os.environ['TOKENIZERS_PARALLELISM'] = 'true'
-    return score
-
-
-def check_locals(**kwargs):
-    # ensure everything in evaluate is here
-    can_skip_because_locally_generated = [  # evaluate
-        'instruction',
-        'iinput',
-        'context',
-        'instruction_nochat',
-        'iinput_nochat',
-        # get_model:
-        'reward_type'
-    ]
-    for k in eval_func_param_names:
-        if k in can_skip_because_locally_generated:
-            continue
-        assert k in kwargs, "Missing %s" % k
-    for k in inputs_kwargs_list:
-        if k in can_skip_because_locally_generated:
-            continue
-        assert k in kwargs, "Missing %s" % k
-
-    for k in list(inspect.signature(get_model).parameters):
-        if k in can_skip_because_locally_generated:
-            continue
-        assert k in kwargs, "Missing %s" % k
-
-
-if __name__ == "__main__":
-    """
-    Examples:
-
-    WORLD_SIZE=4 CUDA_VISIBLE_DEVICES="0,1,2,3" torchrun --nproc_per_node=4 --master_port=1234 generate.py --base_model='EleutherAI/gpt-j-6B' --lora_weights=lora-alpaca_6B
-    python generate.py --base_model='EleutherAI/gpt-j-6B' --lora_weights='lora-alpaca_6B'
-    python generate.py --base_model='EleutherAI/gpt-neox-20b' --lora_weights='lora-alpaca_20B'
-    
-    # generate without lora weights, no prompt
-    python generate.py --base_model='EleutherAI/gpt-neox-20b' --prompt_type='plain'
-    python generate.py --base_model='togethercomputer/GPT-NeoXT-Chat-Base-20B' --prompt_type='dai_faq'
-
-    python generate.py --base_model='togethercomputer/GPT-NeoXT-Chat-Base-20B' --prompt_type='dai_faq' --lora_weights='lora_20B_daifaq'
-    # OpenChatKit settings:
-    python generate.py --base_model='togethercomputer/GPT-NeoXT-Chat-Base-20B' --prompt_type='human_bot --debug=True --num_beams=1 --temperature=0.6 --top_k=40 --top_p=1.0
-
-    python generate.py --base_model='distilgpt2' --prompt_type='plain' --debug=True --num_beams=1 --temperature=0.6 --top_k=40 --top_p=1.0 --share=False
-    python generate.py --base_model='t5-large' --prompt_type='simple_instruct'
-    python generate.py --base_model='philschmid/bart-large-cnn-samsum'
-    python generate.py --base_model='philschmid/flan-t5-base-samsum'
-    python generate.py --base_model='facebook/mbart-large-50-many-to-many-mmt'
-
-    python generate.py --base_model='togethercomputer/GPT-NeoXT-Chat-Base-20B' --prompt_type='human_bot' --lora_weights='GPT-NeoXT-Chat-Base-20B.merged.json.8_epochs.57b2892c53df5b8cefac45f84d019cace803ef26.28'
-
-    must have 4*48GB GPU and run without 8bit in order for sharding to work with infer_devices=False
-    can also pass --prompt_type='human_bot' and model can somewhat handle instructions without being instruct tuned
-    python generate.py --base_model=decapoda-research/llama-65b-hf --load_8bit=False --infer_devices=False --prompt_type='human_bot'
-
-    python generate.py --base_model=h2oai/h2ogpt-oig-oasst1-512-6_9b
-    """
-    fire.Fire(main)

generate.py

@@ -0,0 +1 @@
+../../generate.py

gpt4all_llm.py

@@ -1,255 +0,0 @@
-import inspect
-import os
-import sys
-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
-
-
-class FakeTokenizer:
-
-    def encode(self, x, *args, **kwargs):
-        return dict(input_ids=[x])
-
-    def decode(self, x, *args, **kwargs):
-        return x
-
-    def __call__(self, x, *args, **kwargs):
-        return self.encode(x, *args, **kwargs)
-
-
-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))
-
-    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):
-    # default from class
-    model_kwargs = {k: v.default for k, v in dict(inspect.signature(cls).parameters).items()}
-    # 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,
-                    verbose=False):
-    env_gpt4all_file = ".env_gpt4all"
-    env_kwargs = dotenv_values(env_gpt4all_file)
-    callbacks = [H2OStreamingStdOutCallbackHandler()]
-    n_ctx = env_kwargs.pop('n_ctx', 2048 - max_new_tokens)
-    default_kwargs = dict(context_erase=0.5,
-                          n_batch=1,
-                          n_ctx=n_ctx,
-                          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)
-        model_kwargs.update(dict(model_path=model_path, callbacks=callbacks))
-        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)
-        model_kwargs.update(dict(model=model_path, backend='llama', callbacks=callbacks))
-        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)
-        model_kwargs.update(dict(model=model_path, backend='gptj', callbacks=callbacks))
-        llm = cls(**model_kwargs)
-    else:
-        raise RuntimeError("No such model_name %s" % model_name)
-    return llm
-
-
-class H2OGPT4All(gpt4all.GPT4All):
-    model: 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,
-                )
-            else:
-                values["client"] = values["model"]
-            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,
-    ) -> str:
-        # Roughly 4 chars per token if natural language
-        prompt = prompt[-self.n_ctx * 4:]
-        verbose = False
-        if verbose:
-            print("_call prompt: %s" % prompt, flush=True)
-        return super()._call(prompt, stop=stop, run_manager=run_manager)
-
-
-from langchain.llms import LlamaCpp
-
-
-class H2OLlamaCpp(LlamaCpp):
-    model_path: 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,
-    ) -> str:
-        verbose = False
-        # tokenize twice, just to count tokens, since llama cpp python wrapper has no way to truncate
-        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)
-        if verbose:
-            print("_call prompt: %s" % prompt, flush=True)
-        return super()._call(prompt, stop=stop, run_manager=run_manager)

gpt4all_llm.py

@@ -0,0 +1 @@
+../../gpt4all_llm.py

gpt_langchain.py

@@ -1,1471 +0,0 @@
-import glob
-import inspect
-import os
-import pathlib
-import pickle
-import queue
-import shutil
-import subprocess
-import sys
-import tempfile
-import traceback
-import uuid
-import zipfile
-from collections import defaultdict
-from datetime import datetime
-from functools import reduce
-from operator import concat
-
-from joblib import Parallel, delayed
-from tqdm import tqdm
-
-from prompter import non_hf_types
-from utils import wrapped_partial, EThread, import_matplotlib, sanitize_filename, makedirs, get_url, flatten_list, \
-    get_device, ProgressParallel, remove, hash_file
-
-import_matplotlib()
-
-import numpy as np
-import pandas as pd
-import requests
-from langchain.chains.qa_with_sources import load_qa_with_sources_chain
-# , GCSDirectoryLoader, GCSFileLoader
-# , OutlookMessageLoader # GPL3
-# ImageCaptionLoader, # use our own wrapper
-#  ReadTheDocsLoader,  # no special file, some path, so have to give as special option
-from langchain.document_loaders import PyPDFLoader, TextLoader, CSVLoader, PythonLoader, TomlLoader, \
-    UnstructuredURLLoader, UnstructuredHTMLLoader, UnstructuredWordDocumentLoader, UnstructuredMarkdownLoader, \
-    EverNoteLoader, UnstructuredEmailLoader, UnstructuredODTLoader, UnstructuredPowerPointLoader, \
-    UnstructuredEPubLoader, UnstructuredImageLoader, UnstructuredRTFLoader, ArxivLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.chains.question_answering import load_qa_chain
-from langchain.docstore.document import Document
-from langchain import PromptTemplate
-from langchain.vectorstores import Chroma
-
-
-def get_db(sources, use_openai_embedding=False, db_type='faiss', persist_directory="db_dir", langchain_mode='notset',
-           collection_name=None,
-           hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2"):
-    if not sources:
-        return None
-    # get embedding model
-    embedding = get_embedding(use_openai_embedding, hf_embedding_model=hf_embedding_model)
-    assert collection_name is not None or langchain_mode != 'notset'
-    if collection_name is None:
-        collection_name = langchain_mode.replace(' ', '_')
-
-    # Create vector database
-    if db_type == 'faiss':
-        from langchain.vectorstores import FAISS
-        db = FAISS.from_documents(sources, embedding)
-
-    elif db_type == 'weaviate':
-        import weaviate
-        from weaviate.embedded import EmbeddedOptions
-        from langchain.vectorstores import Weaviate
-
-        # TODO: add support for connecting via docker compose
-        client = weaviate.Client(
-            embedded_options=EmbeddedOptions()
-        )
-        index_name = collection_name.capitalize()
-        db = Weaviate.from_documents(documents=sources, embedding=embedding, client=client, by_text=False,
-                                     index_name=index_name)
-
-    elif db_type == 'chroma':
-        assert persist_directory is not None
-        os.makedirs(persist_directory, exist_ok=True)
-        db = Chroma.from_documents(documents=sources,
-                                   embedding=embedding,
-                                   persist_directory=persist_directory,
-                                   collection_name=collection_name,
-                                   anonymized_telemetry=False)
-        db.persist()
-    else:
-        raise RuntimeError("No such db_type=%s" % db_type)
-
-    return db
-
-
-def _get_unique_sources_in_weaviate(db):
-    batch_size = 100
-    id_source_list = []
-    result = db._client.data_object.get(class_name=db._index_name, limit=batch_size)
-
-    while result['objects']:
-        id_source_list += [(obj['id'], obj['properties']['source']) for obj in result['objects']]
-        last_id = id_source_list[-1][0]
-        result = db._client.data_object.get(class_name=db._index_name, limit=batch_size, after=last_id)
-
-    unique_sources = {source for _, source in id_source_list}
-    return unique_sources
-
-
-def add_to_db(db, sources, db_type='faiss',
-              avoid_dup_by_file=False,
-              avoid_dup_by_content=True):
-    num_new_sources = len(sources)
-    if not sources:
-        return db, num_new_sources, []
-    if db_type == 'faiss':
-        db.add_documents(sources)
-    elif db_type == 'weaviate':
-        # FIXME: only control by file name, not hash yet
-        if avoid_dup_by_file or avoid_dup_by_content:
-            unique_sources = _get_unique_sources_in_weaviate(db)
-            sources = [x for x in sources if x.metadata['source'] not in unique_sources]
-        num_new_sources = len(sources)
-        if num_new_sources == 0:
-            return db, num_new_sources, []
-        db.add_documents(documents=sources)
-    elif db_type == 'chroma':
-        collection = db.get()
-        # files we already have:
-        metadata_files = set([x['source'] for x in collection['metadatas']])
-        if avoid_dup_by_file:
-            # Too weak in case file changed content, assume parent shouldn't pass true for this for now
-            raise RuntimeError("Not desired code path")
-            sources = [x for x in sources if x.metadata['source'] not in metadata_files]
-        if avoid_dup_by_content:
-            # look at hash, instead of page_content
-            # migration: If no hash previously, avoid updating,
-            #  since don't know if need to update and may be expensive to redo all unhashed files
-            metadata_hash_ids = set(
-                [x['hashid'] for x in collection['metadatas'] if 'hashid' in x and x['hashid'] not in ["None", None]])
-            # avoid sources with same hash
-            sources = [x for x in sources if x.metadata.get('hashid') not in metadata_hash_ids]
-            # get new file names that match existing file names.  delete existing files we are overridding
-            dup_metadata_files = set([x.metadata['source'] for x in sources if x.metadata['source'] in metadata_files])
-            print("Removing %s duplicate files from db because ingesting those as new documents" % len(
-                dup_metadata_files), flush=True)
-            client_collection = db._client.get_collection(name=db._collection.name)
-            for dup_file in dup_metadata_files:
-                dup_file_meta = dict(source=dup_file)
-                try:
-                    client_collection.delete(where=dup_file_meta)
-                except KeyError:
-                    pass
-        num_new_sources = len(sources)
-        if num_new_sources == 0:
-            return db, num_new_sources, []
-        db.add_documents(documents=sources)
-        db.persist()
-    else:
-        raise RuntimeError("No such db_type=%s" % db_type)
-
-    new_sources_metadata = [x.metadata for x in sources]
-
-    return db, num_new_sources, new_sources_metadata
-
-
-def create_or_update_db(db_type, persist_directory, collection_name,
-                        sources, use_openai_embedding, add_if_exists, verbose, hf_embedding_model):
-    if db_type == 'weaviate':
-        import weaviate
-        from weaviate.embedded import EmbeddedOptions
-
-        # TODO: add support for connecting via docker compose
-        client = weaviate.Client(
-            embedded_options=EmbeddedOptions()
-        )
-        index_name = collection_name.replace(' ', '_').capitalize()
-        if client.schema.exists(index_name) and not add_if_exists:
-            client.schema.delete_class(index_name)
-            if verbose:
-                print("Removing %s" % index_name, flush=True)
-    elif db_type == 'chroma':
-        if not os.path.isdir(persist_directory) or not add_if_exists:
-            if os.path.isdir(persist_directory):
-                if verbose:
-                    print("Removing %s" % persist_directory, flush=True)
-                remove(persist_directory)
-            if verbose:
-                print("Generating db", flush=True)
-
-    if not add_if_exists:
-        if verbose:
-            print("Generating db", flush=True)
-    else:
-        if verbose:
-            print("Loading and updating db", flush=True)
-
-    db = get_db(sources,
-                use_openai_embedding=use_openai_embedding,
-                db_type=db_type,
-                persist_directory=persist_directory,
-                langchain_mode=collection_name,
-                hf_embedding_model=hf_embedding_model)
-
-    return db
-
-
-def get_embedding(use_openai_embedding, hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2"):
-    # Get embedding model
-    if use_openai_embedding:
-        assert os.getenv("OPENAI_API_KEY") is not None, "Set ENV OPENAI_API_KEY"
-        from langchain.embeddings import OpenAIEmbeddings
-        embedding = OpenAIEmbeddings()
-    else:
-        # to ensure can fork without deadlock
-        from langchain.embeddings import HuggingFaceEmbeddings
-
-        device, torch_dtype, context_class = get_device_dtype()
-        model_kwargs = dict(device=device)
-        embedding = HuggingFaceEmbeddings(model_name=hf_embedding_model, model_kwargs=model_kwargs)
-    return embedding
-
-
-def get_answer_from_sources(chain, sources, question):
-    return chain(
-        {
-            "input_documents": sources,
-            "question": question,
-        },
-        return_only_outputs=True,
-    )["output_text"]
-
-
-def get_llm(use_openai_model=False, model_name=None, model=None,
-            tokenizer=None, stream_output=False,
-            max_new_tokens=256,
-            temperature=0.1,
-            repetition_penalty=1.0,
-            top_k=40,
-            top_p=0.7,
-            prompt_type=None,
-            prompter=None,
-            verbose=False,
-            ):
-    if use_openai_model:
-        from langchain.llms import OpenAI
-        llm = OpenAI(temperature=0)
-        model_name = 'openai'
-        streamer = None
-        prompt_type = 'plain'
-    elif model_name in non_hf_types:
-        from gpt4all_llm import get_llm_gpt4all
-        llm = get_llm_gpt4all(model_name, model=model, max_new_tokens=max_new_tokens,
-                              temperature=temperature,
-                              repetition_penalty=repetition_penalty,
-                              top_k=top_k,
-                              top_p=top_p,
-                              verbose=verbose,
-                              )
-        streamer = None
-        prompt_type = 'plain'
-    else:
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-
-        if model is None:
-            # only used if didn't pass model in
-            assert model_name is None
-            assert tokenizer is None
-            prompt_type = 'human_bot'
-            model_name = 'h2oai/h2ogpt-oasst1-512-12b'
-            # model_name = 'h2oai/h2ogpt-oig-oasst1-512-6_9b'
-            # model_name = 'h2oai/h2ogpt-oasst1-512-20b'
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            device, torch_dtype, context_class = get_device_dtype()
-
-            with context_class(device):
-                load_8bit = True
-                # FIXME: for now not to spread across hetero GPUs
-                # device_map={"": 0} if load_8bit and device == 'cuda' else "auto"
-                device_map = {"": 0} if device == 'cuda' else "auto"
-                model = AutoModelForCausalLM.from_pretrained(model_name,
-                                                             device_map=device_map,
-                                                             torch_dtype=torch_dtype,
-                                                             load_in_8bit=load_8bit)
-
-        max_max_tokens = tokenizer.model_max_length
-        gen_kwargs = dict(max_new_tokens=max_new_tokens,
-                          return_full_text=True,
-                          early_stopping=False,
-                          handle_long_generation='hole')
-
-        if stream_output:
-            skip_prompt = False
-            from generate import H2OTextIteratorStreamer
-            decoder_kwargs = {}
-            streamer = H2OTextIteratorStreamer(tokenizer, skip_prompt=skip_prompt, block=False, **decoder_kwargs)
-            gen_kwargs.update(dict(streamer=streamer))
-        else:
-            streamer = None
-
-        from h2oai_pipeline import H2OTextGenerationPipeline
-        pipe = H2OTextGenerationPipeline(model=model, use_prompter=True,
-                                         prompter=prompter,
-                                         prompt_type=prompt_type,
-                                         sanitize_bot_response=True,
-                                         chat=False, stream_output=stream_output,
-                                         tokenizer=tokenizer,
-                                         max_input_tokens=max_max_tokens - max_new_tokens,
-                                         **gen_kwargs)
-        # pipe.task = "text-generation"
-        # below makes it listen only to our prompt removal,
-        # not built in prompt removal that is less general and not specific for our model
-        pipe.task = "text2text-generation"
-
-        from langchain.llms import HuggingFacePipeline
-        llm = HuggingFacePipeline(pipeline=pipe)
-    return llm, model_name, streamer, prompt_type
-
-
-def get_device_dtype():
-    # torch.device("cuda") leads to cuda:x cuda:y mismatches for multi-GPU consistently
-    import torch
-    n_gpus = torch.cuda.device_count() if torch.cuda.is_available else 0
-    device = 'cpu' if n_gpus == 0 else 'cuda'
-    # from utils import NullContext
-    # context_class = NullContext if n_gpus > 1 or n_gpus == 0 else context_class
-    context_class = torch.device
-    torch_dtype = torch.float16 if device == 'cuda' else torch.float32
-    return device, torch_dtype, context_class
-
-
-def get_wiki_data(title, first_paragraph_only, text_limit=None, take_head=True):
-    """
-    Get wikipedia data from online
-    :param title:
-    :param first_paragraph_only:
-    :param text_limit:
-    :param take_head:
-    :return:
-    """
-    filename = 'wiki_%s_%s_%s_%s.data' % (first_paragraph_only, title, text_limit, take_head)
-    url = f"https://en.wikipedia.org/w/api.php?format=json&action=query&prop=extracts&explaintext=1&titles={title}"
-    if first_paragraph_only:
-        url += "&exintro=1"
-    import json
-    if not os.path.isfile(filename):
-        data = requests.get(url).json()
-        json.dump(data, open(filename, 'wt'))
-    else:
-        data = json.load(open(filename, "rt"))
-    page_content = list(data["query"]["pages"].values())[0]["extract"]
-    if take_head is not None and text_limit is not None:
-        page_content = page_content[:text_limit] if take_head else page_content[:-text_limit]
-    title_url = str(title).replace(' ', '_')
-    return Document(
-        page_content=page_content,
-        metadata={"source": f"https://en.wikipedia.org/wiki/{title_url}"},
-    )
-
-
-def get_wiki_sources(first_para=True, text_limit=None):
-    """
-    Get specific named sources from wikipedia
-    :param first_para:
-    :param text_limit:
-    :return:
-    """
-    default_wiki_sources = ['Unix', 'Microsoft_Windows', 'Linux']
-    wiki_sources = list(os.getenv('WIKI_SOURCES', default_wiki_sources))
-    return [get_wiki_data(x, first_para, text_limit=text_limit) for x in wiki_sources]
-
-
-def get_github_docs(repo_owner, repo_name):
-    """
-    Access github from specific repo
-    :param repo_owner:
-    :param repo_name:
-    :return:
-    """
-    with tempfile.TemporaryDirectory() as d:
-        subprocess.check_call(
-            f"git clone --depth 1 https://github.com/{repo_owner}/{repo_name}.git .",
-            cwd=d,
-            shell=True,
-        )
-        git_sha = (
-            subprocess.check_output("git rev-parse HEAD", shell=True, cwd=d)
-            .decode("utf-8")
-            .strip()
-        )
-        repo_path = pathlib.Path(d)
-        markdown_files = list(repo_path.glob("*/*.md")) + list(
-            repo_path.glob("*/*.mdx")
-        )
-        for markdown_file in markdown_files:
-            with open(markdown_file, "r") as f:
-                relative_path = markdown_file.relative_to(repo_path)
-                github_url = f"https://github.com/{repo_owner}/{repo_name}/blob/{git_sha}/{relative_path}"
-                yield Document(page_content=f.read(), metadata={"source": github_url})
-
-
-def get_dai_pickle(dest="."):
-    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.pickle', token=token, repo_type='dataset')
-    shutil.copy(path_to_zip_file, dest)
-
-
-def get_dai_docs(from_hf=False, get_pickle=True):
-    """
-    Consume DAI documentation, or consume from public pickle
-    :param from_hf: get DAI docs from HF, then generate pickle for later use by LangChain
-    :param get_pickle: Avoid raw DAI docs, just get pickle directly from HF
-    :return:
-    """
-    import pickle
-
-    if get_pickle:
-        get_dai_pickle()
-
-    dai_store = 'dai_docs.pickle'
-    dst = "working_dir_docs"
-    if not os.path.isfile(dai_store):
-        from create_data import setup_dai_docs
-        dst = setup_dai_docs(dst=dst, from_hf=from_hf)
-
-        import glob
-        files = list(glob.glob(os.path.join(dst, '*rst'), recursive=True))
-
-        basedir = os.path.abspath(os.getcwd())
-        from create_data import rst_to_outputs
-        new_outputs = rst_to_outputs(files)
-        os.chdir(basedir)
-
-        pickle.dump(new_outputs, open(dai_store, 'wb'))
-    else:
-        new_outputs = pickle.load(open(dai_store, 'rb'))
-
-    sources = []
-    for line, file in new_outputs:
-        # gradio requires any linked file to be with app.py
-        sym_src = os.path.abspath(os.path.join(dst, file))
-        sym_dst = os.path.abspath(os.path.join(os.getcwd(), file))
-        if os.path.lexists(sym_dst):
-            os.remove(sym_dst)
-        os.symlink(sym_src, sym_dst)
-        itm = Document(page_content=line, metadata={"source": file})
-        # NOTE: yield has issues when going into db, loses metadata
-        # yield itm
-        sources.append(itm)
-    return sources
-
-
-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
-
-image_types = ["png", "jpg", "jpeg"]
-non_image_types = ["pdf", "txt", "csv", "toml", "py", "rst", "rtf",
-                   "md", "html",
-                   "enex", "eml", "epub", "odt", "pptx", "ppt",
-                   "zip", "urls",
-                   ]
-# "msg",  GPL3
-
-if have_libreoffice:
-    non_image_types.extend(["docx", "doc"])
-
-file_types = non_image_types + image_types
-
-
-def add_meta(docs1, file):
-    file_extension = pathlib.Path(file).suffix
-    hashid = hash_file(file)
-    if not isinstance(docs1, list):
-        docs1 = [docs1]
-    [x.metadata.update(dict(input_type=file_extension, date=str(datetime.now), hashid=hashid)) for x in docs1]
-
-
-def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, chunk=True, chunk_size=512,
-                is_url=False, is_txt=False,
-                enable_captions=True,
-                captions_model=None,
-                enable_ocr=False, caption_loader=None,
-                headsize=50):
-    if file is None:
-        if fail_any_exception:
-            raise RuntimeError("Unexpected None file")
-        else:
-            return []
-    doc1 = []  # in case no support, or disabled support
-    if base_path is None and not is_txt and not is_url:
-        # then assume want to persist but don't care which path used
-        # can't be in base_path
-        dir_name = os.path.dirname(file)
-        base_name = os.path.basename(file)
-        # if from gradio, will have its own temp uuid too, but that's ok
-        base_name = sanitize_filename(base_name) + "_" + str(uuid.uuid4())[:10]
-        base_path = os.path.join(dir_name, base_name)
-    if is_url:
-        if file.lower().startswith('arxiv:'):
-            query = file.lower().split('arxiv:')
-            if len(query) == 2 and have_arxiv:
-                query = query[1]
-                docs1 = ArxivLoader(query=query, load_max_docs=20, load_all_available_meta=True).load()
-                # ensure string, sometimes None
-                [[x.metadata.update({k: str(v)}) for k, v in x.metadata.items()] for x in docs1]
-                query_url = f"https://arxiv.org/abs/{query}"
-                [x.metadata.update(
-                    dict(source=x.metadata.get('entry_id', query_url), query=query_url,
-                         input_type='arxiv', head=x.metadata.get('Title', ''), date=str(datetime.now))) for x in
-                    docs1]
-            else:
-                docs1 = []
-        else:
-            docs1 = UnstructuredURLLoader(urls=[file]).load()
-            [x.metadata.update(dict(input_type='url', date=str(datetime.now))) for x in docs1]
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif is_txt:
-        base_path = "user_paste"
-        source_file = os.path.join(base_path, "_%s" % str(uuid.uuid4())[:10])
-        makedirs(os.path.dirname(source_file), exist_ok=True)
-        with open(source_file, "wt") as f:
-            f.write(file)
-        metadata = dict(source=source_file, date=str(datetime.now()), input_type='pasted txt')
-        doc1 = Document(page_content=file, metadata=metadata)
-    elif file.lower().endswith('.html') or file.lower().endswith('.mhtml'):
-        docs1 = UnstructuredHTMLLoader(file_path=file).load()
-        add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif (file.lower().endswith('.docx') or file.lower().endswith('.doc')) and have_libreoffice:
-        docs1 = UnstructuredWordDocumentLoader(file_path=file).load()
-        add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif file.lower().endswith('.odt'):
-        docs1 = UnstructuredODTLoader(file_path=file).load()
-        add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif file.lower().endswith('pptx') or file.lower().endswith('ppt'):
-        docs1 = UnstructuredPowerPointLoader(file_path=file).load()
-        add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif file.lower().endswith('.txt'):
-        # use UnstructuredFileLoader ?
-        docs1 = TextLoader(file, encoding="utf8", autodetect_encoding=True).load()
-        # makes just one, but big one
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-        add_meta(doc1, file)
-    elif file.lower().endswith('.rtf'):
-        docs1 = UnstructuredRTFLoader(file).load()
-        add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif file.lower().endswith('.md'):
-        docs1 = UnstructuredMarkdownLoader(file).load()
-        add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif file.lower().endswith('.enex'):
-        docs1 = EverNoteLoader(file).load()
-        add_meta(doc1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif file.lower().endswith('.epub'):
-        docs1 = UnstructuredEPubLoader(file).load()
-        add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif file.lower().endswith('.jpeg') or file.lower().endswith('.jpg') or file.lower().endswith('.png'):
-        docs1 = []
-        if have_tesseract and enable_ocr:
-            # OCR, somewhat works, but not great
-            docs1.extend(UnstructuredImageLoader(file).load())
-            add_meta(docs1, file)
-        if enable_captions:
-            # BLIP
-            if caption_loader is not None and not isinstance(caption_loader, (str, bool)):
-                # assumes didn't fork into this process with joblib, else can deadlock
-                caption_loader.set_image_paths([file])
-                docs1c = caption_loader.load()
-                add_meta(docs1c, file)
-                [x.metadata.update(dict(head=x.page_content[:headsize].strip())) for x in docs1c]
-                docs1.extend(docs1c)
-            else:
-                from image_captions import H2OImageCaptionLoader
-                caption_loader = H2OImageCaptionLoader(caption_gpu=caption_loader == 'gpu',
-                                                       blip_model=captions_model,
-                                                       blip_processor=captions_model)
-                caption_loader.set_image_paths([file])
-                docs1c = caption_loader.load()
-                add_meta(docs1c, file)
-                [x.metadata.update(dict(head=x.page_content[:headsize].strip())) for x in docs1c]
-                docs1.extend(docs1c)
-            for doci in docs1:
-                doci.metadata['source'] = doci.metadata['image_path']
-                doci.metadata['hash'] = hash_file(doci.metadata['source'])
-            if docs1:
-                doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif file.lower().endswith('.msg'):
-        raise RuntimeError("Not supported, GPL3 license")
-        # docs1 = OutlookMessageLoader(file).load()
-        # docs1[0].metadata['source'] = file
-    elif file.lower().endswith('.eml'):
-        try:
-            docs1 = UnstructuredEmailLoader(file).load()
-            add_meta(docs1, file)
-            doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-        except ValueError as e:
-            if 'text/html content not found in email' in str(e):
-                # e.g. plain/text dict key exists, but not
-                # doc1 = TextLoader(file, encoding="utf8").load()
-                docs1 = UnstructuredEmailLoader(file, content_source="text/plain").load()
-                add_meta(docs1, file)
-                doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-            else:
-                raise
-    # elif file.lower().endswith('.gcsdir'):
-    #    doc1 = GCSDirectoryLoader(project_name, bucket, prefix).load()
-    # elif file.lower().endswith('.gcsfile'):
-    # doc1 = GCSFileLoader(project_name, bucket, blob).load()
-    elif file.lower().endswith('.rst'):
-        with open(file, "r") as f:
-            doc1 = Document(page_content=f.read(), metadata={"source": file})
-        add_meta(doc1, file)
-    elif file.lower().endswith('.pdf'):
-        env_gpt4all_file = ".env_gpt4all"
-        from dotenv import dotenv_values
-        env_kwargs = dotenv_values(env_gpt4all_file)
-        pdf_class_name = env_kwargs.get('PDF_CLASS_NAME', 'PyMuPDFParser')
-        if have_pymupdf and pdf_class_name == 'PyMuPDFParser':
-            # GPL, only use if installed
-            from langchain.document_loaders import PyMuPDFLoader
-            doc1 = PyMuPDFLoader(file).load_and_split()
-        else:
-            # open-source fallback
-            doc1 = PyPDFLoader(file).load_and_split()
-        # Some PDFs return nothing or junk from PDFMinerLoader
-        add_meta(doc1, file)
-    elif file.lower().endswith('.csv'):
-        doc1 = CSVLoader(file).load()
-        add_meta(doc1, file)
-    elif file.lower().endswith('.py'):
-        doc1 = PythonLoader(file).load()
-        add_meta(doc1, file)
-    elif file.lower().endswith('.toml'):
-        doc1 = TomlLoader(file).load()
-        add_meta(doc1, file)
-    elif file.lower().endswith('.urls'):
-        with open(file, "r") as f:
-            docs1 = UnstructuredURLLoader(urls=f.readlines()).load()
-        add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
-    elif file.lower().endswith('.zip'):
-        with zipfile.ZipFile(file, 'r') as zip_ref:
-            # don't put into temporary path, since want to keep references to docs inside zip
-            # so just extract in path where
-            zip_ref.extractall(base_path)
-            # recurse
-            doc1 = path_to_docs(base_path, verbose=verbose, fail_any_exception=fail_any_exception)
-    else:
-        raise RuntimeError("No file handler for %s" % os.path.basename(file))
-
-    # allow doc1 to be list or not.  If not list, did not chunk yet, so chunk now
-    # if list of length one, don't trust and chunk it
-    if not isinstance(doc1, list):
-        if chunk:
-            docs = chunk_sources([doc1], chunk_size=chunk_size)
-        else:
-            docs = [doc1]
-    elif isinstance(doc1, list) and len(doc1) == 1:
-        if chunk:
-            docs = chunk_sources(doc1, chunk_size=chunk_size)
-        else:
-            docs = doc1
-    else:
-        docs = doc1
-
-    assert isinstance(docs, list)
-    return docs
-
-
-def path_to_doc1(file, verbose=False, fail_any_exception=False, return_file=True, chunk=True, chunk_size=512,
-                 is_url=False, is_txt=False,
-                 enable_captions=True,
-                 captions_model=None,
-                 enable_ocr=False, caption_loader=None):
-    if verbose:
-        if is_url:
-            print("Ingesting URL: %s" % file, flush=True)
-        elif is_txt:
-            print("Ingesting Text: %s" % file, flush=True)
-        else:
-            print("Ingesting file: %s" % file, flush=True)
-    res = None
-    try:
-        # don't pass base_path=path, would infinitely recurse
-        res = file_to_doc(file, base_path=None, verbose=verbose, fail_any_exception=fail_any_exception,
-                          chunk=chunk, chunk_size=chunk_size,
-                          is_url=is_url, is_txt=is_txt,
-                          enable_captions=enable_captions,
-                          captions_model=captions_model,
-                          enable_ocr=enable_ocr,
-                          caption_loader=caption_loader)
-    except BaseException as e:
-        print("Failed to ingest %s due to %s" % (file, traceback.format_exc()))
-        if fail_any_exception:
-            raise
-        else:
-            exception_doc = Document(
-                page_content='',
-                metadata={"source": file, "exception": str(e), "traceback": traceback.format_exc()})
-            res = [exception_doc]
-    if return_file:
-        base_tmp = "temp_path_to_doc1"
-        if not os.path.isdir(base_tmp):
-            os.makedirs(base_tmp, exist_ok=True)
-        filename = os.path.join(base_tmp, str(uuid.uuid4()) + ".tmp.pickle")
-        with open(filename, 'wb') as f:
-            pickle.dump(res, f)
-        return filename
-    return res
-
-
-def path_to_docs(path_or_paths, verbose=False, fail_any_exception=False, n_jobs=-1,
-                 chunk=True, chunk_size=512,
-                 url=None, text=None,
-                 enable_captions=True,
-                 captions_model=None,
-                 caption_loader=None,
-                 enable_ocr=False,
-                 existing_files=[],
-                 existing_hash_ids={},
-                 ):
-    globs_image_types = []
-    globs_non_image_types = []
-    if not path_or_paths and not url and not text:
-        return []
-    elif url:
-        globs_non_image_types = [url]
-    elif text:
-        globs_non_image_types = [text]
-    elif isinstance(path_or_paths, str):
-        # single path, only consume allowed files
-        path = path_or_paths
-        # Below globs should match patterns in file_to_doc()
-        [globs_image_types.extend(glob.glob(os.path.join(path, "./**/*.%s" % ftype), recursive=True))
-         for ftype in image_types]
-        [globs_non_image_types.extend(glob.glob(os.path.join(path, "./**/*.%s" % ftype), recursive=True))
-         for ftype in non_image_types]
-    else:
-        # list/tuple of files (consume what can, and exception those that selected but cannot consume so user knows)
-        assert isinstance(path_or_paths, (list, tuple)), "Wrong type for path_or_paths: %s" % type(path_or_paths)
-        # reform out of allowed types
-        globs_image_types.extend(flatten_list([[x for x in path_or_paths if x.endswith(y)] for y in image_types]))
-        # could do below:
-        # globs_non_image_types = flatten_list([[x for x in path_or_paths if x.endswith(y)] for y in non_image_types])
-        # But instead, allow fail so can collect unsupported too
-        set_globs_image_types = set(globs_image_types)
-        globs_non_image_types.extend([x for x in path_or_paths if x not in set_globs_image_types])
-
-    # filter out any files to skip (e.g. if already processed them)
-    # this is easy, but too aggressive in case a file changed, so parent probably passed existing_files=[]
-    assert not existing_files, "DEV: assume not using this approach"
-    if existing_files:
-        set_skip_files = set(existing_files)
-        globs_image_types = [x for x in globs_image_types if x not in set_skip_files]
-        globs_non_image_types = [x for x in globs_non_image_types if x not in set_skip_files]
-    if existing_hash_ids:
-        # assume consistent with add_meta() use of hash_file(file)
-        # also assume consistent with get_existing_hash_ids for dict creation
-        # assume hashable values
-        existing_hash_ids_set = set(existing_hash_ids.items())
-        hash_ids_all_image = set({x: hash_file(x) for x in globs_image_types}.items())
-        hash_ids_all_non_image = set({x: hash_file(x) for x in globs_non_image_types}.items())
-        # don't use symmetric diff.  If file is gone, ignore and don't remove or something
-        #  just consider existing files (key) having new hash or not (value)
-        new_files_image = set(dict(hash_ids_all_image - existing_hash_ids_set).keys())
-        new_files_non_image = set(dict(hash_ids_all_non_image - existing_hash_ids_set).keys())
-        globs_image_types = [x for x in globs_image_types if x in new_files_image]
-        globs_non_image_types = [x for x in globs_non_image_types if x in new_files_non_image]
-
-    # could use generator, but messes up metadata handling in recursive case
-    if caption_loader and not isinstance(caption_loader, (bool, str)) and \
-            caption_loader.device != 'cpu' or \
-            get_device() == 'cuda':
-        # to avoid deadlocks, presume was preloaded and so can't fork due to cuda context
-        n_jobs_image = 1
-    else:
-        n_jobs_image = n_jobs
-
-    return_file = True  # local choice
-    is_url = url is not None
-    is_txt = text is not None
-    kwargs = dict(verbose=verbose, fail_any_exception=fail_any_exception,
-                  return_file=return_file,
-                  chunk=chunk, chunk_size=chunk_size,
-                  is_url=is_url,
-                  is_txt=is_txt,
-                  enable_captions=enable_captions,
-                  captions_model=captions_model,
-                  caption_loader=caption_loader,
-                  enable_ocr=enable_ocr,
-                  )
-
-    if n_jobs != 1 and len(globs_non_image_types) > 1:
-        # avoid nesting, e.g. upload 1 zip and then inside many files
-        # harder to handle if upload many zips with many files, inner parallel one will be disabled by joblib
-        documents = ProgressParallel(n_jobs=n_jobs, verbose=10 if verbose else 0, backend='multiprocessing')(
-            delayed(path_to_doc1)(file, **kwargs) for file in globs_non_image_types
-        )
-    else:
-        documents = [path_to_doc1(file, **kwargs) for file in tqdm(globs_non_image_types)]
-
-    # do images separately since can't fork after cuda in parent, so can't be parallel
-    if n_jobs_image != 1 and len(globs_image_types) > 1:
-        # avoid nesting, e.g. upload 1 zip and then inside many files
-        # harder to handle if upload many zips with many files, inner parallel one will be disabled by joblib
-        image_documents = ProgressParallel(n_jobs=n_jobs, verbose=10 if verbose else 0, backend='multiprocessing')(
-            delayed(path_to_doc1)(file, **kwargs) for file in globs_image_types
-        )
-    else:
-        image_documents = [path_to_doc1(file, **kwargs) for file in tqdm(globs_image_types)]
-
-    # add image docs in
-    documents += image_documents
-
-    if return_file:
-        # then documents really are files
-        files = documents.copy()
-        documents = []
-        for fil in files:
-            with open(fil, 'rb') as f:
-                documents.extend(pickle.load(f))
-            # remove temp pickle
-            os.remove(fil)
-    else:
-        documents = reduce(concat, documents)
-    return documents
-
-
-def prep_langchain(persist_directory,
-                   load_db_if_exists,
-                   db_type, use_openai_embedding, langchain_mode, user_path,
-                   hf_embedding_model, n_jobs=-1, kwargs_make_db={}):
-    """
-    do prep first time, involving downloads
-    # FIXME: Add github caching then add here
-    :return:
-    """
-    assert langchain_mode not in ['MyData'], "Should not prep scratch data"
-
-    db_dir_exists = os.path.isdir(persist_directory)
-
-    if db_dir_exists and user_path is None:
-        print("Prep: persist_directory=%s exists, using" % persist_directory, flush=True)
-        db = get_existing_db(persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
-                             hf_embedding_model)
-    else:
-        if db_dir_exists and user_path is not None:
-            print("Prep: persist_directory=%s exists, user_path=%s passed, adding any changed or new documents" % (
-            persist_directory, user_path), flush=True)
-        elif not db_dir_exists:
-            print("Prep: persist_directory=%s does not exist, regenerating" % persist_directory, flush=True)
-        db = None
-        if langchain_mode in ['All', 'DriverlessAI docs']:
-            # FIXME: Could also just use dai_docs.pickle directly and upload that
-            get_dai_docs(from_hf=True)
-
-        if langchain_mode in ['All', 'wiki']:
-            get_wiki_sources(first_para=kwargs_make_db['first_para'], text_limit=kwargs_make_db['text_limit'])
-
-        langchain_kwargs = kwargs_make_db.copy()
-        langchain_kwargs.update(locals())
-        db, num_new_sources, new_sources_metadata = make_db(**langchain_kwargs)
-
-    return db
-
-
-import posthog
-
-posthog.disabled = True
-
-
-class FakeConsumer(object):
-    def __init__(self, *args, **kwargs):
-        pass
-
-    def run(self):
-        pass
-
-    def pause(self):
-        pass
-
-    def upload(self):
-        pass
-
-    def next(self):
-        pass
-
-    def request(self, batch):
-        pass
-
-
-posthog.Consumer = FakeConsumer
-
-
-def get_existing_db(persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
-                    hf_embedding_model):
-    if load_db_if_exists and db_type == 'chroma' and os.path.isdir(persist_directory) and os.path.isdir(
-            os.path.join(persist_directory, 'index')):
-        print("DO Loading db: %s" % langchain_mode, flush=True)
-        embedding = get_embedding(use_openai_embedding, hf_embedding_model=hf_embedding_model)
-        from chromadb.config import Settings
-        client_settings = Settings(anonymized_telemetry=False,
-                                   chroma_db_impl="duckdb+parquet",
-                                   persist_directory=persist_directory)
-        db = Chroma(persist_directory=persist_directory, embedding_function=embedding,
-                    collection_name=langchain_mode.replace(' ', '_'),
-                    client_settings=client_settings)
-        print("DONE Loading db: %s" % langchain_mode, flush=True)
-        return db
-    return None
-
-
-def make_db(**langchain_kwargs):
-    func_names = list(inspect.signature(_make_db).parameters)
-    missing_kwargs = [x for x in func_names if x not in langchain_kwargs]
-    defaults_db = {k: v.default for k, v in dict(inspect.signature(run_qa_db).parameters).items()}
-    for k in missing_kwargs:
-        if k in defaults_db:
-            langchain_kwargs[k] = defaults_db[k]
-    # final check for missing
-    missing_kwargs = [x for x in func_names if x not in langchain_kwargs]
-    assert not missing_kwargs, "Missing kwargs: %s" % missing_kwargs
-    # only keep actual used
-    langchain_kwargs = {k: v for k, v in langchain_kwargs.items() if k in func_names}
-    return _make_db(**langchain_kwargs)
-
-
-def _make_db(use_openai_embedding=False,
-             hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
-             first_para=False, text_limit=None, chunk=False, chunk_size=1024,
-             langchain_mode=None,
-             user_path=None,
-             db_type='faiss',
-             load_db_if_exists=True,
-             db=None,
-             n_jobs=-1,
-             verbose=False):
-    persist_directory = 'db_dir_%s' % langchain_mode  # single place, no special names for each case
-    if not db and load_db_if_exists and db_type == 'chroma' and os.path.isdir(persist_directory) and os.path.isdir(
-            os.path.join(persist_directory, 'index')):
-        assert langchain_mode not in ['MyData'], "Should not load MyData db this way"
-        print("Loading existing db", flush=True)
-        embedding = get_embedding(use_openai_embedding, hf_embedding_model=hf_embedding_model)
-        from chromadb.config import Settings
-        client_settings = Settings(anonymized_telemetry=False,
-                                   chroma_db_impl="duckdb+parquet",
-                                   persist_directory=persist_directory)
-        db = Chroma(persist_directory=persist_directory, embedding_function=embedding,
-                    collection_name=langchain_mode.replace(' ', '_'),
-                    client_settings=client_settings)
-    sources = []
-    if not db and langchain_mode not in ['MyData'] or \
-            user_path is not None and \
-            langchain_mode in ['UserData']:
-        # Should not make MyData db this way, why avoided, only upload from UI
-        assert langchain_mode not in ['MyData'], "Should not make MyData db this way"
-        if verbose:
-            if langchain_mode in ['UserData']:
-                if user_path is not None:
-                    print("Checking if changed or new sources in %s, and generating sources them" % user_path,
-                          flush=True)
-                elif db is None:
-                    print("user_path not passed and no db, no sources", flush=True)
-                else:
-                    print("user_path not passed, using only existing db, no new sources", flush=True)
-            else:
-                print("Generating %s sources" % langchain_mode, flush=True)
-        if langchain_mode in ['wiki_full', 'All', "'All'"]:
-            from read_wiki_full import get_all_documents
-            small_test = None
-            print("Generating new wiki", flush=True)
-            sources1 = get_all_documents(small_test=small_test, n_jobs=os.cpu_count() // 2)
-            print("Got new wiki", flush=True)
-            if chunk:
-                sources1 = chunk_sources(sources1, chunk_size=chunk_size)
-                print("Chunked new wiki", flush=True)
-            sources.extend(sources1)
-        if langchain_mode in ['wiki', 'All', "'All'"]:
-            sources1 = get_wiki_sources(first_para=first_para, text_limit=text_limit)
-            if chunk:
-                sources1 = chunk_sources(sources1, chunk_size=chunk_size)
-            sources.extend(sources1)
-        if langchain_mode in ['github h2oGPT', 'All', "'All'"]:
-            # sources = get_github_docs("dagster-io", "dagster")
-            sources1 = get_github_docs("h2oai", "h2ogpt")
-            # FIXME: always chunk for now
-            sources1 = chunk_sources(sources1, chunk_size=chunk_size)
-            sources.extend(sources1)
-        if langchain_mode in ['DriverlessAI docs', 'All', "'All'"]:
-            sources1 = get_dai_docs(from_hf=True)
-            if chunk and False:  # FIXME: DAI docs are already chunked well, should only chunk more if over limit
-                sources1 = chunk_sources(sources1, chunk_size=chunk_size)
-            sources.extend(sources1)
-        if langchain_mode in ['All', 'UserData']:
-            if user_path:
-                if db is not None:
-                    # NOTE: Ignore file names for now, only go by hash ids
-                    # existing_files = get_existing_files(db)
-                    existing_files = []
-                    existing_hash_ids = get_existing_hash_ids(db)
-                else:
-                    # pretend no existing files so won't filter
-                    existing_files = []
-                    existing_hash_ids = []
-                # chunk internally for speed over multiple docs
-                sources1 = path_to_docs(user_path, n_jobs=n_jobs, chunk=chunk, chunk_size=chunk_size,
-                                        existing_files=existing_files, existing_hash_ids=existing_hash_ids)
-                new_metadata_sources = set([x.metadata['source'] for x in sources1])
-                if new_metadata_sources:
-                    print("Loaded %s new files as sources to add to UserData" % len(new_metadata_sources), flush=True)
-                    if verbose:
-                        print("Files added: %s" % '\n'.join(new_metadata_sources), flush=True)
-                sources.extend(sources1)
-                print("Loaded %s sources for potentially adding to UserData" % len(sources), flush=True)
-            else:
-                print("Chose UserData but user_path is empty/None", flush=True)
-        if False and langchain_mode in ['urls', 'All', "'All'"]:
-            # from langchain.document_loaders import UnstructuredURLLoader
-            # loader = UnstructuredURLLoader(urls=urls)
-            urls = ["https://www.birdsongsf.com/who-we-are/"]
-            from langchain.document_loaders import PlaywrightURLLoader
-            loader = PlaywrightURLLoader(urls=urls, remove_selectors=["header", "footer"])
-            sources1 = loader.load()
-            sources.extend(sources1)
-        if not sources:
-            if verbose:
-                if db is not None:
-                    print("langchain_mode %s has no new sources, nothing to add to db" % langchain_mode, flush=True)
-                else:
-                    print("langchain_mode %s has no sources, not making new db" % langchain_mode, flush=True)
-            return db, 0, []
-        if verbose:
-            if db is not None:
-                print("Generating db", flush=True)
-            else:
-                print("Adding to db", flush=True)
-    if not db:
-        if sources:
-            db = get_db(sources, use_openai_embedding=use_openai_embedding, db_type=db_type,
-                        persist_directory=persist_directory, langchain_mode=langchain_mode,
-                        hf_embedding_model=hf_embedding_model)
-            if verbose:
-                print("Generated db", flush=True)
-        else:
-            print("Did not generate db since no sources", flush=True)
-        new_sources_metadata = [x.metadata for x in sources]
-    elif user_path is not None and langchain_mode in ['UserData']:
-        print("Existing db, potentially adding %s sources from user_path=%s" % (len(sources), user_path), flush=True)
-        db, num_new_sources, new_sources_metadata = add_to_db(db, sources, db_type=db_type)
-        print("Existing db, added %s new sources from user_path=%s" % (num_new_sources, user_path), flush=True)
-    else:
-        new_sources_metadata = [x.metadata for x in sources]
-
-    return db, len(new_sources_metadata), new_sources_metadata
-
-
-def get_existing_files(db):
-    collection = db.get()
-    metadata_sources = set([x['source'] for x in collection['metadatas']])
-    return metadata_sources
-
-
-def get_existing_hash_ids(db):
-    collection = db.get()
-    # assume consistency, that any prior hashed source was single hashed file at the time among all source chunks
-    metadata_hash_ids = {x['source']: x.get('hashid') for x in collection['metadatas']}
-    return metadata_hash_ids
-
-
-source_prefix = "Sources [Score | Link]:"
-source_postfix = "End Sources<p>"
-
-
-def run_qa_db(**kwargs):
-    func_names = list(inspect.signature(_run_qa_db).parameters)
-    # hard-coded defaults
-    kwargs['answer_with_sources'] = True
-    kwargs['sanitize_bot_response'] = True
-    kwargs['show_rank'] = False
-    missing_kwargs = [x for x in func_names if x not in kwargs]
-    assert not missing_kwargs, "Missing kwargs: %s" % missing_kwargs
-    # only keep actual used
-    kwargs = {k: v for k, v in kwargs.items() if k in func_names}
-    return _run_qa_db(**kwargs)
-
-
-def _run_qa_db(query=None,
-               use_openai_model=False, use_openai_embedding=False,
-               first_para=False, text_limit=None, k=4, chunk=False, chunk_size=1024,
-               user_path=None,
-               detect_user_path_changes_every_query=False,
-               db_type='faiss',
-               model_name=None, model=None, tokenizer=None,
-               hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
-               stream_output=False,
-               prompter=None,
-               prompt_type=None,
-               answer_with_sources=True,
-               cut_distanct=1.1,
-               sanitize_bot_response=True,
-               show_rank=False,
-               load_db_if_exists=False,
-               db=None,
-               max_new_tokens=256,
-               temperature=0.1,
-               repetition_penalty=1.0,
-               top_k=40,
-               top_p=0.7,
-               langchain_mode=None,
-               document_choice=['All'],
-               n_jobs=-1,
-               verbose=False,
-               cli=False):
-    """
-
-    :param query:
-    :param use_openai_model:
-    :param use_openai_embedding:
-    :param first_para:
-    :param text_limit:
-    :param k:
-    :param chunk:
-    :param chunk_size:
-    :param user_path: user path to glob recursively from
-    :param db_type: 'faiss' for in-memory db or 'chroma' or 'weaviate' for persistent db
-    :param model_name: model name, used to switch behaviors
-    :param model: pre-initialized model, else will make new one
-    :param tokenizer: pre-initialized tokenizer, else will make new one.  Required not None if model is not None
-    :param answer_with_sources
-    :return:
-    """
-    assert query is not None
-    assert prompter is not None or prompt_type is not None or model is None  # if model is None, then will generate
-    if prompter is not None:
-        prompt_type = prompter.prompt_type
-    if model is not None:
-        assert prompt_type is not None
-    llm, model_name, streamer, prompt_type_out = get_llm(use_openai_model=use_openai_model, model_name=model_name,
-                                                         model=model, tokenizer=tokenizer,
-                                                         stream_output=stream_output,
-                                                         max_new_tokens=max_new_tokens,
-                                                         temperature=temperature,
-                                                         repetition_penalty=repetition_penalty,
-                                                         top_k=top_k,
-                                                         top_p=top_p,
-                                                         prompt_type=prompt_type,
-                                                         prompter=prompter,
-                                                         verbose=verbose,
-                                                         )
-
-    if model_name in non_hf_types:
-        # FIXME: for now, streams to stdout/stderr currently
-        stream_output = False
-
-    use_context = False
-    scores = []
-    chain = None
-
-    func_names = list(inspect.signature(get_similarity_chain).parameters)
-    sim_kwargs = {k: v for k, v in locals().items() if k in func_names}
-    missing_kwargs = [x for x in func_names if x not in sim_kwargs]
-    assert not missing_kwargs, "Missing: %s" % missing_kwargs
-    docs, chain, scores, use_context = get_similarity_chain(**sim_kwargs)
-    if len(document_choice) > 0 and document_choice[0] == 'Only':
-        formatted_doc_chunks = '\n\n'.join([get_url(x) + '\n\n' + x.page_content for x in docs])
-        yield formatted_doc_chunks, ''
-        return
-    if chain is None and model_name not in non_hf_types:
-        # can only return if HF type
-        return
-
-    if stream_output:
-        answer = None
-        assert streamer is not None
-        import queue
-        bucket = queue.Queue()
-        thread = EThread(target=chain, streamer=streamer, bucket=bucket)
-        thread.start()
-        outputs = ""
-        prompt = None  # FIXME
-        try:
-            for new_text in streamer:
-                # print("new_text: %s" % new_text, flush=True)
-                if bucket.qsize() > 0 or thread.exc:
-                    thread.join()
-                outputs += new_text
-                if prompter:  # and False:  # FIXME: pipeline can already use prompter
-                    output1 = prompter.get_response(outputs, prompt=prompt,
-                                                    sanitize_bot_response=sanitize_bot_response)
-                    yield output1, ''
-                else:
-                    yield outputs, ''
-        except BaseException:
-            # if any exception, raise that exception if was from thread, first
-            if thread.exc:
-                raise thread.exc
-            raise
-        finally:
-            # in case no exception and didn't join with thread yet, then join
-            if not thread.exc:
-                answer = thread.join()
-        # in case raise StopIteration or broke queue loop in streamer, but still have exception
-        if thread.exc:
-            raise thread.exc
-        # FIXME: answer is not string outputs from streamer.  How to get actual final output?
-        # answer = outputs
-    else:
-        answer = chain()
-
-    if not use_context:
-        ret = answer['output_text']
-        extra = ''
-        yield ret, extra
-    elif answer is not None:
-        ret, extra = get_sources_answer(query, answer, scores, show_rank, answer_with_sources, verbose=verbose)
-        yield ret, extra
-    return
-
-
-def get_similarity_chain(query=None,
-                         use_openai_model=False, use_openai_embedding=False,
-                         first_para=False, text_limit=None, k=4, chunk=False, chunk_size=1024,
-                         user_path=None,
-                         detect_user_path_changes_every_query=False,
-                         db_type='faiss',
-                         model_name=None,
-                         hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
-                         prompt_type=None,
-                         cut_distanct=1.1,
-                         load_db_if_exists=False,
-                         db=None,
-                         langchain_mode=None,
-                         document_choice=['All'],
-                         n_jobs=-1,
-                         # beyond run_db_query:
-                         llm=None,
-                         verbose=False,
-                         ):
-    # determine whether use of context out of docs is planned
-    if not use_openai_model and prompt_type not in ['plain'] or model_name in non_hf_types:
-        if langchain_mode in ['Disabled', 'ChatLLM', 'LLM']:
-            use_context = False
-        else:
-            use_context = True
-    else:
-        use_context = True
-
-    # https://github.com/hwchase17/langchain/issues/1946
-    # FIXME: Seems to way to get size of chroma db to limit k to avoid
-    # Chroma collection MyData contains fewer than 4 elements.
-    # type logger error
-    k_db = 1000 if db_type == 'chroma' else k  # k=100 works ok too for
-
-    # FIXME: For All just go over all dbs instead of a separate db for All
-    if not detect_user_path_changes_every_query and db is not None:
-        # avoid looking at user_path during similarity search db handling,
-        # if already have db and not updating from user_path every query
-        # but if db is None, no db yet loaded (e.g. from prep), so allow user_path to be whatever it was
-        user_path = None
-    db, num_new_sources, new_sources_metadata = make_db(use_openai_embedding=use_openai_embedding,
-                                                        hf_embedding_model=hf_embedding_model,
-                                                        first_para=first_para, text_limit=text_limit, chunk=chunk,
-                                                        chunk_size=chunk_size,
-                                                        langchain_mode=langchain_mode,
-                                                        user_path=user_path,
-                                                        db_type=db_type,
-                                                        load_db_if_exists=load_db_if_exists,
-                                                        db=db,
-                                                        n_jobs=n_jobs,
-                                                        verbose=verbose)
-
-    if db and use_context:
-        if isinstance(document_choice, str):
-            # support string as well
-            document_choice = [document_choice]
-        if not isinstance(db, Chroma) or \
-                len(document_choice) == 0 or \
-                len(document_choice) <= 1 and document_choice[0] == 'All':
-            # treat empty list as All for now, not 'None'
-            filter_kwargs = {}
-        elif len(document_choice) > 0 and document_choice[0] == 'Only':
-            # Only means All docs, but only will return sources, not LLM response
-            filter_kwargs = {}
-        else:
-            if len(document_choice) >= 2:
-                or_filter = [{"source": {"$eq": x}} for x in document_choice]
-                filter_kwargs = dict(filter={"$or": or_filter})
-            elif len(document_choice) > 0:
-                one_filter = [{"source": {"$eq": x}} for x in document_choice][0]
-                filter_kwargs = dict(filter=one_filter)
-            else:
-                filter_kwargs = {}
-            if len(document_choice) == 1 and document_choice[0] == 'None':
-                k_db = 1
-                k = 0
-        docs_with_score = db.similarity_search_with_score(query, k=k_db, **filter_kwargs)[:k]
-        # cut off so no high distance docs/sources considered
-        docs = [x[0] for x in docs_with_score if x[1] < cut_distanct]
-        scores = [x[1] for x in docs_with_score if x[1] < cut_distanct]
-        if len(scores) > 0 and verbose:
-            print("Distance: min: %s max: %s mean: %s median: %s" %
-                  (scores[0], scores[-1], np.mean(scores), np.median(scores)), flush=True)
-    else:
-        docs = []
-        scores = []
-
-    if not docs and use_context and model_name not in non_hf_types:
-        # if HF type and have no docs, can bail out
-        return docs, None, [], False
-
-    if len(document_choice) > 0 and document_choice[0] == 'Only':
-        # no LLM use
-        return docs, None, [], False
-
-    common_words_file = "data/NGSL_1.2_stats.csv.zip"
-    if os.path.isfile(common_words_file):
-        df = pd.read_csv("data/NGSL_1.2_stats.csv.zip")
-        import string
-        reduced_query = query.translate(str.maketrans(string.punctuation, ' ' * len(string.punctuation))).strip()
-        reduced_query_words = reduced_query.split(' ')
-        set_common = set(df['Lemma'].values.tolist())
-        num_common = len([x.lower() in set_common for x in reduced_query_words])
-        frac_common = num_common / len(reduced_query) if reduced_query else 0
-        # FIXME: report to user bad query that uses too many common words
-        if verbose:
-            print("frac_common: %s" % frac_common, flush=True)
-
-    if len(docs) == 0:
-        # avoid context == in prompt then
-        use_context = False
-
-    if not use_openai_model and prompt_type not in ['plain'] or model_name in non_hf_types:
-        # instruct-like, rather than few-shot prompt_type='plain' as default
-        # but then sources confuse the model with how inserted among rest of text, so avoid
-        prefix = ""
-        if langchain_mode in ['Disabled', 'ChatLLM', 'LLM'] or not use_context:
-            template = """%s{context}{question}""" % prefix
-        else:
-            template = """%s
-==
-{context}
-==
-{question}""" % prefix
-        prompt = PromptTemplate(
-            # input_variables=["summaries", "question"],
-            input_variables=["context", "question"],
-            template=template,
-        )
-        chain = load_qa_chain(llm, prompt=prompt)
-    else:
-        chain = load_qa_with_sources_chain(llm)
-
-    if not use_context:
-        chain_kwargs = dict(input_documents=[], question=query)
-    else:
-        chain_kwargs = dict(input_documents=docs, question=query)
-
-    target = wrapped_partial(chain, chain_kwargs)
-    return docs, target, scores, use_context
-
-
-def get_sources_answer(query, answer, scores, show_rank, answer_with_sources, verbose=False):
-    if verbose:
-        print("query: %s" % query, flush=True)
-        print("answer: %s" % answer['output_text'], flush=True)
-
-    if len(answer['input_documents']) == 0:
-        extra = ''
-        ret = answer['output_text'] + extra
-        return ret, extra
-
-    # link
-    answer_sources = [(max(0.0, 1.5 - score) / 1.5, get_url(doc)) for score, doc in
-                      zip(scores, answer['input_documents'])]
-    answer_sources_dict = defaultdict(list)
-    [answer_sources_dict[url].append(score) for score, url in answer_sources]
-    answers_dict = {}
-    for url, scores_url in answer_sources_dict.items():
-        answers_dict[url] = np.max(scores_url)
-    answer_sources = [(score, url) for url, score in answers_dict.items()]
-    answer_sources.sort(key=lambda x: x[0], reverse=True)
-    if show_rank:
-        # answer_sources = ['%d | %s' % (1 + rank, url) for rank, (score, url) in enumerate(answer_sources)]
-        # sorted_sources_urls = "Sources [Rank | Link]:<br>" + "<br>".join(answer_sources)
-        answer_sources = ['%s' % url for rank, (score, url) in enumerate(answer_sources)]
-        sorted_sources_urls = "Ranked Sources:<br>" + "<br>".join(answer_sources)
-    else:
-        answer_sources = ['<li>%.2g | %s</li>' % (score, url) for score, url in answer_sources]
-        sorted_sources_urls = f"{source_prefix}<p><ul>" + "<p>".join(answer_sources)
-        sorted_sources_urls += f"</ul></p>{source_postfix}"
-
-    if not answer['output_text'].endswith('\n'):
-        answer['output_text'] += '\n'
-
-    if answer_with_sources:
-        extra = '\n' + sorted_sources_urls
-    else:
-        extra = ''
-    ret = answer['output_text'] + extra
-    return ret, extra
-
-
-def chunk_sources(sources, chunk_size=1024):
-    source_chunks = []
-    # Below for known separator
-    # splitter = CharacterTextSplitter(separator=" ", chunk_size=chunk_size, chunk_overlap=0)
-    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=0)
-    for source in sources:
-        # print(source.metadata['source'], flush=True)
-        for chunky in splitter.split_text(source.page_content):
-            source_chunks.append(Document(page_content=chunky, metadata=source.metadata))
-    return source_chunks
-
-
-def get_db_from_hf(dest=".", db_dir='db_dir_DriverlessAI_docs.zip'):
-    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/db_dirs', db_dir, token=token, repo_type='dataset')
-    import zipfile
-    with zipfile.ZipFile(path_to_zip_file, 'r') as zip_ref:
-        zip_ref.extractall(dest)
-    return path_to_zip_file
-
-
-# Note dir has space in some cases, while zip does not
-some_db_zips = [['db_dir_DriverlessAI_docs.zip', 'db_dir_DriverlessAI docs', 'CC-BY-NC license'],
-                ['db_dir_UserData.zip', 'db_dir_UserData', 'CC-BY license for ArXiv'],
-                ['db_dir_github_h2oGPT.zip', 'db_dir_github h2oGPT', 'ApacheV2 license'],
-                ['db_dir_wiki.zip', 'db_dir_wiki', 'CC-BY-SA Wikipedia license'],
-                # ['db_dir_wiki_full.zip', 'db_dir_wiki_full.zip', '23GB, 05/04/2023 CC-BY-SA Wiki license'],
-                ]
-
-all_db_zips = some_db_zips + \
-              [['db_dir_wiki_full.zip', 'db_dir_wiki_full.zip', '23GB, 05/04/2023 CC-BY-SA Wiki license'],
-               ]
-
-
-def get_some_dbs_from_hf(dest='.', db_zips=None):
-    if db_zips is None:
-        db_zips = some_db_zips
-    for db_dir, dir_expected, license1 in db_zips:
-        path_to_zip_file = get_db_from_hf(dest=dest, db_dir=db_dir)
-        assert os.path.isfile(path_to_zip_file), "Missing zip in %s" % path_to_zip_file
-        if dir_expected:
-            assert os.path.isdir(os.path.join(dest, dir_expected)), "Missing path for %s" % dir_expected
-            assert os.path.isdir(os.path.join(dest, dir_expected, 'index')), "Missing index in %s" % dir_expected
-
-
-if __name__ == '__main__':
-    pass

gpt_langchain.py

@@ -0,0 +1 @@
+../../gpt_langchain.py

gradio_runner.py

@@ -1,1741 +0,0 @@
-import copy
-import functools
-import inspect
-import json
-import os
-import random
-import sys
-import traceback
-import uuid
-import filelock
-import pandas as pd
-import requests
-import tabulate
-
-# This is a hack to prevent Gradio from phoning home when it gets imported
-os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False'
-
-
-def my_get(url, **kwargs):
-    print('Gradio HTTP request redirected to localhost :)', flush=True)
-    kwargs.setdefault('allow_redirects', True)
-    return requests.api.request('get', 'http://127.0.0.1/', **kwargs)
-
-
-original_get = requests.get
-requests.get = my_get
-import gradio as gr
-
-requests.get = original_get
-
-from gradio_themes import H2oTheme, SoftTheme, get_h2o_title, get_simple_title, get_dark_js
-from prompter import Prompter, \
-    prompt_type_to_model_name, prompt_types_strings, inv_prompt_type_to_model_lower, generate_prompt, non_hf_types
-from utils import get_githash, flatten_list, zip_data, s3up, clear_torch_cache, get_torch_allocated, system_info_print, \
-    ping, get_short_name, get_url, makedirs, get_kwargs
-from generate import get_model, languages_covered, evaluate, eval_func_param_names, score_qa, langchain_modes, \
-    inputs_kwargs_list, get_cutoffs, scratch_base_dir
-
-from apscheduler.schedulers.background import BackgroundScheduler
-
-
-def go_gradio(**kwargs):
-    allow_api = kwargs['allow_api']
-    is_public = kwargs['is_public']
-    is_hf = kwargs['is_hf']
-    memory_restriction_level = kwargs['memory_restriction_level']
-    n_gpus = kwargs['n_gpus']
-    admin_pass = kwargs['admin_pass']
-    model_state0 = kwargs['model_state0']
-    score_model_state0 = kwargs['score_model_state0']
-    dbs = kwargs['dbs']
-    db_type = kwargs['db_type']
-    visible_langchain_modes = kwargs['visible_langchain_modes']
-    allow_upload_to_user_data = kwargs['allow_upload_to_user_data']
-    allow_upload_to_my_data = kwargs['allow_upload_to_my_data']
-    enable_sources_list = kwargs['enable_sources_list']
-    enable_url_upload = kwargs['enable_url_upload']
-    enable_text_upload = kwargs['enable_text_upload']
-    use_openai_embedding = kwargs['use_openai_embedding']
-    hf_embedding_model = kwargs['hf_embedding_model']
-    enable_captions = kwargs['enable_captions']
-    captions_model = kwargs['captions_model']
-    enable_ocr = kwargs['enable_ocr']
-    caption_loader = kwargs['caption_loader']
-
-    # easy update of kwargs needed for evaluate() etc.
-    queue = True
-    allow_upload = allow_upload_to_user_data or allow_upload_to_my_data
-    kwargs.update(locals())
-
-    if 'mbart-' in kwargs['model_lower']:
-        instruction_label_nochat = "Text to translate"
-    else:
-        instruction_label_nochat = "Instruction (Shift-Enter or push Submit to send message," \
-                                   " use Enter for multiple input lines)"
-    if kwargs['input_lines'] > 1:
-        instruction_label = "You (Shift-Enter or push Submit to send message, use Enter for multiple input lines)"
-    else:
-        instruction_label = "You (Enter or push Submit to send message, shift-enter for more lines)"
-
-    title = 'h2oGPT'
-    if 'h2ogpt-research' in kwargs['base_model']:
-        title += " [Research demonstration]"
-    more_info = """For more information, visit our GitHub pages: [h2oGPT](https://github.com/h2oai/h2ogpt) and [H2O-LLMStudio](https://github.com/h2oai/h2o-llmstudio)<br>"""
-    if is_public:
-        more_info += """<iframe src="https://ghbtns.com/github-btn.html?user=h2oai&repo=h2ogpt&type=star&count=true&size=small" frameborder="0" scrolling="0" width="150" height="20" title="GitHub"></iframe>"""
-    if kwargs['verbose']:
-        description = f"""Model {kwargs['base_model']} Instruct dataset.
-                      For more information, visit our GitHub pages: [h2oGPT](https://github.com/h2oai/h2ogpt) and [H2O LLM Studio](https://github.com/h2oai/h2o-llmstudio).
-                      Command: {str(' '.join(sys.argv))}
-                      Hash: {get_githash()}
-                      """
-    else:
-        description = more_info
-    description += "If this host is busy, try [12B](https://gpt.h2o.ai), [Falcon 40B](http://falcon.h2o.ai), [HF Spaces1 12B](https://huggingface.co/spaces/h2oai/h2ogpt-chatbot) or [HF Spaces2 12B](https://huggingface.co/spaces/h2oai/h2ogpt-chatbot2)<br>"
-    description += """<p>By using h2oGPT, you accept our [Terms of Service](https://github.com/h2oai/h2ogpt/blob/main/docs/tos.md)</p>"""
-    if is_hf:
-        description += '''<a href="https://huggingface.co/spaces/h2oai/h2ogpt-chatbot?duplicate=true"><img src="https://bit.ly/3gLdBN6" style="white-space: nowrap" alt="Duplicate Space"></a>'''
-
-    if kwargs['verbose']:
-        task_info_md = f"""
-        ### Task: {kwargs['task_info']}"""
-    else:
-        task_info_md = ''
-
-    if kwargs['h2ocolors']:
-        css_code = """footer {visibility: hidden;}
-    body{background:linear-gradient(#f5f5f5,#e5e5e5);}
-    body.dark{background:linear-gradient(#000000,#0d0d0d);}
-    """
-    else:
-        css_code = """footer {visibility: hidden}"""
-    css_code += """
-@import url('https://fonts.googleapis.com/css2?family=Source+Sans+Pro:wght@400;600&display=swap');
-body.dark{#warning {background-color: #555555};}
-#small_btn {
-    margin: 0.6em 0em 0.55em 0;
-    max-width: 20em;
-    min-width: 5em !important;
-    height: 5em;
-    font-size: 14px !important
-}"""
-
-    if kwargs['gradio_avoid_processing_markdown']:
-        from gradio_client import utils as client_utils
-        from gradio.components import Chatbot
-
-        # gradio has issue with taking too long to process input/output for markdown etc.
-        # Avoid for now, allow raw html to render, good enough for chatbot.
-        def _postprocess_chat_messages(self, chat_message: str):
-            if chat_message is None:
-                return None
-            elif isinstance(chat_message, (tuple, list)):
-                filepath = chat_message[0]
-                mime_type = client_utils.get_mimetype(filepath)
-                filepath = self.make_temp_copy_if_needed(filepath)
-                return {
-                    "name": filepath,
-                    "mime_type": mime_type,
-                    "alt_text": chat_message[1] if len(chat_message) > 1 else None,
-                    "data": None,  # These last two fields are filled in by the frontend
-                    "is_file": True,
-                }
-            elif isinstance(chat_message, str):
-                return chat_message
-            else:
-                raise ValueError(f"Invalid message for Chatbot component: {chat_message}")
-
-        Chatbot._postprocess_chat_messages = _postprocess_chat_messages
-
-    if kwargs['gradio_offline_level'] >= 0:
-        # avoid GoogleFont that pulls from internet
-        if kwargs['gradio_offline_level'] == 1:
-            # front end would still have to download fonts or have cached it at some point
-            base_font = 'Source Sans Pro'
-        else:
-            base_font = 'Helvetica'
-        theme_kwargs = dict(font=(base_font, 'ui-sans-serif', 'system-ui', 'sans-serif'),
-                            font_mono=('IBM Plex Mono', 'ui-monospace', 'Consolas', 'monospace'))
-    else:
-        theme_kwargs = dict()
-
-    theme = H2oTheme(**theme_kwargs) if kwargs['h2ocolors'] else SoftTheme(**theme_kwargs)
-    demo = gr.Blocks(theme=theme, css=css_code, title="h2oGPT", analytics_enabled=False)
-    callback = gr.CSVLogger()
-
-    model_options = flatten_list(list(prompt_type_to_model_name.values())) + kwargs['extra_model_options']
-    if kwargs['base_model'].strip() not in model_options:
-        lora_options = [kwargs['base_model'].strip()] + model_options
-    lora_options = kwargs['extra_lora_options']
-    if kwargs['lora_weights'].strip() not in lora_options:
-        lora_options = [kwargs['lora_weights'].strip()] + lora_options
-    # always add in no lora case
-    # add fake space so doesn't go away in gradio dropdown
-    no_lora_str = no_model_str = '[None/Remove]'
-    lora_options = [no_lora_str] + kwargs['extra_lora_options']  # FIXME: why double?
-    # always add in no model case so can free memory
-    # add fake space so doesn't go away in gradio dropdown
-    model_options = [no_model_str] + model_options
-
-    # transcribe, will be detranscribed before use by evaluate()
-    if not kwargs['lora_weights'].strip():
-        kwargs['lora_weights'] = no_lora_str
-
-    if not kwargs['base_model'].strip():
-        kwargs['base_model'] = no_model_str
-
-    # transcribe for gradio
-    kwargs['gpu_id'] = str(kwargs['gpu_id'])
-
-    no_model_msg = 'h2oGPT [   !!! Please Load Model in Models Tab !!!   ]'
-    output_label0 = f'h2oGPT [Model: {kwargs.get("base_model")}]' if kwargs.get(
-        'base_model') else no_model_msg
-    output_label0_model2 = no_model_msg
-
-    with demo:
-        # avoid actual model/tokenizer here or anything that would be bad to deepcopy
-        # https://github.com/gradio-app/gradio/issues/3558
-        model_state = gr.State(['model', 'tokenizer', kwargs['device'], kwargs['base_model']])
-        model_state2 = gr.State([None, None, None, None])
-        model_options_state = gr.State([model_options])
-        lora_options_state = gr.State([lora_options])
-        my_db_state = gr.State([None, None])
-        chat_state = gr.State({})
-        # make user default first and default choice, dedup
-        docs_state00 = kwargs['document_choice'] + ['All', 'Only', 'None']
-        docs_state0 = []
-        [docs_state0.append(x) for x in docs_state00 if x not in docs_state0]
-        docs_state = gr.State(docs_state0)  # first is chosen as default
-        gr.Markdown(f"""
-            {get_h2o_title(title) if kwargs['h2ocolors'] else get_simple_title(title)}
-
-            {description}
-            {task_info_md}
-            """)
-        if is_hf:
-            gr.HTML(
-            )
-
-        # go button visible if
-        base_wanted = kwargs['base_model'] != no_model_str and kwargs['login_mode_if_model0']
-        go_btn = gr.Button(value="ENTER", visible=base_wanted, variant="primary")
-        normal_block = gr.Row(visible=not base_wanted)
-        with normal_block:
-            with gr.Tabs():
-                with gr.Row():
-                    col_nochat = gr.Column(visible=not kwargs['chat'])
-                    with col_nochat:  # FIXME: for model comparison, and check rest
-                        text_output_nochat = gr.Textbox(lines=5, label=output_label0).style(show_copy_button=True)
-                        instruction_nochat = gr.Textbox(
-                            lines=kwargs['input_lines'],
-                            label=instruction_label_nochat,
-                            placeholder=kwargs['placeholder_instruction'],
-                        )
-                        iinput_nochat = gr.Textbox(lines=4, label="Input context for Instruction",
-                                                   placeholder=kwargs['placeholder_input'])
-                        submit_nochat = gr.Button("Submit")
-                        flag_btn_nochat = gr.Button("Flag")
-                        if not kwargs['auto_score']:
-                            with gr.Column(visible=kwargs['score_model']):
-                                score_btn_nochat = gr.Button("Score last prompt & response")
-                                score_text_nochat = gr.Textbox("Response Score: NA", show_label=False)
-                        else:
-                            with gr.Column(visible=kwargs['score_model']):
-                                score_text_nochat = gr.Textbox("Response Score: NA", show_label=False)
-                    col_chat = gr.Column(visible=kwargs['chat'])
-                    with col_chat:
-                        with gr.Row():
-                            text_output = gr.Chatbot(label=output_label0).style(height=kwargs['height'] or 400)
-                            text_output2 = gr.Chatbot(label=output_label0_model2, visible=False).style(
-                                height=kwargs['height'] or 400)
-                        with gr.Row():
-                            with gr.Column(scale=50):
-                                instruction = gr.Textbox(
-                                    lines=kwargs['input_lines'],
-                                    label=instruction_label,
-                                    placeholder=kwargs['placeholder_instruction'],
-                                )
-                            with gr.Row():
-                                submit = gr.Button(value='Submit').style(full_width=False, size='sm')
-                                stop_btn = gr.Button(value="Stop").style(full_width=False, size='sm')
-                        with gr.Row():
-                            clear = gr.Button("Save Chat / New Chat")
-                            flag_btn = gr.Button("Flag")
-                            if not kwargs['auto_score']:  # FIXME: For checkbox model2
-                                with gr.Column(visible=kwargs['score_model']):
-                                    with gr.Row():
-                                        score_btn = gr.Button("Score last prompt & response").style(
-                                            full_width=False, size='sm')
-                                        score_text = gr.Textbox("Response Score: NA", show_label=False)
-                                    score_res2 = gr.Row(visible=False)
-                                    with score_res2:
-                                        score_btn2 = gr.Button("Score last prompt & response 2").style(
-                                            full_width=False, size='sm')
-                                        score_text2 = gr.Textbox("Response Score2: NA", show_label=False)
-                            else:
-                                with gr.Column(visible=kwargs['score_model']):
-                                    score_text = gr.Textbox("Response Score: NA", show_label=False)
-                                    score_text2 = gr.Textbox("Response Score2: NA", show_label=False, visible=False)
-                            retry = gr.Button("Regenerate")
-                            undo = gr.Button("Undo")
-                with gr.TabItem("Chat"):
-                    with gr.Row():
-                        if 'mbart-' in kwargs['model_lower']:
-                            src_lang = gr.Dropdown(list(languages_covered().keys()),
-                                                   value=kwargs['src_lang'],
-                                                   label="Input Language")
-                            tgt_lang = gr.Dropdown(list(languages_covered().keys()),
-                                                   value=kwargs['tgt_lang'],
-                                                   label="Output Language")
-                    radio_chats = gr.Radio(value=None, label="Saved Chats", visible=True, interactive=True,
-                                           type='value')
-                    with gr.Row():
-                        clear_chat_btn = gr.Button(value="Clear Chat", visible=True).style(size='sm')
-                        export_chats_btn = gr.Button(value="Export Chats to Download").style(size='sm')
-                        remove_chat_btn = gr.Button(value="Remove Selected Chat", visible=True).style(size='sm')
-                        add_to_chats_btn = gr.Button("Import Chats from Upload").style(size='sm')
-                    with gr.Row():
-                        chats_file = gr.File(interactive=False, label="Download Exported Chats")
-                        chatsup_output = gr.File(label="Upload Chat File(s)",
-                                                 file_types=['.json'],
-                                                 file_count='multiple',
-                                                 elem_id="warning", elem_classes="feedback")
-                with gr.TabItem("Data Source"):
-                    langchain_readme = get_url('https://github.com/h2oai/h2ogpt/blob/main/docs/README_LangChain.md',
-                                               from_str=True)
-                    gr.HTML(value=f"""LangChain Support Disabled<p>
-                            Run:<p>
-                            <code>
-                            python generate.py --langchain_mode=MyData
-                            </code>
-                            <p>
-                            For more options see: {langchain_readme}""",
-                            visible=kwargs['langchain_mode'] == 'Disabled', interactive=False)
-                    data_row1 = gr.Row(visible=kwargs['langchain_mode'] != 'Disabled')
-                    with data_row1:
-                        if is_hf:
-                            # don't show 'wiki' since only usually useful for internal testing at moment
-                            no_show_modes = ['Disabled', 'wiki']
-                        else:
-                            no_show_modes = ['Disabled']
-                        allowed_modes = visible_langchain_modes.copy()
-                        allowed_modes = [x for x in allowed_modes if x in dbs]
-                        allowed_modes += ['ChatLLM', 'LLM']
-                        if allow_upload_to_my_data and 'MyData' not in allowed_modes:
-                            allowed_modes += ['MyData']
-                        if allow_upload_to_user_data and 'UserData' not in allowed_modes:
-                            allowed_modes += ['UserData']
-                        langchain_mode = gr.Radio(
-                            [x for x in langchain_modes if x in allowed_modes and x not in no_show_modes],
-                            value=kwargs['langchain_mode'],
-                            label="Data Collection of Sources",
-                            visible=kwargs['langchain_mode'] != 'Disabled')
-                    data_row2 = gr.Row(visible=kwargs['langchain_mode'] != 'Disabled')
-                    with data_row2:
-                        with gr.Column(scale=50):
-                            document_choice = gr.Dropdown(docs_state.value,
-                                                          label="Choose Subset of Doc(s) in Collection [click get sources to update]",
-                                                          value=docs_state.value[0],
-                                                          interactive=True,
-                                                          multiselect=True,
-                                                          )
-                        with gr.Row(visible=kwargs['langchain_mode'] != 'Disabled' and enable_sources_list):
-                            get_sources_btn = gr.Button(value="Get Sources",
-                                                        ).style(full_width=False, size='sm')
-                            show_sources_btn = gr.Button(value="Show Sources",
-                                                         ).style(full_width=False, size='sm')
-                            refresh_sources_btn = gr.Button(value="Refresh Sources",
-                                                            ).style(full_width=False, size='sm')
-
-                    # import control
-                    if kwargs['langchain_mode'] != 'Disabled':
-                        from gpt_langchain import file_types, have_arxiv
-                    else:
-                        have_arxiv = False
-                        file_types = []
-
-                    upload_row = gr.Row(visible=kwargs['langchain_mode'] != 'Disabled' and allow_upload).style(
-                        equal_height=False)
-                    with upload_row:
-                        with gr.Column():
-                            file_types_str = '[' + ' '.join(file_types) + ']'
-                            fileup_output = gr.File(label=f'Upload {file_types_str}',
-                                                    file_types=file_types,
-                                                    file_count="multiple",
-                                                    elem_id="warning", elem_classes="feedback")
-                            with gr.Row():
-                                add_to_shared_db_btn = gr.Button("Add File(s) to UserData",
-                                                                 visible=allow_upload_to_user_data, elem_id='small_btn')
-                                add_to_my_db_btn = gr.Button("Add File(s) to Scratch MyData",
-                                                             visible=allow_upload_to_my_data,
-                                                             elem_id='small_btn' if allow_upload_to_user_data else None,
-                                                             ).style(
-                                    size='sm' if not allow_upload_to_user_data else None)
-                        with gr.Column(
-                                visible=kwargs['langchain_mode'] != 'Disabled' and allow_upload and enable_url_upload):
-                            url_label = 'URL (http/https) or ArXiv:' if have_arxiv else 'URL (http/https)'
-                            url_text = gr.Textbox(label=url_label, interactive=True)
-                            with gr.Row():
-                                url_user_btn = gr.Button(value='Add URL content to Shared UserData',
-                                                         visible=allow_upload_to_user_data, elem_id='small_btn')
-                                url_my_btn = gr.Button(value='Add URL content to Scratch MyData',
-                                                       visible=allow_upload_to_my_data,
-                                                       elem_id='small_btn' if allow_upload_to_user_data else None,
-                                                       ).style(size='sm' if not allow_upload_to_user_data else None)
-                        with gr.Column(
-                                visible=kwargs['langchain_mode'] != 'Disabled' and allow_upload and enable_text_upload):
-                            user_text_text = gr.Textbox(label='Paste Text [Shift-Enter more lines]', interactive=True)
-                            with gr.Row():
-                                user_text_user_btn = gr.Button(value='Add Text to Shared UserData',
-                                                               visible=allow_upload_to_user_data,
-                                                               elem_id='small_btn')
-                                user_text_my_btn = gr.Button(value='Add Text to Scratch MyData',
-                                                             visible=allow_upload_to_my_data,
-                                                             elem_id='small_btn' if allow_upload_to_user_data else None,
-                                                             ).style(
-                                    size='sm' if not allow_upload_to_user_data else None)
-                        with gr.Column(visible=False):
-                            # WIP:
-                            with gr.Row(visible=False).style(equal_height=False):
-                                github_textbox = gr.Textbox(label="Github URL")
-                                with gr.Row(visible=True):
-                                    github_shared_btn = gr.Button(value="Add Github to Shared UserData",
-                                                                  visible=allow_upload_to_user_data,
-                                                                  elem_id='small_btn')
-                                    github_my_btn = gr.Button(value="Add Github to Scratch MyData",
-                                                              visible=allow_upload_to_my_data, elem_id='small_btn')
-                    sources_row3 = gr.Row(visible=kwargs['langchain_mode'] != 'Disabled' and enable_sources_list).style(
-                        equal_height=False)
-                    with sources_row3:
-                        with gr.Column(scale=1):
-                            file_source = gr.File(interactive=False,
-                                                  label="Download File w/Sources [click get sources to make file]")
-                        with gr.Column(scale=2):
-                            pass
-                    sources_row = gr.Row(visible=kwargs['langchain_mode'] != 'Disabled' and enable_sources_list).style(
-                        equal_height=False)
-                    with sources_row:
-                        sources_text = gr.HTML(label='Sources Added', interactive=False)
-
-                with gr.TabItem("Expert"):
-                    with gr.Row():
-                        with gr.Column():
-                            stream_output = gr.components.Checkbox(label="Stream output",
-                                                                   value=kwargs['stream_output'])
-                            prompt_type = gr.Dropdown(prompt_types_strings,
-                                                      value=kwargs['prompt_type'], label="Prompt Type",
-                                                      visible=not is_public)
-                            prompt_type2 = gr.Dropdown(prompt_types_strings,
-                                                       value=kwargs['prompt_type'], label="Prompt Type Model 2",
-                                                       visible=not is_public and False)
-                            do_sample = gr.Checkbox(label="Sample",
-                                                    info="Enable sampler, required for use of temperature, top_p, top_k",
-                                                    value=kwargs['do_sample'])
-                            temperature = gr.Slider(minimum=0.01, maximum=3,
-                                                    value=kwargs['temperature'],
-                                                    label="Temperature",
-                                                    info="Lower is deterministic (but may lead to repeats), Higher more creative (but may lead to hallucinations)")
-                            top_p = gr.Slider(minimum=0, maximum=1,
-                                              value=kwargs['top_p'], label="Top p",
-                                              info="Cumulative probability of tokens to sample from")
-                            top_k = gr.Slider(
-                                minimum=0, maximum=100, step=1,
-                                value=kwargs['top_k'], label="Top k",
-                                info='Num. tokens to sample from'
-                            )
-                            # FIXME: https://github.com/h2oai/h2ogpt/issues/106
-                            if os.getenv('TESTINGFAIL'):
-                                max_beams = 8 if not (memory_restriction_level or is_public) else 1
-                            else:
-                                max_beams = 1
-                            num_beams = gr.Slider(minimum=1, maximum=max_beams, step=1,
-                                                  value=min(max_beams, kwargs['num_beams']), label="Beams",
-                                                  info="Number of searches for optimal overall probability.  "
-                                                       "Uses more GPU memory/compute")
-                            # FIXME: 2048 should be tokenizer.model_max_length, but may not even have model yet
-                            if kwargs['max_new_tokens']:
-                                max_max_new_tokens = kwargs['max_new_tokens']
-                            elif memory_restriction_level == 1:
-                                max_max_new_tokens = 768
-                            elif memory_restriction_level == 2:
-                                max_max_new_tokens = 512
-                            elif memory_restriction_level >= 3:
-                                max_max_new_tokens = 256
-                            else:
-                                max_max_new_tokens = 2048
-                            max_new_tokens = gr.Slider(
-                                minimum=1, maximum=max_max_new_tokens, step=1,
-                                value=min(max_max_new_tokens, kwargs['max_new_tokens']), label="Max output length",
-                            )
-                            min_new_tokens = gr.Slider(
-                                minimum=0, maximum=max_max_new_tokens, step=1,
-                                value=min(max_max_new_tokens, kwargs['min_new_tokens']), label="Min output length",
-                            )
-                            early_stopping = gr.Checkbox(label="EarlyStopping", info="Stop early in beam search",
-                                                         value=kwargs['early_stopping'])
-                            max_max_time = 60 * 5 if not is_public else 60 * 2
-                            if is_hf:
-                                max_max_time = min(max_max_time, 60 * 1)
-                            max_time = gr.Slider(minimum=0, maximum=max_max_time, step=1,
-                                                 value=min(max_max_time, kwargs['max_time']), label="Max. time",
-                                                 info="Max. time to search optimal output.")
-                            repetition_penalty = gr.Slider(minimum=0.01, maximum=3.0,
-                                                           value=kwargs['repetition_penalty'],
-                                                           label="Repetition Penalty")
-                            num_return_sequences = gr.Slider(minimum=1, maximum=10, step=1,
-                                                             value=kwargs['num_return_sequences'],
-                                                             label="Number Returns", info="Must be <= num_beams",
-                                                             visible=not is_public)
-                            iinput = gr.Textbox(lines=4, label="Input",
-                                                placeholder=kwargs['placeholder_input'],
-                                                visible=not is_public)
-                            context = gr.Textbox(lines=3, label="System Pre-Context",
-                                                 info="Directly pre-appended without prompt processing",
-                                                 visible=not is_public)
-                            chat = gr.components.Checkbox(label="Chat mode", value=kwargs['chat'],
-                                                          visible=not is_public)
-                            count_chat_tokens_btn = gr.Button(value="Count Chat Tokens", visible=not is_public)
-                            chat_token_count = gr.Textbox(label="Chat Token Count", value=None,
-                                                          visible=not is_public, interactive=False)
-                            top_k_docs = gr.Slider(minimum=0, maximum=20, step=1,
-                                                   value=kwargs['top_k_docs'],
-                                                   label="Number of document chunks",
-                                                   info="For LangChain",
-                                                   visible=not is_public)
-
-                with gr.TabItem("Models"):
-                    load_msg = "Load-Unload Model/LORA [unload works if did not use --base_model]" if not is_public \
-                        else "LOAD-UNLOAD DISABLED FOR HOSTED DEMO"
-                    load_msg2 = "Load-Unload Model/LORA 2 [unload works if did not use --base_model]" if not is_public \
-                        else "LOAD-UNLOAD DISABLED FOR HOSTED DEMO 2"
-                    compare_checkbox = gr.components.Checkbox(label="Compare Mode",
-                                                              value=False, visible=not is_public)
-                    with gr.Row():
-                        n_gpus_list = [str(x) for x in list(range(-1, n_gpus))]
-                        with gr.Column():
-                            with gr.Row():
-                                with gr.Column(scale=50):
-                                    model_choice = gr.Dropdown(model_options_state.value[0], label="Choose Model",
-                                                               value=kwargs['base_model'])
-                                    lora_choice = gr.Dropdown(lora_options_state.value[0], label="Choose LORA",
-                                                              value=kwargs['lora_weights'], visible=kwargs['show_lora'])
-                                with gr.Column(scale=1):
-                                    load_model_button = gr.Button(load_msg).style(full_width=False, size='sm')
-                                    model_load8bit_checkbox = gr.components.Checkbox(
-                                        label="Load 8-bit [requires support]",
-                                        value=kwargs['load_8bit'])
-                                    model_infer_devices_checkbox = gr.components.Checkbox(
-                                        label="Choose Devices [If not Checked, use all GPUs]",
-                                        value=kwargs['infer_devices'])
-                                    model_gpu = gr.Dropdown(n_gpus_list,
-                                                            label="GPU ID [-1 = all GPUs, if Choose is enabled]",
-                                                            value=kwargs['gpu_id'])
-                                    model_used = gr.Textbox(label="Current Model", value=kwargs['base_model'],
-                                                            interactive=False)
-                                    lora_used = gr.Textbox(label="Current LORA", value=kwargs['lora_weights'],
-                                                           visible=kwargs['show_lora'], interactive=False)
-                        col_model2 = gr.Column(visible=False)
-                        with col_model2:
-                            with gr.Row():
-                                with gr.Column(scale=50):
-                                    model_choice2 = gr.Dropdown(model_options_state.value[0], label="Choose Model 2",
-                                                                value=no_model_str)
-                                    lora_choice2 = gr.Dropdown(lora_options_state.value[0], label="Choose LORA 2",
-                                                               value=no_lora_str,
-                                                               visible=kwargs['show_lora'])
-                                with gr.Column(scale=1):
-                                    load_model_button2 = gr.Button(load_msg2).style(full_width=False, size='sm')
-                                    model_load8bit_checkbox2 = gr.components.Checkbox(
-                                        label="Load 8-bit 2 [requires support]",
-                                        value=kwargs['load_8bit'])
-                                    model_infer_devices_checkbox2 = gr.components.Checkbox(
-                                        label="Choose Devices 2 [If not Checked, use all GPUs]",
-                                        value=kwargs[
-                                            'infer_devices'])
-                                    model_gpu2 = gr.Dropdown(n_gpus_list,
-                                                             label="GPU ID 2 [-1 = all GPUs, if choose is enabled]",
-                                                             value=kwargs['gpu_id'])
-                                    # no model/lora loaded ever in model2 by default
-                                    model_used2 = gr.Textbox(label="Current Model 2", value=no_model_str)
-                                    lora_used2 = gr.Textbox(label="Current LORA 2", value=no_lora_str,
-                                                            visible=kwargs['show_lora'])
-                    with gr.Row():
-                        with gr.Column(scale=50):
-                            new_model = gr.Textbox(label="New Model HF name/path")
-                        with gr.Row():
-                            add_model_button = gr.Button("Add new model name").style(full_width=False, size='sm')
-                        with gr.Column(scale=50):
-                            new_lora = gr.Textbox(label="New LORA HF name/path", visible=kwargs['show_lora'])
-                        with gr.Row():
-                            add_lora_button = gr.Button("Add new LORA name", visible=kwargs['show_lora']).style(
-                                full_width=False, size='sm')
-                with gr.TabItem("System"):
-                    admin_row = gr.Row()
-                    with admin_row:
-                        admin_pass_textbox = gr.Textbox(label="Admin Password", type='password', visible=is_public)
-                        admin_btn = gr.Button(value="Admin Access", visible=is_public)
-                    system_row = gr.Row(visible=not is_public)
-                    with system_row:
-                        with gr.Column():
-                            with gr.Row():
-                                system_btn = gr.Button(value='Get System Info')
-                                system_text = gr.Textbox(label='System Info', interactive=False).style(
-                                    show_copy_button=True)
-
-                            with gr.Row():
-                                zip_btn = gr.Button("Zip")
-                                zip_text = gr.Textbox(label="Zip file name", interactive=False)
-                                file_output = gr.File(interactive=False, label="Zip file to Download")
-                            with gr.Row():
-                                s3up_btn = gr.Button("S3UP")
-                                s3up_text = gr.Textbox(label='S3UP result', interactive=False)
-                with gr.TabItem("Disclaimers"):
-                    description = ""
-                    description += """<p><b> DISCLAIMERS: </b><ul><i><li>The model was trained on The Pile and other data, which may contain objectionable content.  Use at own risk.</i></li>"""
-                    if kwargs['load_8bit']:
-                        description += """<i><li> Model is loaded in 8-bit and has other restrictions on this host. UX can be worse than non-hosted version.</i></li>"""
-                    description += """<i><li>Conversations may be used to improve h2oGPT.  Do not share sensitive information.</i></li>"""
-                    if 'h2ogpt-research' in kwargs['base_model']:
-                        description += """<i><li>Research demonstration only, not used for commercial purposes.</i></li>"""
-                    description += """<i><li>By using h2oGPT, you accept our <a href="https://github.com/h2oai/h2ogpt/blob/main/docs/tos.md">Terms of Service</a></i></li></ul></p>"""
-                    gr.Markdown(value=description, show_label=False, interactive=False)
-
-        # Get flagged data
-        zip_data1 = functools.partial(zip_data, root_dirs=['flagged_data_points', kwargs['save_dir']])
-        zip_btn.click(zip_data1, inputs=None, outputs=[file_output, zip_text], queue=False,
-                      api_name='zip_data' if allow_api else None)
-        s3up_btn.click(s3up, inputs=zip_text, outputs=s3up_text, queue=False,
-                       api_name='s3up_data' if allow_api else None)
-
-        def make_add_visible(x):
-            return gr.update(visible=x is not None)
-
-        def clear_file_list():
-            return None
-
-        def make_invisible():
-            return gr.update(visible=False)
-
-        def make_visible():
-            return gr.update(visible=True)
-
-        # Add to UserData
-        update_user_db_func = functools.partial(update_user_db, dbs=dbs, db_type=db_type, langchain_mode='UserData',
-                                                use_openai_embedding=use_openai_embedding,
-                                                hf_embedding_model=hf_embedding_model,
-                                                enable_captions=enable_captions,
-                                                captions_model=captions_model,
-                                                enable_ocr=enable_ocr,
-                                                caption_loader=caption_loader,
-                                                )
-
-        # note for update_user_db_func output is ignored for db
-        add_to_shared_db_btn.click(update_user_db_func,
-                                   inputs=[fileup_output, my_db_state, add_to_shared_db_btn, add_to_my_db_btn],
-                                   outputs=[add_to_shared_db_btn, add_to_my_db_btn, sources_text], queue=queue,
-                                   api_name='add_to_shared' if allow_api else None) \
-            .then(clear_file_list, outputs=fileup_output, queue=queue)
-
-        # .then(make_invisible, outputs=add_to_shared_db_btn, queue=queue)
-        # .then(make_visible, outputs=upload_button, queue=queue)
-
-        def clear_textbox():
-            return gr.Textbox.update(value='')
-
-        update_user_db_url_func = functools.partial(update_user_db_func, is_url=True)
-        url_user_btn.click(update_user_db_url_func,
-                           inputs=[url_text, my_db_state, add_to_shared_db_btn, add_to_my_db_btn],
-                           outputs=[add_to_shared_db_btn, add_to_my_db_btn, sources_text], queue=queue,
-                           api_name='add_url_to_shared' if allow_api else None) \
-            .then(clear_textbox, outputs=url_text, queue=queue)
-
-        update_user_db_txt_func = functools.partial(update_user_db_func, is_txt=True)
-        user_text_user_btn.click(update_user_db_txt_func,
-                                 inputs=[user_text_text, my_db_state, add_to_shared_db_btn, add_to_my_db_btn],
-                                 outputs=[add_to_shared_db_btn, add_to_my_db_btn, sources_text], queue=queue,
-                                 api_name='add_text_to_shared' if allow_api else None) \
-            .then(clear_textbox, outputs=user_text_text, queue=queue)
-
-        # Add to MyData
-        update_my_db_func = functools.partial(update_user_db, dbs=dbs, db_type=db_type, langchain_mode='MyData',
-                                              use_openai_embedding=use_openai_embedding,
-                                              hf_embedding_model=hf_embedding_model,
-                                              enable_captions=enable_captions,
-                                              captions_model=captions_model,
-                                              enable_ocr=enable_ocr,
-                                              caption_loader=caption_loader,
-                                              )
-
-        add_to_my_db_btn.click(update_my_db_func,
-                               inputs=[fileup_output, my_db_state, add_to_shared_db_btn, add_to_my_db_btn],
-                               outputs=[my_db_state, add_to_shared_db_btn, add_to_my_db_btn, sources_text], queue=queue,
-                               api_name='add_to_my' if allow_api else None) \
-            .then(clear_file_list, outputs=fileup_output, queue=queue)
-        # .then(make_invisible, outputs=add_to_shared_db_btn, queue=queue)
-        # .then(make_visible, outputs=upload_button, queue=queue)
-
-        update_my_db_url_func = functools.partial(update_my_db_func, is_url=True)
-        url_my_btn.click(update_my_db_url_func,
-                         inputs=[url_text, my_db_state, add_to_shared_db_btn, add_to_my_db_btn],
-                         outputs=[my_db_state, add_to_shared_db_btn, add_to_my_db_btn, sources_text], queue=queue,
-                         api_name='add_url_to_my' if allow_api else None) \
-            .then(clear_textbox, outputs=url_text, queue=queue)
-
-        update_my_db_txt_func = functools.partial(update_my_db_func, is_txt=True)
-        user_text_my_btn.click(update_my_db_txt_func,
-                               inputs=[user_text_text, my_db_state, add_to_shared_db_btn, add_to_my_db_btn],
-                               outputs=[my_db_state, add_to_shared_db_btn, add_to_my_db_btn, sources_text], queue=queue,
-                               api_name='add_txt_to_my' if allow_api else None) \
-            .then(clear_textbox, outputs=user_text_text, queue=queue)
-
-        get_sources1 = functools.partial(get_sources, dbs=dbs, docs_state0=docs_state0)
-
-        # if change collection source, must clear doc selections from it to avoid inconsistency
-        def clear_doc_choice():
-            return gr.Dropdown.update(choices=docs_state0, value=[docs_state0[0]])
-
-        langchain_mode.change(clear_doc_choice, inputs=None, outputs=document_choice)
-
-        def update_dropdown(x):
-            return gr.Dropdown.update(choices=x, value=[docs_state0[0]])
-
-        get_sources_btn.click(get_sources1, inputs=[my_db_state, langchain_mode], outputs=[file_source, docs_state],
-                              queue=queue,
-                              api_name='get_sources' if allow_api else None) \
-            .then(fn=update_dropdown, inputs=docs_state, outputs=document_choice)
-        # show button, else only show when add.  Could add to above get_sources for download/dropdown, but bit much maybe
-        show_sources1 = functools.partial(get_source_files_given_langchain_mode, dbs=dbs)
-        show_sources_btn.click(fn=show_sources1, inputs=[my_db_state, langchain_mode], outputs=sources_text,
-                               api_name='show_sources' if allow_api else None)
-
-        # Get inputs to evaluate() and make_db()
-        # don't deepcopy, can contain model itself
-        all_kwargs = kwargs.copy()
-        all_kwargs.update(locals())
-
-        refresh_sources1 = functools.partial(update_and_get_source_files_given_langchain_mode,
-                                             **get_kwargs(update_and_get_source_files_given_langchain_mode,
-                                                          exclude_names=['db1', 'langchain_mode'],
-                                                          **all_kwargs))
-        refresh_sources_btn.click(fn=refresh_sources1, inputs=[my_db_state, langchain_mode], outputs=sources_text,
-                                  api_name='refresh_sources' if allow_api else None)
-
-        def check_admin_pass(x):
-            return gr.update(visible=x == admin_pass)
-
-        def close_admin(x):
-            return gr.update(visible=not (x == admin_pass))
-
-        admin_btn.click(check_admin_pass, inputs=admin_pass_textbox, outputs=system_row, queue=False) \
-            .then(close_admin, inputs=admin_pass_textbox, outputs=admin_row, queue=False)
-
-        inputs_list = get_inputs_list(all_kwargs, kwargs['model_lower'])
-        from functools import partial
-        kwargs_evaluate = {k: v for k, v in all_kwargs.items() if k in inputs_kwargs_list}
-        # ensure present
-        for k in inputs_kwargs_list:
-            assert k in kwargs_evaluate, "Missing %s" % k
-        fun = partial(evaluate,
-                      **kwargs_evaluate)
-        fun2 = partial(evaluate,
-                       **kwargs_evaluate)
-
-        dark_mode_btn = gr.Button("Dark Mode", variant="primary").style(
-            size="sm",
-        )
-        # FIXME: Could add exceptions for non-chat but still streaming
-        exception_text = gr.Textbox(value="", visible=kwargs['chat'], label='Chat Exceptions', interactive=False)
-        dark_mode_btn.click(
-            None,
-            None,
-            None,
-            _js=get_dark_js(),
-            api_name="dark" if allow_api else None,
-            queue=False,
-        )
-
-        # Control chat and non-chat blocks, which can be independently used by chat checkbox swap
-        def col_nochat_fun(x):
-            return gr.Column.update(visible=not x)
-
-        def col_chat_fun(x):
-            return gr.Column.update(visible=x)
-
-        def context_fun(x):
-            return gr.Textbox.update(visible=not x)
-
-        chat.select(col_nochat_fun, chat, col_nochat, api_name="chat_checkbox" if allow_api else None) \
-            .then(col_chat_fun, chat, col_chat) \
-            .then(context_fun, chat, context) \
-            .then(col_chat_fun, chat, exception_text)
-
-        # examples after submit or any other buttons for chat or no chat
-        if kwargs['examples'] is not None and kwargs['show_examples']:
-            gr.Examples(examples=kwargs['examples'], inputs=inputs_list)
-
-        # Score
-        def score_last_response(*args, nochat=False, model2=False):
-            """ Similar to user() """
-            args_list = list(args)
-
-            if memory_restriction_level > 0:
-                max_length_tokenize = 768 - 256 if memory_restriction_level <= 2 else 512 - 256
-            else:
-                max_length_tokenize = 2048 - 256
-            cutoff_len = max_length_tokenize * 4  # restrict deberta related to max for LLM
-            smodel = score_model_state0[0]
-            stokenizer = score_model_state0[1]
-            sdevice = score_model_state0[2]
-            if not nochat:
-                history = args_list[-1]
-                if history is None:
-                    if not model2:
-                        # maybe only doing first model, no need to complain
-                        print("Bad history in scoring last response, fix for now", flush=True)
-                    history = []
-                if smodel is not None and \
-                        stokenizer is not None and \
-                        sdevice is not None and \
-                        history is not None and len(history) > 0 and \
-                        history[-1] is not None and \
-                        len(history[-1]) >= 2:
-                    os.environ['TOKENIZERS_PARALLELISM'] = 'false'
-
-                    question = history[-1][0]
-
-                    answer = history[-1][1]
-                else:
-                    return 'Response Score: NA'
-            else:
-                answer = args_list[-1]
-                instruction_nochat_arg_id = eval_func_param_names.index('instruction_nochat')
-                question = args_list[instruction_nochat_arg_id]
-
-            if question is None:
-                return 'Response Score: Bad Question'
-            if answer is None:
-                return 'Response Score: Bad Answer'
-            score = score_qa(smodel, stokenizer, max_length_tokenize, question, answer, cutoff_len)
-            if isinstance(score, str):
-                return 'Response Score: NA'
-            return 'Response Score: {:.1%}'.format(score)
-
-        def noop_score_last_response(*args, **kwargs):
-            return "Response Score: Disabled"
-
-        if kwargs['score_model']:
-            score_fun = score_last_response
-        else:
-            score_fun = noop_score_last_response
-
-        score_args = dict(fn=score_fun,
-                          inputs=inputs_list + [text_output],
-                          outputs=[score_text],
-                          )
-        score_args2 = dict(fn=partial(score_fun, model2=True),
-                           inputs=inputs_list + [text_output2],
-                           outputs=[score_text2],
-                           )
-
-        score_args_nochat = dict(fn=partial(score_fun, nochat=True),
-                                 inputs=inputs_list + [text_output_nochat],
-                                 outputs=[score_text_nochat],
-                                 )
-        if not kwargs['auto_score']:
-            score_event = score_btn.click(**score_args, queue=queue, api_name='score' if allow_api else None) \
-                .then(**score_args2, queue=queue, api_name='score2' if allow_api else None)
-            score_event_nochat = score_btn_nochat.click(**score_args_nochat, queue=queue,
-                                                        api_name='score_nochat' if allow_api else None)
-
-        def user(*args, undo=False, sanitize_user_prompt=True, model2=False):
-            """
-            User that fills history for bot
-            :param args:
-            :param undo:
-            :param sanitize_user_prompt:
-            :param model2:
-            :return:
-            """
-            args_list = list(args)
-            user_message = args_list[eval_func_param_names.index('instruction')]  # chat only
-            input1 = args_list[eval_func_param_names.index('iinput')]  # chat only
-            context1 = args_list[eval_func_param_names.index('context')]
-            prompt_type1 = args_list[eval_func_param_names.index('prompt_type')]
-            chat1 = args_list[eval_func_param_names.index('chat')]
-            stream_output1 = args_list[eval_func_param_names.index('stream_output')]
-            if input1 and not user_message.endswith(':'):
-                user_message1 = user_message + ":" + input1
-            elif input1:
-                user_message1 = user_message + input1
-            else:
-                user_message1 = user_message
-            if sanitize_user_prompt:
-                from better_profanity import profanity
-                user_message1 = profanity.censor(user_message1)
-            # FIXME: WIP to use desired seperator when user enters nothing
-            prompter = Prompter(prompt_type1, debug=kwargs['debug'], chat=chat1, stream_output=stream_output1)
-            if user_message1 in ['']:
-                # e.g. when user just hits enter in textbox,
-                # else will have <human>: <bot>: on single line, which seems to be "ok" for LLM but not usual
-                user_message1 = '\n'
-            # ensure good visually, else markdown ignores multiple \n
-            user_message1 = user_message1.replace('\n', '<br>')
-
-            history = args_list[-1]
-            if undo and history:
-                history.pop()
-            args_list = args_list[:-1]  # FYI, even if unused currently
-            if history is None:
-                if not model2:
-                    # no need to complain so often unless model1
-                    print("Bad history, fix for now", flush=True)
-                history = []
-            # ensure elements not mixed across models as output,
-            # even if input is currently same source
-            history = history.copy()
-            if undo:
-                return history
-            else:
-                # FIXME: compare, same history for now
-                return history + [[user_message1, None]]
-
-        def history_to_context(history, langchain_mode1, prompt_type1, chat1):
-            # ensure output will be unique to models
-            # FIXME: hard-coded 2048 implicitly passed:
-            _, _, _, max_prompt_length = get_cutoffs(memory_restriction_level, for_context=True)
-            history = copy.deepcopy(history)
-
-            context1 = ''
-            if max_prompt_length is not None and langchain_mode1 not in ['LLM']:
-                context1 = ''
-                # - 1 below because current instruction already in history from user()
-                for histi in range(0, len(history) - 1):
-                    data_point = dict(instruction=history[histi][0], input='', output=history[histi][1])
-                    prompt, pre_response, terminate_response, chat_sep = generate_prompt(data_point, prompt_type1,
-                                                                                         chat1, reduced=True)
-                    # md -> back to text, maybe not super important if model trained enough
-                    if not kwargs['keep_sources_in_context']:
-                        from gpt_langchain import source_prefix, source_postfix
-                        import re
-                        prompt = re.sub(f'{re.escape(source_prefix)}.*?{re.escape(source_postfix)}', '', prompt,
-                                        flags=re.DOTALL)
-                        if prompt.endswith('\n<p>'):
-                            prompt = prompt[:-4]
-                    prompt = prompt.replace('<br>', chat_sep)
-                    if not prompt.endswith(chat_sep):
-                        prompt += chat_sep
-                    # most recent first, add older if can
-                    # only include desired chat history
-                    if len(prompt + context1) > max_prompt_length:
-                        break
-                    context1 = prompt + context1
-
-                _, pre_response, terminate_response, chat_sep = generate_prompt({}, prompt_type1, chat1,
-                                                                                reduced=True)
-                if context1 and not context1.endswith(chat_sep):
-                    context1 += chat_sep  # ensure if terminates abruptly, then human continues on next line
-            return context1
-
-        def bot(*args, retry=False):
-            """
-            bot that consumes history for user input
-            instruction (from input_list) itself is not consumed by bot
-            :param args:
-            :param retry:
-            :return:
-            """
-            # don't deepcopy, can contain model itself
-            args_list = list(args).copy()
-            model_state1 = args_list[-3]
-            my_db_state1 = args_list[-2]
-            history = args_list[-1]
-
-            if model_state1[0] is None or model_state1[0] == no_model_str:
-                history = []
-                yield history, ''
-                return
-
-            args_list = args_list[:-3]  # only keep rest needed for evaluate()
-            langchain_mode1 = args_list[eval_func_param_names.index('langchain_mode')]
-            if retry and history:
-                history.pop()
-                if not args_list[eval_func_param_names.index('do_sample')]:
-                    # if was not sampling, no point in retry unless change to sample
-                    args_list[eval_func_param_names.index('do_sample')] = True
-            if not history:
-                print("No history", flush=True)
-                history = []
-                yield history, ''
-                return
-            instruction1 = history[-1][0]
-            if not instruction1:
-                # reject empty query, can sometimes go nuts
-                history = []
-                yield history, ''
-                return
-            prompt_type1 = args_list[eval_func_param_names.index('prompt_type')]
-            chat1 = args_list[eval_func_param_names.index('chat')]
-            context1 = history_to_context(history, langchain_mode1, prompt_type1, chat1)
-            args_list[0] = instruction1  # override original instruction with history from user
-            args_list[2] = context1
-            fun1 = partial(evaluate,
-                           model_state1,
-                           my_db_state1,
-                           **kwargs_evaluate)
-            try:
-                for output_fun in fun1(*tuple(args_list)):
-                    output = output_fun['response']
-                    extra = output_fun['sources']  # FIXME: can show sources in separate text box etc.
-                    # ensure good visually, else markdown ignores multiple \n
-                    bot_message = output.replace('\n', '<br>')
-                    history[-1][1] = bot_message
-                    yield history, ''
-            except StopIteration:
-                yield history, ''
-            except RuntimeError as e:
-                if "generator raised StopIteration" in str(e):
-                    # assume last entry was bad, undo
-                    history.pop()
-                    yield history, ''
-                else:
-                    if history and len(history) > 0 and len(history[0]) > 1 and history[-1][1] is None:
-                        history[-1][1] = ''
-                    yield history, str(e)
-                    raise
-            except Exception as e:
-                # put error into user input
-                ex = "Exception: %s" % str(e)
-                if history and len(history) > 0 and len(history[0]) > 1 and history[-1][1] is None:
-                    history[-1][1] = ''
-                yield history, ex
-                raise
-            return
-
-        # NORMAL MODEL
-        user_args = dict(fn=functools.partial(user, sanitize_user_prompt=kwargs['sanitize_user_prompt']),
-                         inputs=inputs_list + [text_output],
-                         outputs=text_output,
-                         )
-        bot_args = dict(fn=bot,
-                        inputs=inputs_list + [model_state, my_db_state] + [text_output],
-                        outputs=[text_output, exception_text],
-                        )
-        retry_bot_args = dict(fn=functools.partial(bot, retry=True),
-                              inputs=inputs_list + [model_state, my_db_state] + [text_output],
-                              outputs=[text_output, exception_text],
-                              )
-        undo_user_args = dict(fn=functools.partial(user, undo=True),
-                              inputs=inputs_list + [text_output],
-                              outputs=text_output,
-                              )
-
-        # MODEL2
-        user_args2 = dict(fn=functools.partial(user, sanitize_user_prompt=kwargs['sanitize_user_prompt'], model2=True),
-                          inputs=inputs_list + [text_output2],
-                          outputs=text_output2,
-                          )
-        bot_args2 = dict(fn=bot,
-                         inputs=inputs_list + [model_state2, my_db_state] + [text_output2],
-                         outputs=[text_output2, exception_text],
-                         )
-        retry_bot_args2 = dict(fn=functools.partial(bot, retry=True),
-                               inputs=inputs_list + [model_state2, my_db_state] + [text_output2],
-                               outputs=[text_output2, exception_text],
-                               )
-        undo_user_args2 = dict(fn=functools.partial(user, undo=True),
-                               inputs=inputs_list + [text_output2],
-                               outputs=text_output2,
-                               )
-
-        def clear_instruct():
-            return gr.Textbox.update(value='')
-
-        if kwargs['auto_score']:
-            score_args_submit = score_args
-            score_args2_submit = score_args2
-        else:
-            score_args_submit = dict(fn=lambda: None, inputs=None, outputs=None)
-            score_args2_submit = dict(fn=lambda: None, inputs=None, outputs=None)
-
-        # in case 2nd model, consume instruction first, so can clear quickly
-        # bot doesn't consume instruction itself, just history from user, so why works
-        submit_event1a = instruction.submit(**user_args, queue=queue,
-                                            api_name='instruction' if allow_api else None)
-        submit_event1b = submit_event1a.then(**user_args2, api_name='instruction2' if allow_api else None)
-        submit_event1c = submit_event1b.then(clear_instruct, None, instruction) \
-            .then(clear_instruct, None, iinput)
-        submit_event1d = submit_event1c.then(**bot_args, api_name='instruction_bot' if allow_api else None,
-                                             queue=queue)
-        submit_event1e = submit_event1d.then(**score_args_submit,
-                                             api_name='instruction_bot_score' if allow_api else None,
-                                             queue=queue)
-        submit_event1f = submit_event1e.then(**bot_args2, api_name='instruction_bot2' if allow_api else None,
-                                             queue=queue)
-        submit_event1g = submit_event1f.then(**score_args2_submit,
-                                             api_name='instruction_bot_score2' if allow_api else None, queue=queue)
-        submit_event1h = submit_event1g.then(clear_torch_cache)
-
-        submit_event2a = submit.click(**user_args, api_name='submit' if allow_api else None)
-        submit_event2b = submit_event2a.then(**user_args2, api_name='submit2' if allow_api else None)
-        submit_event2c = submit_event2b.then(clear_instruct, None, instruction) \
-            .then(clear_instruct, None, iinput)
-        submit_event2d = submit_event2c.then(**bot_args, api_name='submit_bot' if allow_api else None, queue=queue)
-        submit_event2e = submit_event2d.then(**score_args_submit, api_name='submit_bot_score' if allow_api else None,
-                                             queue=queue)
-        submit_event2f = submit_event2e.then(**bot_args2, api_name='submit_bot2' if allow_api else None, queue=queue)
-        submit_event2g = submit_event2f.then(**score_args2_submit, api_name='submit_bot_score2' if allow_api else None,
-                                             queue=queue)
-        submit_event2h = submit_event2g.then(clear_torch_cache)
-
-        submit_event3a = retry.click(**user_args, api_name='retry' if allow_api else None)
-        submit_event3b = submit_event3a.then(**user_args2, api_name='retry2' if allow_api else None)
-        submit_event3c = submit_event3b.then(clear_instruct, None, instruction) \
-            .then(clear_instruct, None, iinput)
-        submit_event3d = submit_event3c.then(**retry_bot_args, api_name='retry_bot' if allow_api else None,
-                                             queue=queue)
-        submit_event3e = submit_event3d.then(**score_args_submit, api_name='retry_bot_score' if allow_api else None,
-                                             queue=queue)
-        submit_event3f = submit_event3e.then(**retry_bot_args2, api_name='retry_bot2' if allow_api else None,
-                                             queue=queue)
-        submit_event3g = submit_event3f.then(**score_args2_submit, api_name='retry_bot_score2' if allow_api else None,
-                                             queue=queue)
-        submit_event3h = submit_event3g.then(clear_torch_cache)
-
-        submit_event4 = undo.click(**undo_user_args, api_name='undo' if allow_api else None) \
-            .then(**undo_user_args2, api_name='undo2' if allow_api else None) \
-            .then(clear_instruct, None, instruction) \
-            .then(clear_instruct, None, iinput) \
-            .then(**score_args_submit, api_name='undo_score' if allow_api else None) \
-            .then(**score_args2_submit, api_name='undo_score2' if allow_api else None)
-
-        # MANAGE CHATS
-        def dedup(short_chat, short_chats):
-            if short_chat not in short_chats:
-                return short_chat
-            for i in range(1, 1000):
-                short_chat_try = short_chat + "_" + str(i)
-                if short_chat_try not in short_chats:
-                    return short_chat_try
-            # fallback and hope for best
-            short_chat = short_chat + "_" + str(random.random())
-            return short_chat
-
-        def get_short_chat(x, short_chats, short_len=20, words=4):
-            if x and len(x[0]) == 2 and x[0][0] is not None:
-                short_chat = ' '.join(x[0][0][:short_len].split(' ')[:words]).strip()
-                short_chat = dedup(short_chat, short_chats)
-            else:
-                short_chat = None
-            return short_chat
-
-        def is_chat_same(x, y):
-            # <p> etc. added in chat, try to remove some of that to help avoid dup entries when hit new conversation
-            is_same = True
-            # length of conversation has to be same
-            if len(x) != len(y):
-                return False
-            for stepx, stepy in zip(x, y):
-                if len(stepx) != len(stepy):
-                    # something off with a conversation
-                    return False
-                if len(stepx) != 2:
-                    # something off
-                    return False
-                if len(stepy) != 2:
-                    # something off
-                    return False
-                questionx = stepx[0].replace('<p>', '').replace('</p>', '') if stepx[0] is not None else None
-                answerx = stepx[1].replace('<p>', '').replace('</p>', '') if stepx[1] is not None else None
-
-                questiony = stepy[0].replace('<p>', '').replace('</p>', '') if stepy[0] is not None else None
-                answery = stepy[1].replace('<p>', '').replace('</p>', '') if stepy[1] is not None else None
-
-                if questionx != questiony or answerx != answery:
-                    return False
-            return is_same
-
-        def save_chat(chat1, chat2, chat_state1):
-            short_chats = list(chat_state1.keys())
-            for chati in [chat1, chat2]:
-                if chati and len(chati) > 0 and len(chati[0]) == 2 and chati[0][1] is not None:
-                    short_chat = get_short_chat(chati, short_chats)
-                    if short_chat:
-                        already_exists = any([is_chat_same(chati, x) for x in chat_state1.values()])
-                        if not already_exists:
-                            chat_state1[short_chat] = chati
-            return chat_state1
-
-        def update_radio_chats(chat_state1):
-            return gr.update(choices=list(chat_state1.keys()), value=None)
-
-        def deselect_radio_chats():
-            return gr.update(value=None)
-
-        def switch_chat(chat_key, chat_state1):
-            chosen_chat = chat_state1[chat_key]
-            return chosen_chat, chosen_chat
-
-        radio_chats.input(switch_chat, inputs=[radio_chats, chat_state], outputs=[text_output, text_output2])
-
-        def remove_chat(chat_key, chat_state1):
-            chat_state1.pop(chat_key, None)
-            return chat_state1
-
-        remove_chat_btn.click(remove_chat, inputs=[radio_chats, chat_state], outputs=chat_state) \
-            .then(update_radio_chats, inputs=chat_state, outputs=radio_chats)
-
-        def get_chats1(chat_state1):
-            base = 'chats'
-            makedirs(base, exist_ok=True)
-            filename = os.path.join(base, 'chats_%s.json' % str(uuid.uuid4()))
-            with open(filename, "wt") as f:
-                f.write(json.dumps(chat_state1, indent=2))
-            return filename
-
-        export_chats_btn.click(get_chats1, inputs=chat_state, outputs=chats_file, queue=False,
-                               api_name='export_chats' if allow_api else None)
-
-        def add_chats_from_file(file, chat_state1, add_btn):
-            if not file:
-                return chat_state1, add_btn
-            if isinstance(file, str):
-                files = [file]
-            else:
-                files = file
-            if not files:
-                return chat_state1, add_btn
-            for file1 in files:
-                try:
-                    if hasattr(file1, 'name'):
-                        file1 = file1.name
-                    with open(file1, "rt") as f:
-                        new_chats = json.loads(f.read())
-                        for chat1_k, chat1_v in new_chats.items():
-                            # ignore chat1_k, regenerate and de-dup to avoid loss
-                            chat_state1 = save_chat(chat1_v, None, chat_state1)
-                except BaseException as e:
-                    print("Add chats exception: %s" % str(e), flush=True)
-            return chat_state1, add_btn
-
-        # note for update_user_db_func output is ignored for db
-        add_to_chats_btn.click(add_chats_from_file,
-                               inputs=[chatsup_output, chat_state, add_to_chats_btn],
-                               outputs=[chat_state, add_to_my_db_btn], queue=False,
-                               api_name='add_to_chats' if allow_api else None) \
-            .then(clear_file_list, outputs=chatsup_output, queue=False) \
-            .then(update_radio_chats, inputs=chat_state, outputs=radio_chats, queue=False)
-
-        clear_chat_btn.click(lambda: None, None, text_output, queue=False, api_name='clear' if allow_api else None) \
-            .then(lambda: None, None, text_output2, queue=False, api_name='clear2' if allow_api else None) \
-            .then(deselect_radio_chats, inputs=None, outputs=radio_chats, queue=False)
-
-        # does both models
-        clear.click(save_chat, inputs=[text_output, text_output2, chat_state], outputs=chat_state,
-                    api_name='save_chat' if allow_api else None) \
-            .then(update_radio_chats, inputs=chat_state, outputs=radio_chats,
-                  api_name='update_chats' if allow_api else None) \
-            .then(lambda: None, None, text_output, queue=False, api_name='clearB' if allow_api else None) \
-            .then(lambda: None, None, text_output2, queue=False, api_name='clearB2' if allow_api else None)
-        # NOTE: clear of instruction/iinput for nochat has to come after score,
-        # because score for nochat consumes actual textbox, while chat consumes chat history filled by user()
-        submit_event_nochat = submit_nochat.click(fun,
-                                                  inputs=[model_state, my_db_state] + inputs_list,
-                                                  outputs=text_output_nochat,
-                                                  queue=queue,
-                                                  api_name='submit_nochat' if allow_api else None) \
-            .then(**score_args_nochat, api_name='instruction_bot_score_nochat' if allow_api else None, queue=queue) \
-            .then(clear_instruct, None, instruction_nochat) \
-            .then(clear_instruct, None, iinput_nochat) \
-            .then(clear_torch_cache)
-
-        def load_model(model_name, lora_weights, model_state_old, prompt_type_old, load_8bit, infer_devices, gpu_id):
-            # ensure old model removed from GPU memory
-            if kwargs['debug']:
-                print("Pre-switch pre-del GPU memory: %s" % get_torch_allocated(), flush=True)
-
-            model0 = model_state0[0]
-            if isinstance(model_state_old[0], str) and model0 is not None:
-                # best can do, move model loaded at first to CPU
-                model0.cpu()
-
-            if model_state_old[0] is not None and not isinstance(model_state_old[0], str):
-                try:
-                    model_state_old[0].cpu()
-                except Exception as e:
-                    # sometimes hit NotImplementedError: Cannot copy out of meta tensor; no data!
-                    print("Unable to put model on CPU: %s" % str(e), flush=True)
-                del model_state_old[0]
-                model_state_old[0] = None
-
-            if model_state_old[1] is not None and not isinstance(model_state_old[1], str):
-                del model_state_old[1]
-                model_state_old[1] = None
-
-            clear_torch_cache()
-            if kwargs['debug']:
-                print("Pre-switch post-del GPU memory: %s" % get_torch_allocated(), flush=True)
-
-            if model_name is None or model_name == no_model_str:
-                # no-op if no model, just free memory
-                # no detranscribe needed for model, never go into evaluate
-                lora_weights = no_lora_str
-                return [None, None, None, model_name], model_name, lora_weights, prompt_type_old
-
-            # don't deepcopy, can contain model itself
-            all_kwargs1 = all_kwargs.copy()
-            all_kwargs1['base_model'] = model_name.strip()
-            all_kwargs1['load_8bit'] = load_8bit
-            all_kwargs1['infer_devices'] = infer_devices
-            all_kwargs1['gpu_id'] = int(gpu_id)  # detranscribe
-            model_lower = model_name.strip().lower()
-            if model_lower in inv_prompt_type_to_model_lower:
-                prompt_type1 = inv_prompt_type_to_model_lower[model_lower]
-            else:
-                prompt_type1 = prompt_type_old
-
-            # detranscribe
-            if lora_weights == no_lora_str:
-                lora_weights = ''
-
-            all_kwargs1['lora_weights'] = lora_weights.strip()
-            model1, tokenizer1, device1 = get_model(reward_type=False,
-                                                    **get_kwargs(get_model, exclude_names=['reward_type'],
-                                                                 **all_kwargs1))
-            clear_torch_cache()
-
-            if kwargs['debug']:
-                print("Post-switch GPU memory: %s" % get_torch_allocated(), flush=True)
-            return [model1, tokenizer1, device1, model_name], model_name, lora_weights, prompt_type1
-
-        def dropdown_prompt_type_list(x):
-            return gr.Dropdown.update(value=x)
-
-        def chatbot_list(x, model_used_in):
-            return gr.Textbox.update(label=f'h2oGPT [Model: {model_used_in}]')
-
-        load_model_args = dict(fn=load_model,
-                               inputs=[model_choice, lora_choice, model_state, prompt_type,
-                                       model_load8bit_checkbox, model_infer_devices_checkbox, model_gpu],
-                               outputs=[model_state, model_used, lora_used, prompt_type])
-        prompt_update_args = dict(fn=dropdown_prompt_type_list, inputs=prompt_type, outputs=prompt_type)
-        chatbot_update_args = dict(fn=chatbot_list, inputs=[text_output, model_used], outputs=text_output)
-        nochat_update_args = dict(fn=chatbot_list, inputs=[text_output_nochat, model_used], outputs=text_output_nochat)
-        if not is_public:
-            load_model_event = load_model_button.click(**load_model_args, api_name='load_model' if allow_api else None) \
-                .then(**prompt_update_args) \
-                .then(**chatbot_update_args) \
-                .then(**nochat_update_args) \
-                .then(clear_torch_cache)
-
-        load_model_args2 = dict(fn=load_model,
-                                inputs=[model_choice2, lora_choice2, model_state2, prompt_type2,
-                                        model_load8bit_checkbox2, model_infer_devices_checkbox2, model_gpu2],
-                                outputs=[model_state2, model_used2, lora_used2, prompt_type2])
-        prompt_update_args2 = dict(fn=dropdown_prompt_type_list, inputs=prompt_type2, outputs=prompt_type2)
-        chatbot_update_args2 = dict(fn=chatbot_list, inputs=[text_output2, model_used2], outputs=text_output2)
-        if not is_public:
-            load_model_event2 = load_model_button2.click(**load_model_args2,
-                                                         api_name='load_model2' if allow_api else None) \
-                .then(**prompt_update_args2) \
-                .then(**chatbot_update_args2) \
-                .then(clear_torch_cache)
-
-        def dropdown_model_list(list0, x):
-            new_state = [list0[0] + [x]]
-            new_options = [*new_state[0]]
-            return gr.Dropdown.update(value=x, choices=new_options), \
-                gr.Dropdown.update(value=x, choices=new_options), \
-                '', new_state
-
-        add_model_event = add_model_button.click(fn=dropdown_model_list,
-                                                 inputs=[model_options_state, new_model],
-                                                 outputs=[model_choice, model_choice2, new_model, model_options_state],
-                                                 queue=False)
-
-        def dropdown_lora_list(list0, x, model_used1, lora_used1, model_used2, lora_used2):
-            new_state = [list0[0] + [x]]
-            new_options = [*new_state[0]]
-            # don't switch drop-down to added lora if already have model loaded
-            x1 = x if model_used1 == no_model_str else lora_used1
-            x2 = x if model_used2 == no_model_str else lora_used2
-            return gr.Dropdown.update(value=x1, choices=new_options), \
-                gr.Dropdown.update(value=x2, choices=new_options), \
-                '', new_state
-
-        add_lora_event = add_lora_button.click(fn=dropdown_lora_list,
-                                               inputs=[lora_options_state, new_lora, model_used, lora_used, model_used2,
-                                                       lora_used2],
-                                               outputs=[lora_choice, lora_choice2, new_lora, lora_options_state],
-                                               queue=False)
-
-        go_btn.click(lambda: gr.update(visible=False), None, go_btn, api_name="go" if allow_api else None, queue=False) \
-            .then(lambda: gr.update(visible=True), None, normal_block, queue=False) \
-            .then(**load_model_args, queue=False).then(**prompt_update_args, queue=False)
-
-        def compare_textbox_fun(x):
-            return gr.Textbox.update(visible=x)
-
-        def compare_column_fun(x):
-            return gr.Column.update(visible=x)
-
-        def compare_prompt_fun(x):
-            return gr.Dropdown.update(visible=x)
-
-        compare_checkbox.select(compare_textbox_fun, compare_checkbox, text_output2,
-                                api_name="compare_checkbox" if allow_api else None) \
-            .then(compare_column_fun, compare_checkbox, col_model2) \
-            .then(compare_prompt_fun, compare_checkbox, prompt_type2) \
-            .then(compare_textbox_fun, compare_checkbox, score_text2)
-        # FIXME: add score_res2 in condition, but do better
-
-        # callback for logging flagged input/output
-        callback.setup(inputs_list + [text_output, text_output2], "flagged_data_points")
-        flag_btn.click(lambda *args: callback.flag(args), inputs_list + [text_output, text_output2], None,
-                       preprocess=False,
-                       api_name='flag' if allow_api else None, queue=False)
-        flag_btn_nochat.click(lambda *args: callback.flag(args), inputs_list + [text_output_nochat], None,
-                              preprocess=False,
-                              api_name='flag_nochat' if allow_api else None, queue=False)
-
-        def get_system_info():
-            return gr.Textbox.update(value=system_info_print())
-
-        system_event = system_btn.click(get_system_info, outputs=system_text,
-                                        api_name='system_info' if allow_api else None, queue=False)
-
-        # don't pass text_output, don't want to clear output, just stop it
-        # cancel only stops outer generation, not inner generation or non-generation
-        stop_btn.click(lambda: None, None, None,
-                       cancels=[submit_event1d, submit_event1f,
-                                submit_event2d, submit_event2f,
-                                submit_event3d, submit_event3f,
-                                submit_event_nochat],
-                       queue=False, api_name='stop' if allow_api else None).then(clear_torch_cache, queue=False)
-
-        def count_chat_tokens(model_state1, chat1, prompt_type1):
-            if model_state1 and not isinstance(model_state1[1], str):
-                tokenizer = model_state1[1]
-            elif model_state0 and not isinstance(model_state0[1], str):
-                tokenizer = model_state0[1]
-            else:
-                tokenizer = None
-            if tokenizer is not None:
-                langchain_mode1 = 'ChatLLM'
-                # fake user message to mimic bot()
-                chat1 = copy.deepcopy(chat1)
-                chat1 = chat1 + [['user_message1', None]]
-                context1 = history_to_context(chat1, langchain_mode1, prompt_type1, chat1)
-                return str(tokenizer(context1, return_tensors="pt")['input_ids'].shape[1])
-            else:
-                return "N/A"
-
-        count_chat_tokens_btn.click(fn=count_chat_tokens, inputs=[model_state, text_output, prompt_type],
-                                    outputs=chat_token_count, api_name='count_tokens' if allow_api else None)
-
-        demo.load(None, None, None, _js=get_dark_js() if kwargs['h2ocolors'] else None)
-
-    demo.queue(concurrency_count=kwargs['concurrency_count'], api_open=kwargs['api_open'])
-    favicon_path = "h2o-logo.svg"
-
-    scheduler = BackgroundScheduler()
-    scheduler.add_job(func=clear_torch_cache, trigger="interval", seconds=20)
-    if is_public and \
-            kwargs['base_model'] not in non_hf_types:
-        # FIXME: disable for gptj, langchain or gpt4all modify print itself
-        # FIXME: and any multi-threaded/async print will enter model output!
-        scheduler.add_job(func=ping, trigger="interval", seconds=60)
-    scheduler.start()
-
-    # import control
-    if kwargs['langchain_mode'] == 'Disabled' and \
-            os.environ.get("TEST_LANGCHAIN_IMPORT") and \
-            kwargs['base_model'] not in non_hf_types:
-        assert 'gpt_langchain' not in sys.modules, "Dev bug, import of langchain when should not have"
-        assert 'langchain' not in sys.modules, "Dev bug, import of langchain when should not have"
-
-    demo.launch(share=kwargs['share'], server_name="0.0.0.0", show_error=True,
-                favicon_path=favicon_path, prevent_thread_lock=True,
-                auth=kwargs['auth'])
-    if kwargs['verbose']:
-        print("Started GUI", flush=True)
-    if kwargs['block_gradio_exit']:
-        demo.block_thread()
-
-
-input_args_list = ['model_state', 'my_db_state']
-
-
-def get_inputs_list(inputs_dict, model_lower):
-    """
-    map gradio objects in locals() to inputs for evaluate().
-    :param inputs_dict:
-    :param model_lower:
-    :return:
-    """
-    inputs_list_names = list(inspect.signature(evaluate).parameters)
-    inputs_list = []
-    for k in inputs_list_names:
-        if k == 'kwargs':
-            continue
-        if k in input_args_list + inputs_kwargs_list:
-            # these are added at use time for args or partial for kwargs, not taken as input
-            continue
-        if 'mbart-' not in model_lower and k in ['src_lang', 'tgt_lang']:
-            continue
-        inputs_list.append(inputs_dict[k])
-    return inputs_list
-
-
-def get_sources(db1, langchain_mode, dbs=None, docs_state0=None):
-    if langchain_mode in ['ChatLLM', 'LLM']:
-        source_files_added = "NA"
-        source_list = []
-    elif langchain_mode in ['wiki_full']:
-        source_files_added = "Not showing wiki_full, takes about 20 seconds and makes 4MB file." \
-                             "  Ask jon.mckinney@h2o.ai for file if required."
-        source_list = []
-    elif langchain_mode == 'MyData' and len(db1) > 0 and db1[0] is not None:
-        db_get = db1[0].get()
-        source_list = sorted(set([x['source'] for x in db_get['metadatas']]))
-        source_files_added = '\n'.join(source_list)
-    elif langchain_mode in dbs and dbs[langchain_mode] is not None:
-        db1 = dbs[langchain_mode]
-        db_get = db1.get()
-        source_list = sorted(set([x['source'] for x in db_get['metadatas']]))
-        source_files_added = '\n'.join(source_list)
-    else:
-        source_list = []
-        source_files_added = "None"
-    sources_file = 'sources_%s_%s' % (langchain_mode, str(uuid.uuid4()))
-    with open(sources_file, "wt") as f:
-        f.write(source_files_added)
-    source_list = docs_state0 + source_list
-    return sources_file, source_list
-
-
-def update_user_db(file, db1, x, y, *args, dbs=None, langchain_mode='UserData', **kwargs):
-    try:
-        return _update_user_db(file, db1, x, y, *args, dbs=dbs, langchain_mode=langchain_mode, **kwargs)
-    except BaseException as e:
-        print(traceback.format_exc(), flush=True)
-        # gradio has issues if except, so fail semi-gracefully, else would hang forever in processing textbox
-        ex_str = "Exception: %s" % str(e)
-        source_files_added = """\
-        <html>
-          <body>
-            <p>
-               Sources: <br>
-            </p>
-               <div style="overflow-y: auto;height:400px">
-               {0}
-               </div>
-          </body>
-        </html>
-        """.format(ex_str)
-        if langchain_mode == 'MyData':
-            return db1, x, y, source_files_added
-        else:
-            return x, y, source_files_added
-
-
-def _update_user_db(file, db1, x, y, dbs=None, db_type=None, langchain_mode='UserData', use_openai_embedding=False,
-                    hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
-                    caption_loader=None,
-                    enable_captions=True,
-                    captions_model="Salesforce/blip-image-captioning-base",
-                    enable_ocr=False,
-                    verbose=False,
-                    chunk=True, chunk_size=512, is_url=False, is_txt=False):
-    assert isinstance(dbs, dict), "Wrong type for dbs: %s" % str(type(dbs))
-    assert db_type in ['faiss', 'chroma'], "db_type %s not supported" % db_type
-    from gpt_langchain import add_to_db, get_db, path_to_docs
-    # handle case of list of temp buffer
-    if isinstance(file, list) and len(file) > 0 and hasattr(file[0], 'name'):
-        file = [x.name for x in file]
-    # handle single file of temp buffer
-    if hasattr(file, 'name'):
-        file = file.name
-    if verbose:
-        print("Adding %s" % file, flush=True)
-    sources = path_to_docs(file if not is_url and not is_txt else None,
-                           verbose=verbose, chunk=chunk, chunk_size=chunk_size,
-                           url=file if is_url else None,
-                           text=file if is_txt else None,
-                           enable_captions=enable_captions,
-                           captions_model=captions_model,
-                           enable_ocr=enable_ocr,
-                           caption_loader=caption_loader,
-                           )
-    exceptions = [x for x in sources if x.metadata.get('exception')]
-    sources = [x for x in sources if 'exception' not in x.metadata]
-
-    with filelock.FileLock("db_%s.lock" % langchain_mode.replace(' ', '_')):
-        if langchain_mode == 'MyData':
-            if db1[0] is not None:
-                # then add
-                db, num_new_sources, new_sources_metadata = add_to_db(db1[0], sources, db_type=db_type)
-            else:
-                assert len(db1) == 2 and db1[1] is None, "Bad MyData db: %s" % db1
-                # then create
-                # assign fresh hash for this user session, so not shared
-                # if added has to original state and didn't change, then would be shared db for all users
-                db1[1] = str(uuid.uuid4())
-                persist_directory = os.path.join(scratch_base_dir, 'db_dir_%s_%s' % (langchain_mode, db1[1]))
-                db1[0] = get_db(sources, use_openai_embedding=use_openai_embedding,
-                                db_type=db_type,
-                                persist_directory=persist_directory,
-                                langchain_mode=langchain_mode,
-                                hf_embedding_model=hf_embedding_model)
-                if db1[0] is None:
-                    db1[1] = None
-            source_files_added = get_source_files(db=db1[0], exceptions=exceptions)
-            return db1, x, y, source_files_added
-        else:
-            persist_directory = 'db_dir_%s' % langchain_mode
-            if langchain_mode in dbs and dbs[langchain_mode] is not None:
-                # then add
-                db, num_new_sources, new_sources_metadata = add_to_db(dbs[langchain_mode], sources, db_type=db_type)
-            else:
-                # then create
-                db = get_db(sources, use_openai_embedding=use_openai_embedding,
-                            db_type=db_type,
-                            persist_directory=persist_directory,
-                            langchain_mode=langchain_mode,
-                            hf_embedding_model=hf_embedding_model)
-                dbs[langchain_mode] = db
-            # NOTE we do not return db, because function call always same code path
-            # return dbs[langchain_mode], x, y
-            # db in this code path is updated in place
-            source_files_added = get_source_files(db=dbs[langchain_mode], exceptions=exceptions)
-            return x, y, source_files_added
-
-
-def get_db(db1, langchain_mode, dbs=None):
-    with filelock.FileLock("db_%s.lock" % langchain_mode.replace(' ', '_')):
-        if langchain_mode in ['wiki_full']:
-            # NOTE: avoid showing full wiki.  Takes about 30 seconds over about 90k entries, but not useful for now
-            db = None
-        elif langchain_mode == 'MyData' and len(db1) > 0 and db1[0] is not None:
-            db = db1[0]
-        elif langchain_mode in dbs and dbs[langchain_mode] is not None:
-            db = dbs[langchain_mode]
-        else:
-            db = None
-    return db
-
-
-def get_source_files_given_langchain_mode(db1, langchain_mode='UserData', dbs=None):
-    db = get_db(db1, langchain_mode, dbs=dbs)
-    return get_source_files(db=db, exceptions=None)
-
-
-def get_source_files(db=None, exceptions=None, metadatas=None):
-    if exceptions is None:
-        exceptions = []
-
-    # only should be one source, not confused
-    assert db is not None or metadatas is not None
-
-    if metadatas is None:
-        source_label = "Sources:"
-        if db is not None:
-            metadatas = db.get()['metadatas']
-        else:
-            metadatas = []
-        adding_new = False
-    else:
-        source_label = "New Sources:"
-        adding_new = True
-
-    # below automatically de-dups
-    from gpt_langchain import get_url
-    small_dict = {get_url(x['source'], from_str=True, short_name=True): get_short_name(x.get('head')) for x in
-                  metadatas}
-    # if small_dict is empty dict, that's ok
-    df = pd.DataFrame(small_dict.items(), columns=['source', 'head'])
-    df.index = df.index + 1
-    df.index.name = 'index'
-    source_files_added = tabulate.tabulate(df, headers='keys', tablefmt='unsafehtml')
-
-    if exceptions:
-        exception_metadatas = [x.metadata for x in exceptions]
-        small_dict = {get_url(x['source'], from_str=True, short_name=True): get_short_name(x.get('exception')) for x in
-                      exception_metadatas}
-        # if small_dict is empty dict, that's ok
-        df = pd.DataFrame(small_dict.items(), columns=['source', 'exception'])
-        df.index = df.index + 1
-        df.index.name = 'index'
-        exceptions_html = tabulate.tabulate(df, headers='keys', tablefmt='unsafehtml')
-    else:
-        exceptions_html = ''
-
-    if metadatas and exceptions:
-        source_files_added = """\
-        <html>
-          <body>
-            <p>
-               {0} <br>
-            </p>
-               <div style="overflow-y: auto;height:400px">
-               {1}
-               {2}
-               </div>
-          </body>
-        </html>
-        """.format(source_label, source_files_added, exceptions_html)
-    elif metadatas:
-        source_files_added = """\
-        <html>
-          <body>
-            <p>
-               {0} <br>
-            </p>
-               <div style="overflow-y: auto;height:400px">
-               {1}
-               </div>
-          </body>
-        </html>
-        """.format(source_label, source_files_added)
-    elif exceptions_html:
-        source_files_added = """\
-        <html>
-          <body>
-            <p>
-               Exceptions: <br>
-            </p>
-               <div style="overflow-y: auto;height:400px">
-               {0}
-               </div>
-          </body>
-        </html>
-        """.format(exceptions_html)
-    else:
-        if adding_new:
-            source_files_added = "No New Sources"
-        else:
-            source_files_added = "No Sources"
-
-    return source_files_added
-
-
-def update_and_get_source_files_given_langchain_mode(db1, langchain_mode, dbs=None, first_para=None,
-                                                     text_limit=None, chunk=None, chunk_size=None,
-                                                     user_path=None, db_type=None, load_db_if_exists=None,
-                                                     n_jobs=None, verbose=None):
-    db = get_db(db1, langchain_mode, dbs=dbs)
-
-    from gpt_langchain import make_db
-    db, num_new_sources, new_sources_metadata = make_db(use_openai_embedding=False,
-                                                        hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
-                                                        first_para=first_para, text_limit=text_limit, chunk=chunk,
-                                                        chunk_size=chunk_size,
-                                                        langchain_mode=langchain_mode,
-                                                        user_path=user_path,
-                                                        db_type=db_type,
-                                                        load_db_if_exists=load_db_if_exists,
-                                                        db=db,
-                                                        n_jobs=n_jobs,
-                                                        verbose=verbose)
-    # return only new sources with text saying such
-    return get_source_files(db=None, exceptions=None, metadatas=new_sources_metadata)

gradio_runner.py

@@ -0,0 +1 @@
+../../gradio_runner.py

gradio_themes.py

@@ -1,183 +0,0 @@
-from __future__ import annotations
-
-from typing import Iterable
-
-from gradio.themes.soft import Soft
-from gradio.themes import Color
-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",
-)
-
-
-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',
-        )
-
-
-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,
-        )
-
-
-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):
-    return f"""<div style="display:flex; justify-content:center; margin-bottom:30px;">
-                    <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></img>
-                </div>
-                """
-
-
-def get_simple_title(title):
-    return f"""<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');
-        }
-    }"""

gradio_themes.py

@@ -0,0 +1 @@
+../../gradio_themes.py

gradio_ui

@@ -0,0 +1 @@
+../../gradio_ui

h2oai_pipeline.py

@@ -1,128 +0,0 @@
-from transformers import TextGenerationPipeline
-from transformers.pipelines.text_generation import ReturnType
-
-from stopping import get_stopping
-from prompter import Prompter
-
-
-class H2OTextGenerationPipeline(TextGenerationPipeline):
-    def __init__(self, *args, debug=False, chat=False, stream_output=False,
-                 sanitize_bot_response=True,
-                 use_prompter=True, prompter=None, prompt_type=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 max_input_tokens:
-        :param kwargs:
-        """
-        super().__init__(*args, **kwargs)
-        self.prompt_text = None
-        self.use_prompter = use_prompter
-        self.prompt_type = prompt_type
-        self.prompter = prompter
-        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, 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
-
-    def preprocess(self, prompt_text, prefix="", handle_long_generation=None, **generate_kwargs):
-        data_point = dict(context='', instruction=prompt_text, input='')
-        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 = 'hole'
-        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].strip().split(self.human)[0].strip()
-            else:
-                outputs = rec['generated_text']
-            rec['generated_text'] = outputs
-        return records
-
-    def _forward(self, model_inputs, **generate_kwargs):
-        if self.can_stop:
-            stopping_criteria = get_stopping(self.prompt_type, self.tokenizer, self.device, human=self.human,
-                                             bot=self.bot)
-            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}

h2oai_pipeline.py

@@ -0,0 +1 @@
+../../h2oai_pipeline.py

loaders.py

@@ -1,50 +0,0 @@
-def get_loaders(llama_type, model_name, reward_type):
-    # 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 llama_type:
-        from transformers import LlamaForCausalLM, LlamaTokenizer
-        model_loader = LlamaForCausalLM
-        tokenizer_loader = LlamaTokenizer
-    elif 'distilgpt2' in model_name.lower():
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        return AutoModelForCausalLM, AutoTokenizer
-    elif 'gpt2' in model_name.lower():
-        from transformers import GPT2LMHeadModel, GPT2Tokenizer
-        return GPT2LMHeadModel, GPT2Tokenizer
-    elif 'mbart-' in model_name.lower():
-        from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-        return MBartForConditionalGeneration, 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, AutoTokenizer
-    elif 'bigbird' in model_name:
-        from transformers import BigBirdPegasusForConditionalGeneration, AutoTokenizer
-        return BigBirdPegasusForConditionalGeneration, 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, AutoTokenizer
-    else:
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-        model_loader = AutoModelForCausalLM
-        tokenizer_loader = AutoTokenizer
-    return model_loader, 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)
-
-    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

loaders.py

@@ -0,0 +1 @@
+../../loaders.py

prompter.py

@@ -1,576 +0,0 @@
-import time
-from enum import Enum
-
-non_hf_types = ['gpt4all_llama', 'llama', 'gptj']
-
-
-class PromptType(Enum):
-    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
-
-
-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',
-        'mosaicml/mpt-7b-instruct',  # internal code handles instruct
-        'mosaicml/mpt-7b-chat',  # NC, internal code handles instruct
-        'gptj',  # internally handles prompting
-        'llama',  # plain, or need to choose prompt_type for given TheBloke model
-        'gpt4all_llama',  # internally handles prompting
-    ],
-    '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-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',
-    ],
-    'instruct': [],
-    '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-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'],
-}
-
-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, chat, context, reduced):
-    if prompt_type in [PromptType.plain.value, str(PromptType.plain.value),
-                       PromptType.plain.name]:
-        promptA = promptB = PreInstruct = PreInput = PreResponse = ''
-        terminate_response = []
-        chat_sep = ''
-        humanstr = ''
-        botstr = ''
-    elif prompt_type == 'simple_instruct':
-        promptA = promptB = PreInstruct = PreInput = PreResponse = None
-        terminate_response = []
-        chat_sep = '\n'
-        humanstr = ''
-        botstr = ''
-    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_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_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 = human
-        promptB = promptA = '%s%s ' % (preprompt, start)
-
-        PreInstruct = ""
-
-        PreInput = None
-
-        if reduced:
-            # 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 = [start, PreResponse]
-        chat_sep = '\n'
-        humanstr = human  # tag before human talks
-        botstr = bot  # tag before bot talks
-    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_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_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_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 = prompt_tokens
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = ""
-        PreInput = None
-        PreResponse = answer_tokens
-        eos = '<|endoftext|>'  # neox eos
-        terminate_response = [start, PreResponse, eos]
-        chat_sep = eos
-        humanstr = prompt_tokens
-        botstr = answer_tokens
-    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 = prompt_tokens
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = ""
-        PreInput = None
-        PreResponse = answer_tokens
-        pend = "<|prefix_end|>"
-        eos = "</s>"
-        terminate_response = [start, PreResponse, pend, eos]
-        chat_sep = eos
-        humanstr = prompt_tokens
-        botstr = answer_tokens
-    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_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_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_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_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."""
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = """
-### Instruction:
-"""
-        PreInput = None
-        PreResponse = """
-### Response:
-"""
-        terminate_response = [PreResponse]
-        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."""
-        start = ''
-        promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = """USER: """
-        PreInput = None
-        PreResponse = """ASSISTANT: """
-        terminate_response = [PreResponse]
-        chat_sep = '\n'
-        humanstr = PreInstruct
-        botstr = PreResponse
-
-    else:
-        raise RuntimeError("No such prompt_type=%s" % prompt_type)
-
-    return promptA, promptB, PreInstruct, PreInput, PreResponse, terminate_response, chat_sep, humanstr, botstr
-
-
-def generate_prompt(data_point, prompt_type, chat, reduced):
-    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)
-    assert prompt_type in prompt_types, "Bad prompt type: %s" % prompt_type
-    promptA, promptB, PreInstruct, PreInput, PreResponse, \
-        terminate_response, chat_sep, humanstr, botstr = get_prompt(prompt_type, chat, context, reduced)
-
-    prompt = context if not reduced else ''
-
-    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_newline(prompt_type, prompt)
-    elif instruction and input and PreInstruct is None and PreInput is not None:
-        prompt += f"""{PreInput}{instruction}
-{input}"""
-        prompt = inject_newline(prompt_type, prompt)
-    elif input and instruction and PreInput is None and PreInstruct is not None:
-        prompt += f"""{PreInstruct}{instruction}
-{input}"""
-        prompt = inject_newline(prompt_type, prompt)
-    elif instruction and PreInstruct is not None:
-        prompt += f"""{PreInstruct}{instruction}"""
-        prompt = inject_newline(prompt_type, prompt)
-    elif input and PreInput is not None:
-        prompt += f"""{PreInput}{input}"""
-        prompt = inject_newline(prompt_type, prompt)
-    elif input and instruction and PreInput is not None:
-        prompt += f"""{PreInput}{instruction}{input}"""
-        prompt = inject_newline(prompt_type, prompt)
-    elif input and instruction and PreInstruct is not None:
-        prompt += f"""{PreInstruct}{instruction}{input}"""
-        prompt = inject_newline(prompt_type, prompt)
-    elif input and instruction:
-        # i.e. for simple_instruct
-        prompt += f"""{instruction}: {input}"""
-        prompt = inject_newline(prompt_type, prompt)
-    elif input:
-        prompt += f"""{input}"""
-        prompt = inject_newline(prompt_type, prompt)
-    elif instruction:
-        prompt += f"""{instruction}"""
-        prompt = inject_newline(prompt_type, prompt)
-
-    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
-
-
-def inject_newline(prompt_type, prompt):
-    if prompt_type not in [-1, '-1', 'plain', 'simple_instruct']:
-        # only add new line if structured prompt, while 'plain' is just generation of next tokens from input
-        prompt += '\n'
-    return prompt
-
-
-class Prompter(object):
-    def __init__(self, prompt_type, debug=False, chat=False, stream_output=False, repeat_penalty=True,
-                 allowed_repeat_line_length=10):
-        self.prompt_type = prompt_type
-        data_point = dict(instruction='', input='', output='')
-        _, self.pre_response, self.terminate_response, self.chat_sep = \
-            generate_prompt(data_point, prompt_type, chat, False)
-        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
-        self.promptA, self.promptB, self.PreInstruct, self.PreInput, self.PreResponse, \
-            self.terminate_response, self.chat_sep, self.humanstr, self.botstr = \
-            get_prompt(prompt_type, chat, context, reduced)
-
-    def generate_prompt(self, data_point):
-        reduced = False
-        prompt, _, _, _ = generate_prompt(data_point, self.prompt_type, self.chat, reduced)
-        if self.debug:
-            print("prompt: ", prompt, flush=True)
-        self.prompt = prompt
-        return prompt
-
-    def get_response(self, outputs, prompt=None, sanitize_bot_response=True):
-        if isinstance(outputs, str):
-            outputs = [outputs]
-        if self.debug:
-            print("output:\n", '\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)
-            response = response.strip("\n")
-            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 self.botstr in output:
-                    if self.humanstr:
-                        output = clean_response(output.split(self.botstr)[1].strip().split(self.humanstr)[0].strip())
-                    else:
-                        # i.e. use after bot but only up to next bot
-                        output = clean_response(output.split(self.botstr)[1].strip().split(self.botstr)[0].strip())
-                else:
-                    # output = clean_response(output.strip())
-                    # 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).strip()
-                if self.repeat_penalty:
-                    output = clean_repeats(output).strip()
-                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].strip()
-                    else:
-                        output = output.strip()
-                else:
-                    output = output.strip()
-            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'
-            outputs[oi] = output
-        # join all outputs, only one extra new line between outputs
-        output = '\n'.join(outputs)
-        if self.debug:
-            print("outputclean:\n", '\n\n'.join(outputs), flush=True)
-        return output

prompter.py

@@ -0,0 +1 @@
+../../prompter.py

requirements.txt

@@ -1,100 +0,0 @@
-# for generate (gradio server) and finetune
-datasets==2.12.0
-sentencepiece==0.1.97
-gradio==3.31.0
-huggingface_hub==0.14.1
-appdirs==1.4.4
-fire==0.5.0
-docutils==0.19
-torch==2.0.1
-evaluate==0.4.0
-rouge_score==0.1.2
-sacrebleu==2.3.1
-scikit-learn==1.2.2
-alt-profanity-check==1.2.2
-better-profanity==0.6.1
-numpy==1.24.2
-pandas==2.0.0
-matplotlib==3.7.1
-loralib==0.1.1
-bitsandbytes==0.39.0
-accelerate==0.19.0
-git+https://github.com/huggingface/peft.git@3714aa2fff158fdfa637b2b65952580801d890b2
-transformers==4.28.1
-tokenizers==0.13.3
-APScheduler==3.10.1
-
-# optional for generate
-pynvml==11.5.0
-psutil==5.9.4
-boto3==1.26.101
-botocore==1.29.101
-
-# optional for finetune
-tensorboard==2.12.1
-neptune==1.1.1
-
-# for gradio client
-gradio_client==0.2.5
-beautifulsoup4==4.12.2
-markdown==3.4.1
-
-# data and testing
-pytest==7.2.2
-pytest-xdist==3.2.1
-nltk==3.8.1
-textstat==0.7.3
-pandoc==2.3
-#pypandoc==1.11
-pypandoc_binary==1.11
-openpyxl==3.1.2
-lm_dataformat==0.0.20
-bioc==2.0
-
-# falcon
-einops==0.6.1
-# optional for chat with PDF
-langchain==0.0.183
-pypdf==3.8.1
-tiktoken==0.3.3
-# avoid textract, requires old six
-#textract==1.6.5
-
-# for HF embeddings
-sentence_transformers==2.2.2
-# for OpenAI embeddings (requires key)
-openai==0.27.6
-
-# local vector db
-chromadb==0.3.25
-# server vector db
-#pymilvus==2.2.8
-
-# weak url support, if can't install opencv etc. If comment-in this one, then comment-out unstructured[local-inference]==0.6.6
-# unstructured==0.6.6
-
-# strong support for images
-# Requires on Ubuntu: sudo apt-get install libmagic-dev poppler-utils tesseract-ocr libreoffice
-unstructured[local-inference]==0.6.6
-#pdf2image==1.16.3
-#pytesseract==0.3.10
-pillow
-
-pdfminer.six==20221105
-urllib3==1.26.6
-requests_file==1.5.1
-
-#pdf2image==1.16.3
-#pytesseract==0.3.10
-tabulate==0.9.0
-# FYI pandoc already part of requirements.txt
-
-# JSONLoader, but makes some trouble for some users
-# jq==1.4.1
-
-# to check licenses
-# Run: pip-licenses|grep -v 'BSD\|Apache\|MIT'
-pip-licenses==4.3.0
-
-# weaviate vector db
-weaviate-client==3.19.2

stopping.py

@@ -1,72 +0,0 @@
-import torch
-from transformers import StoppingCriteria, StoppingCriteriaList
-
-from prompter import PromptType
-
-
-class StoppingCriteriaSub(StoppingCriteria):
-
-    def __init__(self, stops=[], encounters=[], device="cuda"):
-        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)
-
-    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
-        # 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, tokenizer, device, human='<human>:', bot="<bot>:"):
-    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)])
-    else:
-        stopping_criteria = StoppingCriteriaList()
-    return stopping_criteria

stopping.py

@@ -0,0 +1 @@
+../../stopping.py

utils.py

@@ -1,843 +0,0 @@
-import contextlib
-import functools
-import hashlib
-import inspect
-import os
-import gc
-import pathlib
-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 get_torch_allocated():
-    import torch
-    return torch.cuda.memory_allocated()
-
-
-def get_device():
-    import torch
-    if torch.cuda.is_available():
-        device = "cuda"
-    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
-    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
-
-    # https://github.com/gpuopenanalytics/pynvml/blob/master/help_query_gpu.txt
-    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
-
-    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(output=None, base_model=None, save_dir=None):
-    try:
-        return _save_generate_output(output=output, base_model=base_model, save_dir=save_dir)
-    except Exception as e:
-        traceback.print_exc()
-        print('Exception in saving: %s' % str(e))
-
-
-def _save_generate_output(output=None, base_model=None, save_dir=None):
-    """
-    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
-    """
-    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
-    if output[-10:] == '\n\n<human>:':
-        # remove trailing <human>:
-        output = output[:-10]
-    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(text=output, time=time.ctime(), base_model=base_model)
-                ) + ",\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):
-    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()

utils.py

@@ -0,0 +1 @@
+../../utils.py