Update with h2oGPT hash 13a8343d2a96885985bda8c4480bbb23cf55bb9b

LICENSE

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

LICENSE

@@ -0,0 +1,201 @@
+                                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.

client_test.py

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

client_test.py

@@ -0,0 +1,278 @@
+"""
+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
+
+from enums import DocumentChoices
+
+debug = False
+
+os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
+
+
+def get_client(serialize=True):
+    from gradio_client import Client
+
+    client = Client(os.getenv('HOST', "http://localhost:7860"), serialize=serialize)
+    if debug:
+        print(client.view_api(all_endpoints=True))
+    return client
+
+
+def get_args(prompt, prompt_type, chat=False, stream_output=False,
+             max_new_tokens=50,
+             top_k_docs=3,
+             langchain_mode='Disabled'):
+    from collections import OrderedDict
+    kwargs = OrderedDict(instruction=prompt if chat else '',  # only for chat=True
+                         iinput='',  # only for chat=True
+                         context='',
+                         # streaming output is supported, loops over and outputs each generation in streaming mode
+                         # but leave stream_output=False for simple input/output mode
+                         stream_output=stream_output,
+                         prompt_type=prompt_type,
+                         prompt_dict='',
+                         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=top_k_docs,
+                         chunk=True,
+                         chunk_size=512,
+                         document_choice=[DocumentChoices.All_Relevant.name],
+                         )
+    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(res))
+    print(res_dict)
+    return res_dict
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic_api():
+    return run_client_nochat_api(prompt='Who are you?', prompt_type='human_bot', max_new_tokens=50)
+
+
+def run_client_nochat_api(prompt, prompt_type, max_new_tokens):
+    kwargs, args = get_args(prompt, prompt_type, chat=False, max_new_tokens=max_new_tokens)
+
+    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
+    client = get_client(serialize=True)
+    res = client.predict(
+        str(dict(kwargs)),
+        api_name=api_name,
+    )
+    print("Raw client result: %s" % res, flush=True)
+    res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
+                    response=md_to_text(ast.literal_eval(res)['response']),
+                    sources=ast.literal_eval(res)['sources'])
+    print(res_dict)
+    return res_dict
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic_api_lean():
+    return run_client_nochat_api_lean(prompt='Who are you?', prompt_type='human_bot', max_new_tokens=50)
+
+
+def run_client_nochat_api_lean(prompt, prompt_type, max_new_tokens):
+    kwargs = dict(instruction_nochat=prompt)
+
+    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
+    client = get_client(serialize=True)
+    res = client.predict(
+        str(dict(kwargs)),
+        api_name=api_name,
+    )
+    print("Raw client result: %s" % res, flush=True)
+    res_dict = dict(prompt=kwargs['instruction_nochat'],
+                    response=md_to_text(ast.literal_eval(res)['response']),
+                    sources=ast.literal_eval(res)['sources'])
+    print(res_dict)
+    return res_dict
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic_api_lean_morestuff():
+    return run_client_nochat_api_lean_morestuff(prompt='Who are you?', prompt_type='human_bot', max_new_tokens=50)
+
+
+def run_client_nochat_api_lean_morestuff(prompt, prompt_type, max_new_tokens):
+    kwargs = dict(
+        instruction='',
+        iinput='',
+        context='',
+        stream_output=False,
+        prompt_type='human_bot',
+        temperature=0.1,
+        top_p=0.75,
+        top_k=40,
+        num_beams=1,
+        max_new_tokens=256,
+        min_new_tokens=0,
+        early_stopping=False,
+        max_time=20,
+        repetition_penalty=1.0,
+        num_return_sequences=1,
+        do_sample=True,
+        chat=False,
+        instruction_nochat=prompt,
+        iinput_nochat='',
+        langchain_mode='Disabled',
+        top_k_docs=4,
+        document_choice=['All'],
+    )
+
+    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
+    client = get_client(serialize=True)
+    res = client.predict(
+        str(dict(kwargs)),
+        api_name=api_name,
+    )
+    print("Raw client result: %s" % res, flush=True)
+    res_dict = dict(prompt=kwargs['instruction_nochat'],
+                    response=md_to_text(ast.literal_eval(res)['response']),
+                    sources=ast.literal_eval(res)['sources'])
+    print(res_dict)
+    return res_dict
+
+
+@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()
+    test_client_basic_api()
+    test_client_basic_api_lean()
+    test_client_basic_api_lean_morestuff()

create_data.py

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

create_data.py

@@ -0,0 +1,1809 @@
+"""
+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))

enums.py

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

enums.py

@@ -0,0 +1,46 @@
+from enum import Enum
+
+
+class PromptType(Enum):
+    custom = -1
+    plain = 0
+    instruct = 1
+    quality = 2
+    human_bot = 3
+    dai_faq = 4
+    summarize = 5
+    simple_instruct = 6
+    instruct_vicuna = 7
+    instruct_with_end = 8
+    human_bot_orig = 9
+    prompt_answer = 10
+    open_assistant = 11
+    wizard_lm = 12
+    wizard_mega = 13
+    instruct_vicuna2 = 14
+    instruct_vicuna3 = 15
+    wizard2 = 16
+    wizard3 = 17
+    instruct_simple = 18
+
+
+class DocumentChoices(Enum):
+    All_Relevant = 0
+    All_Relevant_Only_Sources = 1
+    Only_All_Sources = 2
+    Just_LLM = 3
+
+
+class LangChainMode(Enum):
+    """LangChain mode"""
+
+    DISABLED = "Disabled"
+    CHAT_LLM = "ChatLLM"
+    LLM = "LLM"
+    ALL = "All"
+    WIKI = "wiki"
+    WIKI_FULL = "wiki_full"
+    USER_DATA = "UserData"
+    MY_DATA = "MyData"
+    GITHUB_H2OGPT = "github h2oGPT"
+    H2O_DAI_DOCS = "DriverlessAI docs"

finetune.py

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

finetune.py

@@ -0,0 +1,676 @@
+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, PromptType
+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
+    prompt_dict = ''  # only for custom prompt_type
+    assert prompt_type != PromptType.custom.name, "custom not setup for finetune"
+    full_prompt, _, _, _ = generate_prompt(data_point, prompt_type, prompt_dict, 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, prompt_dict, 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)

generate.py

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

generate.py

@@ -0,0 +1,1712 @@
+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 enums import DocumentChoices, LangChainMode
+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, start_faulthandler
+
+start_faulthandler()
+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, PromptType, get_prompt
+from stopping import get_stopping
+
+eval_extra_columns = ['prompt', 'response', 'score']
+
+langchain_modes = [x.value for x in list(LangChainMode)]
+
+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,
+        prompt_dict: typing.Dict = 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 = [DocumentChoices.All_Relevant.name],
+        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 = None,
+        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 prompt_dict: If prompt_type=custom, then expects (some) items returned by get_prompt(..., return_dict=True)
+    :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
+           Default is instructor-large with 768 parameters per embedding if have GPUs, else all-MiniLM-L6-v1 if no GPUs
+           Can also choose simpler model with 384 parameters per embedding: "sentence-transformers/all-MiniLM-L6-v2"
+           Can also choose even better embedding with 1024 parameters: 'hkunlp/instructor-xl'
+           We support automatically changing of embeddings for chroma, with a backup of db made if this is done
+    :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(int(os.getenv("HUGGINGFACE_SPACES", '0')))
+    is_gpth2oai = bool(int(os.getenv("GPT_H2O_AI", '0')))
+    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(int(os.environ.get("allow_upload_to_user_data", allow_upload_to_user_data)))
+    allow_upload_to_my_data = bool(int(os.environ.get("allow_upload_to_my_data", allow_upload_to_my_data)))
+    height = int(os.environ.get("HEIGHT", height))
+    h2ocolors = bool(int(os.getenv('h2ocolors', h2ocolors)))
+
+    # 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 hf_embedding_model is None:
+            hf_embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
+    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
+        if hf_embedding_model is None:
+            # if no GPUs, use simpler embedding model to avoid cost in time
+            hf_embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
+    else:
+        if hf_embedding_model is None:
+            # if still None, then set default
+            hf_embedding_model = 'hkunlp/instructor-large'
+
+    # 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)
+
+    user_set_max_new_tokens = max_new_tokens is not None
+
+    placeholder_instruction, placeholder_input, \
+        stream_output, show_examples, \
+        prompt_type, prompt_dict, \
+        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, prompt_dict,
+                            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,
+                            chunk,
+                            chunk_size,
+                            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
+            try:
+                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())
+            finally:
+                # in case updated embeddings or created new embeddings
+                clear_torch_cache()
+            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_seq_len') and isinstance(config.max_seq_len, int):
+        tokenizer.model_max_length = config.max_seq_len
+    elif 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:
+        if verbose:
+            print("Could not determine model_max_length, setting to 2048", flush=True)
+        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
+
+
+no_default_param_names = [
+    'instruction',
+    'iinput',
+    'context',
+    'instruction_nochat',
+    'iinput_nochat',
+]
+
+eval_func_param_names = ['instruction',
+                         'iinput',
+                         'context',
+                         'stream_output',
+                         'prompt_type',
+                         'prompt_dict',
+                         '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',
+                         'chunk',
+                         'chunk_size',
+                         'document_choice',
+                         ]
+
+# form evaluate defaults for submit_nochat_api
+eval_func_param_names_defaults = eval_func_param_names.copy()
+for k in no_default_param_names:
+    if k in eval_func_param_names_defaults:
+        eval_func_param_names_defaults.remove(k)
+
+
+def evaluate_from_str(
+        model_state,
+        my_db_state,
+        # START NOTE: Examples must have same order of parameters
+        user_kwargs,
+        # END NOTE: Examples must have same order of parameters
+        default_kwargs=None,
+        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,
+):
+    if isinstance(user_kwargs, str):
+        user_kwargs = ast.literal_eval(user_kwargs)
+    # only used for submit_nochat_api
+    user_kwargs['chat'] = False
+    user_kwargs['stream_output'] = False
+
+    assert set(list(default_kwargs.keys())) == set(eval_func_param_names)
+    # correct ordering.  Note some things may not be in default_kwargs, so can't be default of user_kwargs.get()
+    args_list = [user_kwargs[k] if k in user_kwargs else default_kwargs[k] for k in eval_func_param_names]
+
+    ret = evaluate(
+        model_state,
+        my_db_state,
+        # START NOTE: Examples must have same order of parameters
+        *tuple(args_list),
+        # END NOTE: Examples must have same order of parameters
+        src_lang=src_lang,
+        tgt_lang=tgt_lang,
+        debug=debug,
+        concurrency_count=concurrency_count,
+        save_dir=save_dir,
+        sanitize_bot_response=sanitize_bot_response,
+        model_state0=model_state0,
+        memory_restriction_level=memory_restriction_level,
+        raise_generate_gpu_exceptions=raise_generate_gpu_exceptions,
+        chat_context=chat_context,
+        lora_weights=lora_weights,
+        load_db_if_exists=load_db_if_exists,
+        dbs=dbs,
+        user_path=user_path,
+        detect_user_path_changes_every_query=detect_user_path_changes_every_query,
+        use_openai_embedding=use_openai_embedding,
+        use_openai_model=use_openai_model,
+        hf_embedding_model=hf_embedding_model,
+        db_type=db_type,
+        n_jobs=n_jobs,
+        first_para=first_para,
+        text_limit=text_limit,
+        verbose=verbose,
+        cli=cli,
+    )
+    try:
+        for ret1 in ret:
+            yield ret1
+    finally:
+        # clear before return, in finally in case GPU OOM exception
+        clear_torch_cache()
+
+
+def evaluate(
+        model_state,
+        my_db_state,
+        # START NOTE: Examples must have same order of parameters
+        instruction,
+        iinput,
+        context,
+        stream_output,
+        prompt_type,
+        prompt_dict,
+        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,
+        chunk,
+        chunk_size,
+        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,
+        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 db_type is not None
+    assert top_k_docs is not None and isinstance(top_k_docs, int)
+    assert chunk is not None and isinstance(chunk, bool)
+    assert chunk_size is not None and isinstance(chunk_size, 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, prompt_dict, 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,
+                           top_k_docs=top_k_docs,
+                           temperature=temperature,
+                           repetition_penalty=repetition_penalty,
+                           top_k=top_k,
+                           top_p=top_p,
+                           prompt_type=prompt_type,
+                           prompt_dict=prompt_dict,
+                           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
+            # clear before return, since .then() never done if from API
+            clear_torch_cache()
+            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, prompt_dict, 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:
+                        # clear before return, since .then() never done if from API
+                        clear_torch_cache()
+                        # 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:
+                    try:
+                        outputs = model.generate(**gen_kwargs)
+                    finally:
+                        clear_torch_cache()  # has to be here for API submit_nochat_api since.then() not called
+                    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, prompt_dict,
+                        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, chunk, chunk_size,
+                        verbose):
+    use_defaults = False
+    use_default_examples = True
+    examples = []
+    task_info = 'LLM'
+    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"""
+
+    use_placeholder_instruction_as_example = False
+    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
+        use_placeholder_instruction_as_example = True
+        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
+        use_placeholder_instruction_as_example = True
+    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
+        use_placeholder_instruction_as_example = True
+        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
+        num_beams = num_beams or 1
+        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, prompt_dict,
+                   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_placeholder_instruction_as_example:
+        examples += [[placeholder_instruction, ''] + params_list]
+
+    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', top_k_docs, chunk, chunk_size, [DocumentChoices.All_Relevant.name]]
+        # 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))
+
+    if prompt_type == PromptType.custom.name and not prompt_dict:
+        raise ValueError("Unexpected to get non-empty prompt_dict=%s for prompt_type=%s" % (prompt_dict, prompt_type))
+
+    # get prompt_dict from prompt_type, so user can see in UI etc., or for custom do nothing except check format
+    prompt_dict, error0 = get_prompt(prompt_type, prompt_dict,
+                                     chat=False, context='', reduced=False, return_dict=True)
+    if error0:
+        raise RuntimeError("Prompt wrong: %s" % error0)
+
+    return placeholder_instruction, placeholder_input, \
+        stream_output, show_examples, \
+        prompt_type, prompt_dict, \
+        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 = no_default_param_names + [
+        # 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
+
+
+def get_max_max_new_tokens(model_state, **kwargs):
+    if kwargs['max_new_tokens'] and kwargs['user_set_max_new_tokens']:
+        max_max_new_tokens = kwargs['max_new_tokens']
+    elif kwargs['memory_restriction_level'] == 1:
+        max_max_new_tokens = 768
+    elif kwargs['memory_restriction_level'] == 2:
+        max_max_new_tokens = 512
+    elif kwargs['memory_restriction_level'] >= 3:
+        max_max_new_tokens = 256
+    else:
+        if not isinstance(model_state[1], str):
+            max_max_new_tokens = model_state[1].model_max_length
+        else:
+            # FIXME: Need to update after new model loaded, so user can control with slider
+            max_max_new_tokens = 2048
+    return max_max_new_tokens
+
+
+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)

gpt4all_llm.py

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

gpt4all_llm.py

@@ -0,0 +1,258 @@
+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:
+    model_max_length = 2048
+
+    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
+        # still have to avoid crazy sizes, else hit llama_tokenize: too many tokens -- might still hit, not fatal
+        prompt = prompt[-self.n_ctx * 4:]
+        prompt_tokens = self.client.tokenize(b" " + prompt.encode("utf-8"))
+        num_prompt_tokens = len(prompt_tokens)
+        if num_prompt_tokens > self.n_ctx:
+            # conservative by using int()
+            chars_per_token = int(len(prompt) / num_prompt_tokens)
+            prompt = prompt[-self.n_ctx * chars_per_token:]
+            if verbose:
+                print("reducing tokens, assuming average of %s chars/token: %s" % chars_per_token, flush=True)
+                prompt_tokens2 = self.client.tokenize(b" " + prompt.encode("utf-8"))
+                num_prompt_tokens2 = len(prompt_tokens2)
+                print("reduced tokens from %d -> %d" % (num_prompt_tokens, num_prompt_tokens2), flush=True)
+        if verbose:
+            print("_call prompt: %s" % prompt, flush=True)
+        return super()._call(prompt, stop=stop, run_manager=run_manager)

gpt_langchain.py

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

gpt_langchain.py

@@ -0,0 +1,1633 @@
+import glob
+import inspect
+import os
+import pathlib
+import pickle
+import queue
+import random
+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 langchain.embeddings import HuggingFaceInstructEmbeddings
+from tqdm import tqdm
+
+from enums import DocumentChoices
+from prompter import non_hf_types, PromptType
+from utils import wrapped_partial, EThread, import_matplotlib, sanitize_filename, makedirs, get_url, flatten_list, \
+    get_device, ProgressParallel, remove, hash_file, clear_torch_cache
+
+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", load_db_if_exists=True,
+           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
+
+        if os.getenv('WEAVIATE_URL', None):
+            client = _create_local_weaviate_client()
+        else:
+            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)
+
+        # see if already actually have persistent db, and deal with possible changes in embedding
+        db = get_existing_db(None, persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
+                             hf_embedding_model, verbose=False)
+        if db is None:
+            db = Chroma.from_documents(documents=sources,
+                                       embedding=embedding,
+                                       persist_directory=persist_directory,
+                                       collection_name=collection_name,
+                                       anonymized_telemetry=False)
+            db.persist()
+            clear_embedding(db)
+            save_embed(db, use_openai_embedding, hf_embedding_model)
+        else:
+            # then just add
+            db, num_new_sources, new_sources_metadata = add_to_db(db, sources, db_type=db_type,
+                                                                  use_openai_embedding=use_openai_embedding,
+                                                                  hf_embedding_model=hf_embedding_model)
+    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,
+              use_openai_embedding=False,
+              hf_embedding_model=None):
+    assert hf_embedding_model is not None
+    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,
+                                                          embedding_function=db._collection._embedding_function)
+            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()
+        clear_embedding(db)
+        save_embed(db, use_openai_embedding, hf_embedding_model)
+    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
+
+        if os.getenv('WEAVIATE_URL', None):
+            client = _create_local_weaviate_client()
+        else:
+            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)
+        if 'instructor' in hf_embedding_model:
+            encode_kwargs = {'normalize_embeddings': True}
+            embedding = HuggingFaceInstructEmbeddings(model_name=hf_embedding_model,
+                                                      model_kwargs=model_kwargs,
+                                                      encode_kwargs=encode_kwargs)
+        else:
+            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,
+            prompt_dict=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 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,
+                                         prompt_dict=prompt_dict,
+                                         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=chunk, 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=chunk, 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=chunk, chunk_size=chunk_size)
+    elif file.lower().endswith('.odt'):
+        docs1 = UnstructuredODTLoader(file_path=file).load()
+        add_meta(docs1, file)
+        doc1 = chunk_sources(docs1, chunk=chunk, 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=chunk, 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=chunk, 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=chunk, chunk_size=chunk_size)
+    elif file.lower().endswith('.md'):
+        docs1 = UnstructuredMarkdownLoader(file).load()
+        add_meta(docs1, file)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
+    elif file.lower().endswith('.enex'):
+        docs1 = EverNoteLoader(file).load()
+        add_meta(doc1, file)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
+    elif file.lower().endswith('.epub'):
+        docs1 = UnstructuredEPubLoader(file).load()
+        add_meta(docs1, file)
+        doc1 = chunk_sources(docs1, chunk=chunk, 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=chunk, 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=chunk, 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=chunk, 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
+            # load() still chunks by pages, but every page has title at start to help
+            doc1 = PyMuPDFLoader(file).load()
+        else:
+            # open-source fallback
+            # load() still chunks by pages, but every page has title at start to help
+            doc1 = PyPDFLoader(file).load()
+        # 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=chunk, 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=chunk, chunk_size=chunk_size)
+        else:
+            docs = [doc1]
+    elif isinstance(doc1, list) and len(doc1) == 1:
+        if chunk:
+            docs = chunk_sources(doc1, chunk=chunk, 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(None, 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 check_update_chroma_embedding(db, use_openai_embedding, hf_embedding_model, langchain_mode):
+    changed_db = False
+    if load_embed(db) != (use_openai_embedding, hf_embedding_model):
+        print("Detected new embedding, updating db: %s" % langchain_mode, flush=True)
+        # handle embedding changes
+        db_get = db.get()
+        sources = [Document(page_content=result[0], metadata=result[1] or {})
+                   for result in zip(db_get['documents'], db_get['metadatas'])]
+        # delete index, has to be redone
+        persist_directory = db._persist_directory
+        shutil.move(persist_directory, persist_directory + "_" + str(uuid.uuid4()) + ".bak")
+        db_type = 'chroma'
+        load_db_if_exists = False
+        db = get_db(sources, use_openai_embedding=use_openai_embedding, db_type=db_type,
+                    persist_directory=persist_directory, load_db_if_exists=load_db_if_exists,
+                    langchain_mode=langchain_mode,
+                    collection_name=None,
+                    hf_embedding_model=hf_embedding_model)
+        if False:
+            # below doesn't work if db already in memory, so have to switch to new db as above
+            # upsert does new embedding, but if index already in memory, complains about size mismatch etc.
+            client_collection = db._client.get_collection(name=db._collection.name,
+                                                          embedding_function=db._collection._embedding_function)
+            client_collection.upsert(ids=db_get['ids'], metadatas=db_get['metadatas'], documents=db_get['documents'])
+        changed_db = True
+        print("Done updating db for new embedding: %s" % langchain_mode, flush=True)
+
+    return db, changed_db
+
+
+def get_existing_db(db, persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
+                    hf_embedding_model, verbose=False, check_embedding=True):
+    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')):
+        if db is None:
+            if verbose:
+                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)
+            if verbose:
+                print("DONE Loading db: %s" % langchain_mode, flush=True)
+        else:
+            if verbose:
+                print("USING already-loaded db: %s" % langchain_mode, flush=True)
+        if check_embedding:
+            db_trial, changed_db = check_update_chroma_embedding(db, use_openai_embedding, hf_embedding_model,
+                                                                 langchain_mode)
+            if changed_db:
+                db = db_trial
+                # only call persist if really changed db, else takes too long for large db
+                db.persist()
+                clear_embedding(db)
+        save_embed(db, use_openai_embedding, hf_embedding_model)
+        return db
+    return None
+
+
+def clear_embedding(db):
+    # don't keep on GPU, wastes memory, push back onto CPU and only put back on GPU once again embed
+    db._embedding_function.client.cpu()
+    clear_torch_cache()
+
+
+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 save_embed(db, use_openai_embedding, hf_embedding_model):
+    embed_info_file = os.path.join(db._persist_directory, 'embed_info')
+    with open(embed_info_file, 'wb') as f:
+        pickle.dump((use_openai_embedding, hf_embedding_model), f)
+    return use_openai_embedding, hf_embedding_model
+
+
+def load_embed(db):
+    embed_info_file = os.path.join(db._persist_directory, 'embed_info')
+    if os.path.isfile(embed_info_file):
+        with open(embed_info_file, 'rb') as f:
+            use_openai_embedding, hf_embedding_model = pickle.load(f)
+    else:
+        # migration, assume defaults
+        use_openai_embedding, hf_embedding_model = False, "sentence-transformers/all-MiniLM-L6-v2"
+    return use_openai_embedding, hf_embedding_model
+
+
+def get_persist_directory(langchain_mode):
+    return 'db_dir_%s' % langchain_mode  # single place, no special names for each case
+
+
+def _make_db(use_openai_embedding=False,
+             hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
+             first_para=False, text_limit=None,
+             chunk=True, chunk_size=512,
+             langchain_mode=None,
+             user_path=None,
+             db_type='faiss',
+             load_db_if_exists=True,
+             db=None,
+             n_jobs=-1,
+             verbose=False):
+    persist_directory = get_persist_directory(langchain_mode)
+    # see if can get persistent chroma db
+    db_trial = get_existing_db(db, persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
+                               hf_embedding_model, verbose=verbose)
+    if db_trial is not None:
+        db = db_trial
+
+    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=chunk, 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=chunk, 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=chunk, 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=chunk, 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
+                # FIXME: If first had old Hash=None and switch embeddings,
+                #  then re-embed, and then hit here and reload so have hash, and then re-embed.
+                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,
+                                                              use_openai_embedding=use_openai_embedding,
+                                                              hf_embedding_model=hf_embedding_model)
+        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}
+    try:
+        return _run_qa_db(**kwargs)
+    finally:
+        clear_torch_cache()
+
+
+def _run_qa_db(query=None,
+               use_openai_model=False, use_openai_embedding=False,
+               first_para=False, text_limit=None, top_k_docs=4, chunk=True, chunk_size=512,
+               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,
+               prompt_dict=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=[DocumentChoices.All_Relevant.name],
+               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
+        prompt_dict = prompter.prompt_dict
+    if model is not None:
+        assert prompt_type is not None
+        if prompt_type == PromptType.custom.name:
+            assert prompt_dict is not None  # should at least be {} or ''
+        else:
+            prompt_dict = ''
+    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,
+                                                         prompt_dict=prompt_dict,
+                                                         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
+
+    if isinstance(document_choice, str):
+        # support string as well
+        document_choice = [document_choice]
+    # get first DocumentChoices as command to use, ignore others
+    doc_choices_set = set([x.name for x in list(DocumentChoices)])
+    cmd = [x for x in document_choice if x in doc_choices_set]
+    cmd = None if len(cmd) == 0 else cmd[0]
+    # now have cmd, filter out for only docs
+    document_choice = [x for x in document_choice if x not in doc_choices_set]
+
+    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 cmd in [DocumentChoices.All_Relevant_Only_Sources.name, DocumentChoices.Only_All_Sources.name]:
+        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, top_k_docs=4, chunk=True, chunk_size=512,
+                         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,
+                         prompt_dict=None,
+                         cut_distanct=1.1,
+                         load_db_if_exists=False,
+                         db=None,
+                         langchain_mode=None,
+                         document_choice=[DocumentChoices.All_Relevant.name],
+                         n_jobs=-1,
+                         # beyond run_db_query:
+                         llm=None,
+                         verbose=False,
+                         cmd=None,
+                         ):
+    # 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 top_k_docs to avoid
+    # Chroma collection MyData contains fewer than 4 elements.
+    # type logger error
+    k_db = 1000 if db_type == 'chroma' else top_k_docs  # top_k_docs=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 not isinstance(db, Chroma):
+            # only chroma supports filtering
+            filter_kwargs = {}
+        else:
+            # if here then some cmd + documents selected or just documents selected
+            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) == 1:
+                # degenerate UX bug in chroma
+                one_filter = [{"source": {"$eq": x}} for x in document_choice][0]
+                filter_kwargs = dict(filter=one_filter)
+            else:
+                # shouldn't reach
+                filter_kwargs = {}
+        if cmd == DocumentChoices.Just_LLM.name:
+            docs = []
+            scores = []
+        elif cmd == DocumentChoices.Only_All_Sources.name:
+            if isinstance(db, Chroma):
+                db_get = db._collection.get(where=filter_kwargs.get('filter'))
+            else:
+                db_get = db.get()
+            # similar to langchain's chroma's _results_to_docs_and_scores
+            docs_with_score = [(Document(page_content=result[0], metadata=result[1] or {}), 0)
+                               for result in zip(db_get['documents'], db_get['metadatas'])][:top_k_docs]
+            docs = [x[0] for x in docs_with_score]
+            scores = [x[1] for x in docs_with_score]
+        else:
+            docs_with_score = db.similarity_search_with_score(query, k=k_db, **filter_kwargs)[:top_k_docs]
+            # 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 cmd in [DocumentChoices.All_Relevant_Only_Sources.name, DocumentChoices.Only_All_Sources.name]:
+        # 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=True, chunk_size=512):
+    if not chunk:
+        return sources
+    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:
+        persist_directory = os.path.dirname(zip_ref.namelist()[0])
+        remove(persist_directory)
+        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
+
+def _create_local_weaviate_client():
+    WEAVIATE_URL = os.getenv('WEAVIATE_URL', "http://localhost:8080")
+    WEAVIATE_USERNAME = os.getenv('WEAVIATE_USERNAME')
+    WEAVIATE_PASSWORD = os.getenv('WEAVIATE_PASSWORD')
+    WEAVIATE_SCOPE = os.getenv('WEAVIATE_SCOPE', "offline_access")
+
+    resource_owner_config = None
+    if WEAVIATE_USERNAME is not None and WEAVIATE_PASSWORD is not None:
+        resource_owner_config = weaviate.AuthClientPassword(
+            username=WEAVIATE_USERNAME,
+            password=WEAVIATE_PASSWORD,
+            scope=WEAVIATE_SCOPE
+        )
+
+    try:
+        client = weaviate.Client(WEAVIATE_URL, auth_client_secret=resource_owner_config)
+    except Exception as e:
+        print(f"Failed to create Weaviate client: {e}")
+        return None
+
+if __name__ == '__main__':
+    pass

gradio_runner.py

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

gradio_themes.py

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

gradio_themes.py

@@ -0,0 +1,183 @@
+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_ui

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

gradio_ui/css.py

@@ -0,0 +1,51 @@
+def get_css(kwargs) -> str:
+    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 += make_css_base()
+    return css_code
+
+def make_css_base() -> str:
+    return """
+    @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;
+    }
+    
+    #prompt-form {
+        border: 1px solid var(--primary-500) !important;
+    }
+    
+    #prompt-form.block {
+        border-radius: var(--block-radius) !important;
+    }
+    
+    #prompt-form textarea {
+        border: 1px solid rgb(209, 213, 219);
+    }
+    
+    #prompt-form label > div {
+        margin-top: 4px;
+    }
+    
+    button.primary:hover {
+        background-color: var(--primary-600) !important;
+        transition: .2s;
+    }
+    
+    #prompt-form-area {
+        margin-bottom: 2.5rem;
+    }
+    """

gradio_ui/prompt_form.py

@@ -0,0 +1,24 @@
+import gradio as gr
+
+
+def make_prompt_form(kwargs):
+    if kwargs['input_lines'] > 1:
+        instruction_label = "press Shift-Enter or click Submit to send message, press Enter for multiple input lines"
+    else:
+        instruction_label = "press Enter or click Submit to send message, press Shift-Enter for more lines"
+
+    with gr.Row(elem_id='prompt-form-area'):
+        with gr.Column(scale=50):
+            instruction = gr.Textbox(
+                lines=kwargs['input_lines'],
+                label='Ask anything',
+                placeholder=kwargs['placeholder_instruction'],
+                info=instruction_label,
+                elem_id='prompt-form'
+            )
+            instruction.style(container=True)
+        with gr.Row():
+            submit = gr.Button(value='Submit', variant='primary').style(full_width=False, size='sm')
+            stop_btn = gr.Button(value="Stop", variant='secondary').style(full_width=False, size='sm')
+
+    return instruction, submit, stop_btn

h2oai_pipeline.py

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

h2oai_pipeline.py

@@ -0,0 +1,177 @@
+import os
+
+from transformers import TextGenerationPipeline
+from transformers.pipelines.text_generation import ReturnType
+
+from stopping import get_stopping
+from prompter import Prompter, PromptType
+
+
+class H2OTextGenerationPipeline(TextGenerationPipeline):
+    def __init__(self, *args, debug=False, chat=False, stream_output=False,
+                 sanitize_bot_response=True,
+                 use_prompter=True, prompter=None,
+                 prompt_type=None, prompt_dict=None,
+                 max_input_tokens=2048 - 256, **kwargs):
+        """
+        HF-like pipeline, but handle instruction prompting and stopping (for some models)
+        :param args:
+        :param debug:
+        :param chat:
+        :param stream_output:
+        :param sanitize_bot_response:
+        :param use_prompter: Whether to use prompter.  If pass prompt_type, will make prompter
+        :param prompter: prompter, can pass if have already
+        :param prompt_type: prompt_type, e.g. human_bot.  See prompt_type to model mapping in from prompter.py.
+                            If use_prompter, then will make prompter and use it.
+        :param prompt_dict: dict of get_prompt(, return_dict=True) for prompt_type=custom
+        :param max_input_tokens:
+        :param kwargs:
+        """
+        super().__init__(*args, **kwargs)
+        self.prompt_text = None
+        self.use_prompter = use_prompter
+        self.prompt_type = prompt_type
+        self.prompt_dict = prompt_dict
+        self.prompter = prompter
+        if self.use_prompter:
+            if self.prompter is not None:
+                assert self.prompter.prompt_type is not None
+            else:
+                self.prompter = Prompter(self.prompt_type, self.prompt_dict, debug=debug, chat=chat,
+                                         stream_output=stream_output)
+            self.human = self.prompter.humanstr
+            self.bot = self.prompter.botstr
+            self.can_stop = True
+        else:
+            self.prompter = None
+            self.human = None
+            self.bot = None
+            self.can_stop = False
+        self.sanitize_bot_response = sanitize_bot_response
+        self.max_input_tokens = max_input_tokens  # not for generate, so ok that not kwargs
+
+    def preprocess(self, prompt_text, prefix="", handle_long_generation=None, **generate_kwargs):
+        if hasattr(self.tokenizer, 'model_max_length'):
+            # model_max_length only defined for generate.py, not raw use of h2oai_pipeline.py
+            model_max_length = self.tokenizer.model_max_length
+        else:
+            # unknown
+            model_max_length = None
+
+        verbose = bool(int(os.getenv('VERBOSE_PIPELINE', '0')))
+        if model_max_length is not None:
+            num_prompt_tokens = None
+            # can't wait for "hole" if not plain prompt_type, since would lose prefix like <human>:
+            # For https://github.com/h2oai/h2ogpt/issues/192
+            for trial in range(0, 3):
+                prompt_tokens = self.tokenizer(prompt_text)['input_ids']
+                num_prompt_tokens = len(prompt_tokens)
+                if num_prompt_tokens > model_max_length:
+                    # conservative by using int()
+                    chars_per_token = int(len(prompt_text) / num_prompt_tokens)
+                    prompt_text = prompt_text[-model_max_length * chars_per_token:]
+                    if verbose:
+                        print("reducing %s tokens, assuming average of %s chars/token for %s characters" % (
+                            num_prompt_tokens, chars_per_token, len(prompt_text)), flush=True)
+                else:
+                    if verbose:
+                        print("using %s tokens with %s chars" % (num_prompt_tokens, len(prompt_text)), flush=True)
+                    break
+
+            # if input prompt is some number of tokens, despite user request, can't have max_new_tokens more
+            #
+            if self.prompt_type not in [PromptType.plain.name, PromptType.plain.value]:
+                # then give room for prompt
+                fudge = 20
+            else:
+                fudge = 0
+            assert num_prompt_tokens is not None
+            max_new_tokens = max(0, min(generate_kwargs['max_new_tokens'],
+                                        model_max_length - (num_prompt_tokens + fudge)))
+            if max_new_tokens < generate_kwargs['max_new_tokens']:
+                if verbose:
+                    print("Reduced max_new_tokens from %s -> %s" % (generate_kwargs['max_new_tokens'], max_new_tokens))
+                generate_kwargs['max_new_tokens'] = max_new_tokens
+
+        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.prompt_dict,
+                                             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}

loaders.py

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

loaders.py

@@ -0,0 +1,50 @@
+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

prompter.py

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

prompter.py

@@ -0,0 +1,606 @@
+import ast
+import time
+from enums import PromptType  # also supports imports from this file from other files
+
+non_hf_types = ['gpt4all_llama', 'llama', 'gptj']
+
+
+prompt_type_to_model_name = {
+    'plain': [
+        'EleutherAI/gpt-j-6B',
+        'EleutherAI/pythia-6.9b',
+        'EleutherAI/pythia-12b',
+        'EleutherAI/pythia-12b-deduped',
+        'EleutherAI/gpt-neox-20b',
+        'openlm-research/open_llama_7b_700bt_preview',
+        'decapoda-research/llama-7b-hf',
+        'decapoda-research/llama-13b-hf',
+        'decapoda-research/llama-30b-hf',
+        'decapoda-research/llama-65b-hf',
+        'facebook/mbart-large-50-many-to-many-mmt',
+        'philschmid/bart-large-cnn-samsum',
+        'philschmid/flan-t5-base-samsum',
+        'gpt2',
+        'distilgpt2',
+        'mosaicml/mpt-7b-storywriter',
+        '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',
+        'h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b',
+        'h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b-v2',
+    ],
+    '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'],
+    "instruct_simple": ['JosephusCheung/Guanaco'],
+}
+
+inv_prompt_type_to_model_name = {v.strip(): k for k, l in prompt_type_to_model_name.items() for v in l}
+inv_prompt_type_to_model_lower = {v.strip().lower(): k for k, l in prompt_type_to_model_name.items() for v in l}
+
+prompt_types_strings = []
+for p in PromptType:
+    prompt_types_strings.extend([p.name])
+
+prompt_types = []
+for p in PromptType:
+    prompt_types.extend([p.name, p.value, str(p.value)])
+
+
+def get_prompt(prompt_type, prompt_dict, chat, context, reduced, return_dict=False):
+    prompt_dict_error = ''
+    if prompt_type == PromptType.custom.name and not isinstance(prompt_dict, dict):
+        try:
+            prompt_dict = ast.literal_eval(prompt_dict)
+        except BaseException as e:
+            prompt_dict_error = str(e)
+        if prompt_dict_error:
+            return dict(), prompt_dict_error
+
+    if prompt_type in [PromptType.custom.value, str(PromptType.custom.value),
+                       PromptType.custom.name]:
+        promptA = prompt_dict.get('promptA', '')
+        promptB = prompt_dict('promptB', '')
+        PreInstruct = prompt_dict.get('PreInstruct', '')
+        PreInput = prompt_dict.get('PreInput', '')
+        PreResponse = prompt_dict.get('PreResponse', '')
+        terminate_response = prompt_dict.get('terminate_response', None)
+        chat_sep = prompt_dict.get('chat_sep', '\n')
+        humanstr = prompt_dict.get('humanstr', '')
+        botstr = prompt_dict.get('botstr', '')
+    elif prompt_type in [PromptType.plain.value, str(PromptType.plain.value),
+                         PromptType.plain.name]:
+        promptA = promptB = PreInstruct = PreInput = PreResponse = ''
+        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
+
+    elif prompt_type in [PromptType.instruct_simple.value, str(PromptType.instruct_simple.value),
+                         PromptType.instruct_simple.name]:
+        promptA = '' if not (chat and reduced) else ''
+        promptB = '' if not (chat and reduced) else ''
+
+        PreInstruct = """
+### Instruction:
+"""
+
+        PreInput = """
+### Input:
+"""
+
+        PreResponse = """
+### Response:
+"""
+        terminate_response = None
+        chat_sep = '\n'
+        humanstr = PreInstruct
+        botstr = PreResponse
+    else:
+        raise RuntimeError("No such prompt_type=%s" % prompt_type)
+
+    if return_dict:
+        return dict(promptA=promptA, promptB=promptB, PreInstruct=PreInstruct, PreInput=PreInput,
+                    PreResponse=PreResponse, terminate_response=terminate_response, chat_sep=chat_sep,
+                    humanstr=humanstr, botstr=botstr), ''
+    else:
+        return promptA, promptB, PreInstruct, PreInput, PreResponse, terminate_response, chat_sep, humanstr, botstr
+
+
+def generate_prompt(data_point, prompt_type, prompt_dict, 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)
+    prompt_dict = data_point.get('prompt_dict', prompt_dict)
+    assert prompt_type in prompt_types, "Bad prompt type: %s" % prompt_type
+    promptA, promptB, PreInstruct, PreInput, PreResponse, \
+        terminate_response, chat_sep, humanstr, botstr = get_prompt(prompt_type, prompt_dict, 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, prompt_dict, debug=False, chat=False, stream_output=False, repeat_penalty=True,
+                 allowed_repeat_line_length=10):
+        self.prompt_type = prompt_type
+        self.prompt_dict = prompt_dict
+        data_point = dict(instruction='', input='', output='')
+        _, self.pre_response, self.terminate_response, self.chat_sep = \
+            generate_prompt(data_point, self.prompt_type, self.prompt_dict, 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(self.prompt_type, self.prompt_dict, chat, context, reduced)
+
+    def generate_prompt(self, data_point):
+        reduced = False
+        prompt, _, _, _ = generate_prompt(data_point, self.prompt_type, self.prompt_dict, 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

requirements.txt

@@ -0,0 +1,152 @@
+# 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
+instructorembedding==1.0.1
+
+# for gpt4all .env file, but avoid worrying about imports
+python-dotenv==1.0.0# optional for chat with PDF
+langchain==0.0.193
+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# optional for chat with PDF
+langchain==0.0.193
+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.2faiss-gpu==1.7.2
+gpt4all==0.2.3
+llama-cpp-python==0.1.55
+arxiv==1.4.7
+pymupdf==1.22.3 # AGPL license
+# extract-msg==0.41.1  # GPL3

stopping.py

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

stopping.py

@@ -0,0 +1,73 @@
+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, prompt_dict, tokenizer, device, human='<human>:', bot="<bot>:"):
+    # FIXME: prompt_dict unused currently
+    if prompt_type in [PromptType.human_bot.name, PromptType.instruct_vicuna.name, PromptType.instruct_with_end.name]:
+        if prompt_type == PromptType.human_bot.name:
+            # encounters = [prompt.count(human) + 1, prompt.count(bot) + 1]
+            # stopping only starts once output is beyond prompt
+            # 1 human is enough to trigger, but need 2 bots, because very first view back will be bot we added
+            stop_words = [human, bot, '\n' + human, '\n' + bot]
+            encounters = [1, 2]
+        elif prompt_type == PromptType.instruct_vicuna.name:
+            # even below is not enough, generic strings and many ways to encode
+            stop_words = [
+                '### Human:',
+                """
+### Human:""",
+                """
+### Human:
+""",
+                '### Assistant:',
+                """
+### Assistant:""",
+                """
+### Assistant:
+""",
+            ]
+            encounters = [1, 2]
+        else:
+            # some instruct prompts have this as end, doesn't hurt to stop on it since not common otherwise
+            stop_words = ['### End']
+            encounters = [1]
+        stop_words_ids = [
+            tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze() for stop_word in stop_words]
+        # handle single token case
+        stop_words_ids = [x if len(x.shape) > 0 else torch.tensor([x]) for x in stop_words_ids]
+        stop_words_ids = [x for x in stop_words_ids if x.shape[0] > 0]
+        # avoid padding in front of tokens
+        if tokenizer._pad_token:  # use hidden variable to avoid annoying properly logger bug
+            stop_words_ids = [x[1:] if x[0] == tokenizer.pad_token_id and len(x) > 1 else x for x in stop_words_ids]
+        # handle fake \n added
+        stop_words_ids = [x[1:] if y[0] == '\n' else x for x, y in zip(stop_words_ids, stop_words)]
+        # build stopper
+        stopping_criteria = StoppingCriteriaList(
+            [StoppingCriteriaSub(stops=stop_words_ids, encounters=encounters, device=device)])
+    else:
+        stopping_criteria = StoppingCriteriaList()
+    return stopping_criteria

utils.py

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

utils.py

@@ -0,0 +1,856 @@
+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
+from enum import Enum
+
+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()
+
+
+def start_faulthandler():
+    # If hit server or any subprocess with signal SIGUSR1, it'll print out all threads stack trace, but wont't quit or coredump
+    # If more than one fork tries to write at same time, then looks corrupted.
+    import faulthandler
+    import signal
+
+    # SIGUSR1 in h2oai/__init__.py as well
+    faulthandler.enable()
+    faulthandler.register(signal.SIGUSR1)