Upload folder using huggingface_hub

.gitattributes

@@ -25,7 +25,6 @@
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text

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.

README.md

@@ -1,12 +1,12 @@
 ---
-title: Test
-emoji: 📊
+title: test
 colorFrom: yellow
 colorTo: yellow
 sdk: gradio
-sdk_version: 5.9.1
+sdk_version: 3.41.2
 app_file: app.py
 pinned: false
+license: apache-2.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

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

generate.py

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

gradio_utils/css.py

@@ -0,0 +1,148 @@
+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 """
+    #col_container {margin-left: auto; margin-right: auto; text-align: left;}
+
+    @import url('https://fonts.googleapis.com/css2?family=Source+Sans+Pro:wght@400;600&display=swap');
+    
+    body.dark{#warning {background-color: #555555};}
+    
+    #sidebar {
+        order: 1;
+        
+        @media (max-width: 463px) {
+          order: 2;
+        }
+    }
+    
+    #col-tabs {
+        order: 2;
+        
+        @media (max-width: 463px) {
+          order: 1;
+        }
+    }
+    
+    #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;
+    }
+    .chatsmall chatbot {font-size: 10px !important}
+    
+    .gradio-container {
+        max-width: none !important;
+    }
+    
+    div.message {
+        padding: var(--text-lg) !important;
+    }
+    
+    div.message.user > div.icon-button {
+        top: unset;
+        bottom: 0;
+    }
+    
+    div.message.bot > div.icon-button {
+        top: unset;
+        bottom: 0;
+    }
+    
+    #prompt-form-row {
+        position: relative;
+    }
+    
+    #attach-button {
+        position: absolute;
+        top: 45px;
+        right: 20px;
+        
+        display: flex;
+        justify-content: center;
+        border: 1px solid var(--primary-500) !important;
+        
+        @media (max-width: 463px) {
+          width: 56px;
+        }
+    }
+    
+    #attach-button > img {
+        margin-right: 0;
+    }
+    
+    #prompt-form > label > textarea {
+        padding-right: 104px;
+        
+        @media (max-width: 463px) {
+          min-height: 94px;
+          padding-right: 70px;
+        }
+    }
+
+    #visible-models > label > div.wrap > div.wrap-inner > div.secondary-wrap > div.remove-all {
+        display: none !important;
+    }
+    
+    #visible-models > label > div.wrap > div.wrap-inner > div.token {
+        display: none !important;
+    }
+    
+    #visible-models > label > div.wrap > div.wrap-inner > div.secondary-wrap::before {
+        content: "Select";
+        padding: 0 4px;
+        margin-right: 2px;
+    }
+
+    #langchain_agents > label > div.wrap > div.wrap-inner > div.secondary-wrap > div.remove-all {
+        display: none !important;
+    }
+
+    #langchain_agents > label > div.wrap > div.wrap-inner > div.token {
+        display: none !important;
+    }
+
+    #langchain_agents > label > div.wrap > div.wrap-inner > div.secondary-wrap::before {
+        content: "Select";
+        padding: 0 4px;
+        margin-right: 2px;
+    }
+    """

gradio_utils/grclient.py

@@ -0,0 +1,82 @@
+import traceback
+from typing import Callable
+import os
+
+from gradio_client.client import Job
+
+os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
+
+from gradio_client import Client
+
+
+class GradioClient(Client):
+    """
+    Parent class of gradio client
+    To handle automatically refreshing client if detect gradio server changed
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+        super().__init__(*args, **kwargs)
+        self.server_hash = self.get_server_hash()
+
+    def get_server_hash(self):
+        """
+        Get server hash using super without any refresh action triggered
+        Returns: git hash of gradio server
+        """
+        return super().submit(api_name='/system_hash').result()
+
+    def refresh_client_if_should(self):
+        # get current hash in order to update api_name -> fn_index map in case gradio server changed
+        # FIXME: Could add cli api as hash
+        server_hash = self.get_server_hash()
+        if self.server_hash != server_hash:
+            self.refresh_client()
+            self.server_hash = server_hash
+        else:
+            self.reset_session()
+
+    def refresh_client(self):
+        """
+        Ensure every client call is independent
+        Also ensure map between api_name and fn_index is updated in case server changed (e.g. restarted with new code)
+        Returns:
+        """
+        # need session hash to be new every time, to avoid "generator already executing"
+        self.reset_session()
+
+        client = Client(*self.args, **self.kwargs)
+        for k, v in client.__dict__.items():
+            setattr(self, k, v)
+
+    def submit(
+        self,
+        *args,
+        api_name: str | None = None,
+        fn_index: int | None = None,
+        result_callbacks: Callable | list[Callable] | None = None,
+    ) -> Job:
+        # Note predict calls submit
+        try:
+            self.refresh_client_if_should()
+            job = super().submit(*args, api_name=api_name, fn_index=fn_index)
+        except Exception as e:
+            print("Hit e=%s" % str(e), flush=True)
+            # force reconfig in case only that
+            self.refresh_client()
+            job = super().submit(*args, api_name=api_name, fn_index=fn_index)
+
+        # see if immediately failed
+        e = job.future._exception
+        if e is not None:
+            print("GR job failed: %s %s" % (str(e), ''.join(traceback.format_tb(e.__traceback__))), flush=True)
+            # force reconfig in case only that
+            self.refresh_client()
+            job = super().submit(*args, api_name=api_name, fn_index=fn_index)
+            e2 = job.future._exception
+            if e2 is not None:
+                print("GR job failed again: %s\n%s" % (str(e2), ''.join(traceback.format_tb(e2.__traceback__))), flush=True)
+
+        return job

gradio_utils/prompt_form.py

@@ -0,0 +1,108 @@
+import os
+import math
+
+import gradio as gr
+
+
+def make_chatbots(output_label0, output_label0_model2, **kwargs):
+    visible_models = kwargs['visible_models']
+    all_models = kwargs['all_models']
+
+    text_outputs = []
+    chat_kwargs = []
+    for model_state_locki, model_state_lock in enumerate(kwargs['model_states']):
+        if os.environ.get('DEBUG_MODEL_LOCK'):
+            model_name = model_state_lock["base_model"] + " : " + model_state_lock["inference_server"]
+        else:
+            model_name = model_state_lock["base_model"]
+        output_label = f'h2oGPT [{model_name}]'
+        min_width = 250 if kwargs['gradio_size'] in ['small', 'large', 'medium'] else 160
+        chat_kwargs.append(dict(label=output_label, elem_classes='chatsmall',
+                                height=kwargs['height'] or 400, min_width=min_width,
+                                show_copy_button=kwargs['show_copy_button'],
+                                visible=kwargs['model_lock'] and (visible_models is None or
+                                                                  model_state_locki in visible_models or
+                                                                  all_models[model_state_locki] in visible_models
+                                                                  )))
+
+    # base view on initial visible choice
+    if visible_models:
+        len_visible = len(visible_models)
+    else:
+        len_visible = len(kwargs['model_states'])
+    if kwargs['model_lock_columns'] == -1:
+        kwargs['model_lock_columns'] = len_visible
+    if kwargs['model_lock_columns'] is None:
+        kwargs['model_lock_columns'] = 3
+
+    ncols = kwargs['model_lock_columns']
+    if kwargs['model_states'] == 0:
+        nrows = 0
+    else:
+        nrows = math.ceil(len_visible / kwargs['model_lock_columns'])
+
+    if kwargs['model_lock_columns'] == 0:
+        # not using model_lock
+        pass
+    elif nrows <= 1:
+        with gr.Row():
+            for chat_kwargs1, model_state_lock in zip(chat_kwargs, kwargs['model_states']):
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+    elif nrows == kwargs['model_states']:
+        with gr.Row():
+            for chat_kwargs1, model_state_lock in zip(chat_kwargs, kwargs['model_states']):
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+    elif nrows == 2:
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii >= len_visible / 2:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < len_visible / 2:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+    elif nrows == 3:
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii >= 1 * len_visible / 3:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 1 * len_visible / 3 or mii >= 2 * len_visible / 3:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 2 * len_visible / 3:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+    elif nrows >= 4:
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii >= 1 * len_visible / 4:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 1 * len_visible / 4 or mii >= 2 * len_visible / 4:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 2 * len_visible / 4 or mii >= 3 * len_visible / 4:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 3 * len_visible / 4:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+
+    with gr.Row():
+        text_output = gr.Chatbot(label=output_label0, visible=not kwargs['model_lock'], height=kwargs['height'] or 400)
+        text_output2 = gr.Chatbot(label=output_label0_model2,
+                                  visible=False and not kwargs['model_lock'], height=kwargs['height'] or 400)
+    return text_output, text_output2, text_outputs

iterators/__init__.py

@@ -0,0 +1,4 @@
+from .timeout_iterator import TimeoutIterator, AsyncTimeoutIterator
+from .iterator_pipe import IteratorPipe, AsyncIteratorPipe
+
+__all__ = ["TimeoutIterator", "AsyncTimeoutIterator", "IteratorPipe", "AsyncIteratorPipe"]

iterators/iterator_pipe.py

@@ -0,0 +1,93 @@
+import queue
+import asyncio
+
+
+class IteratorPipe:
+    """
+    Iterator Pipe creates an iterator that can be fed in data from another block of code or thread of execution
+    """
+
+    def __init__(self, sentinel=object()):
+        self._q = queue.Queue()
+        self._sentinel = sentinel
+        self._sentinel_pushed = False
+        self._closed = False
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._closed:
+            raise StopIteration
+
+        data = self._q.get(block=True)
+        if data is self._sentinel:
+            self._closed = True
+            raise StopIteration
+
+        return data
+
+    def put(self, data) -> bool:
+        """
+        Pushes next item to Iterator and returns True
+        If iterator has been closed via close(), doesn't push anything and returns False
+        """
+        if self._sentinel_pushed:
+            return False
+
+        self._q.put(data)
+        return True
+
+    def close(self):
+        """
+        Close is idempotent. Calling close multiple times is safe
+        Iterator will raise StopIteration only after all elements pushed before close have been iterated
+        """
+        # make close idempotent
+        if not self._sentinel_pushed:
+            self._sentinel_pushed = True
+        self._q.put(self._sentinel)
+
+
+class AsyncIteratorPipe:
+
+    def __init__(self, sentinel=object()):
+        self._q = asyncio.Queue()
+        self._sentinel = sentinel
+        self._sentinel_pushed = False
+        self._closed = False
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self._closed:
+            raise StopAsyncIteration
+
+        data = await self._q.get()
+        if data is self._sentinel:
+            self._closed = True
+            raise StopAsyncIteration
+
+        return data
+
+    async def put(self, data) -> bool:
+        """
+        Pushes next item to Iterator and returns True
+        If iterator has been closed via close(), doesn't push anything and returns False
+        """
+        if self._sentinel_pushed:
+            return False
+
+        await self._q.put(data)
+        return True
+
+    async def close(self):
+        """
+        Close is idempotent. Calling close multiple times is safe
+        Iterator will raise StopIteration only after all elements pushed before close have been iterated
+        """
+        # make close idempotent
+        if not self._sentinel_pushed:
+            self._sentinel_pushed = True
+            await self._q.put(self._sentinel)

iterators/timeout_iterator.py

@@ -0,0 +1,170 @@
+import queue
+import asyncio
+import threading
+import traceback
+
+
+class TimeoutIterator:
+    """
+    Wrapper class to add timeout feature to synchronous iterators
+    - timeout: timeout for next(). Default=ZERO_TIMEOUT i.e. no timeout or blocking calls to next. Updated using set_timeout() 
+    - sentinel: the object returned by iterator when timeout happens
+    - reset_on_next: if set to True, timeout is reset to the value of ZERO_TIMEOUT on each iteration
+
+    TimeoutIterator uses a thread internally.
+    The thread stops once the iterator exhausts or raises an exception during iteration.
+
+    Any exceptions raised within the wrapped iterator are propagated as it is.
+    Exception is raised when all elements generated by the actual iterator before exception have been consumed
+    Timeout can be set dynamically before going for iteration
+    """
+    ZERO_TIMEOUT = 0.0
+
+    def __init__(self, iterator, timeout=0.0, sentinel=object(), reset_on_next=False, raise_on_exception=True):
+        self._iterator = iterator
+        self._timeout = timeout
+        self._sentinel = sentinel
+        self._reset_on_next = reset_on_next
+        self._raise_on_exception = raise_on_exception
+
+        self._interrupt = False
+        self._done = False
+        self._buffer = queue.Queue()
+        self._thread = threading.Thread(target=self.__lookahead)
+        self._thread.start()
+
+    def get_sentinel(self):
+        return self._sentinel
+
+    def set_reset_on_next(self, reset_on_next):
+        self._reset_on_next = reset_on_next
+
+    def set_timeout(self, timeout: float):
+        """
+        Set timeout for next iteration
+        """
+        self._timeout = timeout
+
+    def interrupt(self):
+        """
+        interrupt and stop the underlying thread.
+        the thread actually dies only after interrupt has been set and
+        the underlying iterator yields a value after that.
+        """
+        self._interrupt = True
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        """
+        yield the result from iterator
+        if timeout > 0:
+            yield data if available.
+            otherwise yield sentinal
+        """
+        if self._done:
+            raise StopIteration
+
+        data = self._sentinel
+        try:
+            if self._timeout > self.ZERO_TIMEOUT:
+                data = self._buffer.get(timeout=self._timeout)
+            else:
+                data = self._buffer.get()
+        except queue.Empty:
+            pass
+        finally:
+            # see if timeout needs to be reset
+            if self._reset_on_next:
+                self._timeout = self.ZERO_TIMEOUT
+
+        # propagate any exceptions including StopIteration
+        if isinstance(data, BaseException):
+            self._done = True
+            if isinstance(data, StopIteration):
+                raise data
+            ex = ''.join(traceback.format_tb(data.__traceback__))
+            print("Generation Failed: %s %s" % (str(data), str(ex)), flush=True)
+            if self._raise_on_exception:
+                raise data
+            else:
+                return data
+
+        return data
+
+    def __lookahead(self):
+        try:
+            while True:
+                self._buffer.put(next(self._iterator))
+                if self._interrupt:
+                    raise StopIteration()
+        except BaseException as e:
+            self._buffer.put(e)
+
+
+class AsyncTimeoutIterator:
+    """
+    Async version of TimeoutIterator. See method documentation of TimeoutIterator
+    """
+    ZERO_TIMEOUT = 0.0
+
+    def __init__(self, iterator, timeout=0.0, sentinel=object(), reset_on_next=False):
+        self._iterator = iterator
+        self._timeout = timeout
+        self._sentinel = sentinel
+        self._reset_on_next = reset_on_next
+
+        self._interrupt = False
+        self._done = False
+        self._buffer = asyncio.Queue()
+        self._task = asyncio.get_event_loop().create_task(self.__lookahead())
+
+    def get_sentinel(self):
+        return self._sentinel
+
+    def set_reset_on_next(self, reset_on_next):
+        self._reset_on_next = reset_on_next
+
+    def set_timeout(self, timeout: float):
+        self._timeout = timeout
+
+    def interrupt(self):
+        self._interrupt = True
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self._done:
+            raise StopAsyncIteration
+
+        data = self._sentinel
+        try:
+            if self._timeout > self.ZERO_TIMEOUT:
+                data = await asyncio.wait_for(self._buffer.get(), self._timeout)
+            else:
+                data = await self._buffer.get()
+        except asyncio.TimeoutError:
+            pass
+        finally:
+            # see if timeout needs to be reset
+            if self._reset_on_next:
+                self._timeout = self.ZERO_TIMEOUT
+
+        # propagate any exceptions including StopIteration
+        if isinstance(data, BaseException):
+            self._done = True
+            raise data
+
+        return data
+
+    async def __lookahead(self):
+        try:
+            while True:
+                data = await self._iterator.__anext__()
+                await self._buffer.put(data)
+                if self._interrupt:
+                    raise StopAsyncIteration()
+        except BaseException as e:
+            await self._buffer.put(e)

requirements.txt

@@ -0,0 +1,192 @@
+# for generate (gradio server) and finetune
+datasets==2.13.0
+sentencepiece==0.1.99
+gradio==3.41.2
+huggingface_hub==0.16.4
+appdirs==1.4.4
+fire==0.5.0
+docutils==0.20.1
+torch==2.0.1; sys_platform != "darwin" and platform_machine != "arm64"
+evaluate==0.4.0
+rouge_score==0.1.2
+sacrebleu==2.3.1
+scikit-learn==1.2.2
+# optional (need to uncomment code in gradio_runner.py for import of better_profanity)
+# alt-profanity-check==1.2.2
+# better-profanity==0.7.0
+numpy==1.24.3
+pandas==2.0.2
+matplotlib==3.7.1
+loralib==0.1.1
+bitsandbytes==0.41.1
+accelerate==0.22.0
+peft==0.5.0
+transformers==4.33.1
+tokenizers==0.13.3
+APScheduler==3.10.1
+
+# optional for generate
+pynvml==11.5.0
+psutil==5.9.5
+boto3==1.26.101
+botocore==1.29.101
+
+# optional for finetune
+tensorboard==2.13.0
+neptune==1.2.0
+
+# for gradio client
+gradio_client==0.5.0
+beautifulsoup4==4.12.2
+markdown==3.4.3
+
+# 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; sys_platform == "darwin" and platform_machine == "arm64"
+pypandoc_binary==1.11; platform_machine == "x86_64"
+pypandoc_binary==1.11; sys_platform == "win32"
+python-magic-bin==0.4.14; sys_platform == "win32"
+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
+
+text-generation==0.6.0
+# for tokenization when don't have HF tokenizer
+tiktoken==0.4.0
+
+requests>=2.31.0
+urllib3>=1.26.16
+filelock>=3.12.2
+joblib>=1.3.1
+tqdm>=4.65.0
+tabulate>=0.9.0
+packaging>=23.1
+# optional for chat with PDF
+langchain==0.0.300
+pypdf==3.14.0
+# avoid textract, requires old six
+#textract==1.6.5
+pypdfium2==4.19.0
+
+# for HF embeddings
+sentence_transformers==2.2.2
+
+# optional: for OpenAI endpoint or embeddings (requires key)
+openai==0.27.8
+replicate==0.10.0
+
+# local vector db
+chromadb==0.4.10
+
+# chroma migration
+chroma-migrate==0.0.7
+duckdb==0.7.1
+https://h2o-release.s3.amazonaws.com/h2ogpt/chromamigdb-0.3.25-py3-none-any.whl
+https://h2o-release.s3.amazonaws.com/h2ogpt/hnswmiglib-0.7.0.tgz
+
+# 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.8.1
+
+# strong support for images
+# Requires on Ubuntu: sudo apt-get install libmagic-dev poppler-utils tesseract-ocr libtesseract-dev libreoffice
+unstructured[local-inference]==0.9.0
+#pdf2image==1.16.3
+#pytesseract==0.3.10
+pillow==9.5.0
+posthog==3.0.1
+
+pdfminer.six==20221105
+urllib3
+requests_file
+
+#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
+# TRY: apt-get install autoconf libtool
+# unclear what happens on windows/mac for now
+jq==1.4.1; platform_machine == "x86_64"
+
+# to check licenses
+# Run: pip-licenses|grep -v 'BSD\|Apache\|MIT'
+pip-licenses==4.3.0
+
+# weaviate vector db
+weaviate-client==3.22.1
+# optional for chat with PDF
+langchain==0.0.300
+pypdf==3.14.0
+# avoid textract, requires old six
+#textract==1.6.5
+pypdfium2==4.19.0
+
+# for HF embeddings
+sentence_transformers==2.2.2
+
+# optional: for OpenAI endpoint or embeddings (requires key)
+openai==0.27.8
+replicate==0.10.0
+
+# local vector db
+chromadb==0.4.10
+
+# chroma migration
+chroma-migrate==0.0.7
+duckdb==0.7.1
+https://h2o-release.s3.amazonaws.com/h2ogpt/chromamigdb-0.3.25-py3-none-any.whl
+https://h2o-release.s3.amazonaws.com/h2ogpt/hnswmiglib-0.7.0.tgz
+
+# 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.8.1
+
+# strong support for images
+# Requires on Ubuntu: sudo apt-get install libmagic-dev poppler-utils tesseract-ocr libtesseract-dev libreoffice
+unstructured[local-inference]==0.9.0
+#pdf2image==1.16.3
+#pytesseract==0.3.10
+pillow==9.5.0
+posthog==3.0.1
+
+pdfminer.six==20221105
+urllib3
+requests_file
+
+#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
+# TRY: apt-get install autoconf libtool
+# unclear what happens on windows/mac for now
+jq==1.4.1; platform_machine == "x86_64"
+
+# to check licenses
+# Run: pip-licenses|grep -v 'BSD\|Apache\|MIT'
+pip-licenses==4.3.0
+
+# weaviate vector db
+weaviate-client==3.22.1
+faiss-gpu==1.7.2
+arxiv==1.4.8
+pymupdf==1.23.1 # AGPL license
+# extract-msg==0.41.1  # GPL3

src/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.

src/client_test.py

@@ -0,0 +1,484 @@
+"""
+Client test.
+
+Run server:
+
+python generate.py  --base_model=h2oai/h2ogpt-oig-oasst1-512-6_9b
+
+NOTE: For private models, add --use-auth_token=True
+
+NOTE: --use_gpu_id=True (default) must be used for multi-GPU in case see failures with cuda:x cuda:y mismatches.
+Currently, this will force model to be on a single GPU.
+
+Then run this client as:
+
+python src/client_test.py
+
+
+
+For HF spaces:
+
+HOST="https://h2oai-h2ogpt-chatbot.hf.space" python src/client_test.py
+
+Result:
+
+Loaded as API: https://h2oai-h2ogpt-chatbot.hf.space ✔
+{'instruction_nochat': 'Who are you?', 'iinput_nochat': '', 'response': 'I am h2oGPT, a large language model developed by LAION.', 'sources': ''}
+
+
+For demo:
+
+HOST="https://gpt.h2o.ai" python src/client_test.py
+
+Result:
+
+Loaded as API: https://gpt.h2o.ai ✔
+{'instruction_nochat': 'Who are you?', 'iinput_nochat': '', 'response': 'I am h2oGPT, a chatbot created by LAION.', 'sources': ''}
+
+NOTE: Raw output from API for nochat case is a string of a python dict and will remain so if other entries are added to dict:
+
+{'response': "I'm h2oGPT, a large language model by H2O.ai, the visionary leader in democratizing AI.", 'sources': ''}
+
+
+"""
+import ast
+import time
+import os
+import markdown  # pip install markdown
+import pytest
+from bs4 import BeautifulSoup  # pip install beautifulsoup4
+
+try:
+    from enums import DocumentSubset, LangChainAction
+except:
+    from src.enums import DocumentSubset, LangChainAction
+
+from tests.utils import get_inf_server
+
+debug = False
+
+os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
+
+
+def get_client(serialize=True):
+    from gradio_client import Client
+
+    client = Client(get_inf_server(), serialize=serialize)
+    if debug:
+        print(client.view_api(all_endpoints=True))
+    return client
+
+
+def get_args(prompt, prompt_type=None, chat=False, stream_output=False,
+             max_new_tokens=50,
+             top_k_docs=3,
+             langchain_mode='Disabled',
+             add_chat_history_to_context=True,
+             langchain_action=LangChainAction.QUERY.value,
+             langchain_agents=[],
+             prompt_dict=None,
+             version=None,
+             h2ogpt_key=None,
+             visible_models=None,
+             system_prompt='',  # default of no system prompt tiggered by empty string
+             add_search_to_context=False,
+             chat_conversation=None,
+             text_context_list=None,
+             ):
+    from collections import OrderedDict
+    kwargs = OrderedDict(instruction=prompt if chat else '',  # only for chat=True
+                         iinput='',  # only for chat=True
+                         context='',
+                         # streaming output is supported, loops over and outputs each generation in streaming mode
+                         # but leave stream_output=False for simple input/output mode
+                         stream_output=stream_output,
+                         prompt_type=prompt_type,
+                         prompt_dict=prompt_dict,
+                         temperature=0.1,
+                         top_p=0.75,
+                         top_k=40,
+                         num_beams=1,
+                         max_new_tokens=max_new_tokens,
+                         min_new_tokens=0,
+                         early_stopping=False,
+                         max_time=20,
+                         repetition_penalty=1.0,
+                         num_return_sequences=1,
+                         do_sample=True,
+                         chat=chat,
+                         instruction_nochat=prompt if not chat else '',
+                         iinput_nochat='',  # only for chat=False
+                         langchain_mode=langchain_mode,
+                         add_chat_history_to_context=add_chat_history_to_context,
+                         langchain_action=langchain_action,
+                         langchain_agents=langchain_agents,
+                         top_k_docs=top_k_docs,
+                         chunk=True,
+                         chunk_size=512,
+                         document_subset=DocumentSubset.Relevant.name,
+                         document_choice=[],
+                         pre_prompt_query=None,
+                         prompt_query=None,
+                         pre_prompt_summary=None,
+                         prompt_summary=None,
+                         system_prompt=system_prompt,
+                         image_loaders=None,
+                         pdf_loaders=None,
+                         url_loaders=None,
+                         jq_schema=None,
+                         visible_models=visible_models,
+                         h2ogpt_key=h2ogpt_key,
+                         add_search_to_context=add_search_to_context,
+                         chat_conversation=chat_conversation,
+                         text_context_list=text_context_list,
+                         docs_ordering_type=None,
+                         min_max_new_tokens=None,
+                         )
+    diff = 0
+    if version is None:
+        # latest
+        version = 1
+    if version == 0:
+        diff = 1
+    if version >= 1:
+        kwargs.update(dict(system_prompt=system_prompt))
+        diff = 0
+
+    from evaluate_params import eval_func_param_names
+    assert len(set(eval_func_param_names).difference(set(list(kwargs.keys())))) == diff
+    if chat:
+        # add chatbot output on end.  Assumes serialize=False
+        kwargs.update(dict(chatbot=[]))
+
+    return kwargs, list(kwargs.values())
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic(prompt_type='human_bot', version=None, visible_models=None, prompt='Who are you?',
+                      h2ogpt_key=None):
+    return run_client_nochat(prompt=prompt, prompt_type=prompt_type, max_new_tokens=50, version=version,
+                             visible_models=visible_models, h2ogpt_key=h2ogpt_key)
+
+
+"""
+time HOST=https://gpt-internal.h2o.ai PYTHONPATH=. pytest -n 20 src/client_test.py::test_client_basic_benchmark
+32 seconds to answer 20 questions at once with 70B llama2 on 4x A100 80GB using TGI 0.9.3
+"""
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+@pytest.mark.parametrize("id", range(20))
+def test_client_basic_benchmark(id, prompt_type='human_bot', version=None):
+    return run_client_nochat(prompt="""
+/nfs4/llm/h2ogpt/h2ogpt/bin/python /home/arno/pycharm-2022.2.2/plugins/python/helpers/pycharm/_jb_pytest_runner.py --target src/client_test.py::test_client_basic
+Testing started at 8:41 AM ...
+Launching pytest with arguments src/client_test.py::test_client_basic --no-header --no-summary -q in /nfs4/llm/h2ogpt
+
+============================= test session starts ==============================
+collecting ...
+src/client_test.py:None (src/client_test.py)
+ImportError while importing test module '/nfs4/llm/h2ogpt/src/client_test.py'.
+Hint: make sure your test modules/packages have valid Python names.
+Traceback:
+h2ogpt/lib/python3.10/site-packages/_pytest/python.py:618: in _importtestmodule
+    mod = import_path(self.path, mode=importmode, root=self.config.rootpath)
+h2ogpt/lib/python3.10/site-packages/_pytest/pathlib.py:533: in import_path
+    importlib.import_module(module_name)
+/usr/lib/python3.10/importlib/__init__.py:126: in import_module
+    return _bootstrap._gcd_import(name[level:], package, level)
+<frozen importlib._bootstrap>:1050: in _gcd_import
+    ???
+<frozen importlib._bootstrap>:1027: in _find_and_load
+    ???
+<frozen importlib._bootstrap>:1006: in _find_and_load_unlocked
+    ???
+<frozen importlib._bootstrap>:688: in _load_unlocked
+    ???
+h2ogpt/lib/python3.10/site-packages/_pytest/assertion/rewrite.py:168: in exec_module
+    exec(co, module.__dict__)
+src/client_test.py:51: in <module>
+    from enums import DocumentSubset, LangChainAction
+E   ModuleNotFoundError: No module named 'enums'
+
+
+collected 0 items / 1 error
+
+=============================== 1 error in 0.14s ===============================
+ERROR: not found: /nfs4/llm/h2ogpt/src/client_test.py::test_client_basic
+(no name '/nfs4/llm/h2ogpt/src/client_test.py::test_client_basic' in any of [<Module client_test.py>])
+
+
+Process finished with exit code 4
+
+What happened?
+""", prompt_type=prompt_type, max_new_tokens=100, version=version)
+
+
+def run_client_nochat(prompt, prompt_type, max_new_tokens, version=None, h2ogpt_key=None, visible_models=None):
+    kwargs, args = get_args(prompt, prompt_type, chat=False, max_new_tokens=max_new_tokens, version=version,
+                            visible_models=visible_models, h2ogpt_key=h2ogpt_key)
+
+    api_name = '/submit_nochat'
+    client = get_client(serialize=True)
+    res = client.predict(
+        *tuple(args),
+        api_name=api_name,
+    )
+    print("Raw client result: %s" % res, flush=True)
+    res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
+                    response=md_to_text(res))
+    print(res_dict)
+    return res_dict, client
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic_api(prompt_type='human_bot', version=None, h2ogpt_key=None):
+    return run_client_nochat_api(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50, version=version,
+                                 h2ogpt_key=h2ogpt_key)
+
+
+def run_client_nochat_api(prompt, prompt_type, max_new_tokens, version=None, h2ogpt_key=None):
+    kwargs, args = get_args(prompt, prompt_type, chat=False, max_new_tokens=max_new_tokens, version=version,
+                            h2ogpt_key=h2ogpt_key)
+
+    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
+    client = get_client(serialize=True)
+    res = client.predict(
+        str(dict(kwargs)),
+        api_name=api_name,
+    )
+    print("Raw client result: %s" % res, flush=True)
+    res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
+                    response=md_to_text(ast.literal_eval(res)['response']),
+                    sources=ast.literal_eval(res)['sources'])
+    print(res_dict)
+    return res_dict, client
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic_api_lean(prompt_type='human_bot', version=None, h2ogpt_key=None):
+    return run_client_nochat_api_lean(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50,
+                                      version=version, h2ogpt_key=h2ogpt_key)
+
+
+def run_client_nochat_api_lean(prompt, prompt_type, max_new_tokens, version=None, h2ogpt_key=None):
+    kwargs = dict(instruction_nochat=prompt, h2ogpt_key=h2ogpt_key)
+
+    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'],
+                    h2ogpt_key=h2ogpt_key)
+    print(res_dict)
+    return res_dict, client
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic_api_lean_morestuff(prompt_type='human_bot', version=None, h2ogpt_key=None):
+    return run_client_nochat_api_lean_morestuff(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50,
+                                                version=version, h2ogpt_key=h2ogpt_key)
+
+
+def run_client_nochat_api_lean_morestuff(prompt, prompt_type='human_bot', max_new_tokens=512, version=None,
+                                         h2ogpt_key=None):
+    kwargs = dict(
+        instruction='',
+        iinput='',
+        context='',
+        stream_output=False,
+        prompt_type=prompt_type,
+        temperature=0.1,
+        top_p=0.75,
+        top_k=40,
+        num_beams=1,
+        max_new_tokens=1024,
+        min_new_tokens=0,
+        early_stopping=False,
+        max_time=20,
+        repetition_penalty=1.0,
+        num_return_sequences=1,
+        do_sample=True,
+        chat=False,
+        instruction_nochat=prompt,
+        iinput_nochat='',
+        langchain_mode='Disabled',
+        add_chat_history_to_context=True,
+        langchain_action=LangChainAction.QUERY.value,
+        langchain_agents=[],
+        top_k_docs=4,
+        document_subset=DocumentSubset.Relevant.name,
+        document_choice=[],
+        h2ogpt_key=h2ogpt_key,
+        add_search_to_context=False,
+    )
+
+    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'],
+                    h2ogpt_key=h2ogpt_key)
+    print(res_dict)
+    return res_dict, client
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_chat(prompt_type='human_bot', version=None, h2ogpt_key=None):
+    return run_client_chat(prompt='Who are you?', prompt_type=prompt_type, stream_output=False, max_new_tokens=50,
+                           langchain_mode='Disabled',
+                           langchain_action=LangChainAction.QUERY.value,
+                           langchain_agents=[],
+                           version=version,
+                           h2ogpt_key=h2ogpt_key)
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_chat_stream(prompt_type='human_bot', version=None, h2ogpt_key=None):
+    return run_client_chat(prompt="Tell a very long kid's story about birds.", prompt_type=prompt_type,
+                           stream_output=True, max_new_tokens=512,
+                           langchain_mode='Disabled',
+                           langchain_action=LangChainAction.QUERY.value,
+                           langchain_agents=[],
+                           version=version,
+                           h2ogpt_key=h2ogpt_key)
+
+
+def run_client_chat(prompt='',
+                    stream_output=None,
+                    max_new_tokens=128,
+                    langchain_mode='Disabled',
+                    langchain_action=LangChainAction.QUERY.value,
+                    langchain_agents=[],
+                    prompt_type=None, prompt_dict=None,
+                    version=None,
+                    h2ogpt_key=None):
+    client = get_client(serialize=False)
+
+    kwargs, args = get_args(prompt, prompt_type, chat=True, stream_output=stream_output,
+                            max_new_tokens=max_new_tokens,
+                            langchain_mode=langchain_mode,
+                            langchain_action=langchain_action,
+                            langchain_agents=langchain_agents,
+                            prompt_dict=prompt_dict,
+                            version=version,
+                            h2ogpt_key=h2ogpt_key)
+    return run_client(client, prompt, args, kwargs)
+
+
+def run_client(client, prompt, args, kwargs, do_md_to_text=True, verbose=False):
+    assert kwargs['chat'], "Chat mode only"
+    res = client.predict(*tuple(args), api_name='/instruction')
+    args[-1] += [res[-1]]
+
+    res_dict = kwargs
+    res_dict['prompt'] = prompt
+    if not kwargs['stream_output']:
+        res = client.predict(*tuple(args), api_name='/instruction_bot')
+        res_dict['response'] = res[0][-1][1]
+        print(md_to_text(res_dict['response'], do_md_to_text=do_md_to_text))
+        return res_dict, client
+    else:
+        job = client.submit(*tuple(args), api_name='/instruction_bot')
+        res1 = ''
+        while not job.done():
+            outputs_list = job.communicator.job.outputs
+            if outputs_list:
+                res = job.communicator.job.outputs[-1]
+                res1 = res[0][-1][-1]
+                res1 = md_to_text(res1, do_md_to_text=do_md_to_text)
+                print(res1)
+            time.sleep(0.1)
+        full_outputs = job.outputs()
+        if verbose:
+            print('job.outputs: %s' % str(full_outputs))
+        # ensure get ending to avoid race
+        # -1 means last response if streaming
+        # 0 means get text_output, ignore exception_text
+        # 0 means get list within text_output that looks like [[prompt], [answer]]
+        # 1 means get bot answer, so will have last bot answer
+        res_dict['response'] = md_to_text(full_outputs[-1][0][0][1], do_md_to_text=do_md_to_text)
+        return res_dict, client
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_nochat_stream(prompt_type='human_bot', version=None, h2ogpt_key=None):
+    return run_client_nochat_gen(prompt="Tell a very long kid's story about birds.", prompt_type=prompt_type,
+                                 stream_output=True, max_new_tokens=512,
+                                 langchain_mode='Disabled',
+                                 langchain_action=LangChainAction.QUERY.value,
+                                 langchain_agents=[],
+                                 version=version,
+                                 h2ogpt_key=h2ogpt_key)
+
+
+def run_client_nochat_gen(prompt, prompt_type, stream_output, max_new_tokens,
+                          langchain_mode, langchain_action, langchain_agents, version=None,
+                          h2ogpt_key=None):
+    client = get_client(serialize=False)
+
+    kwargs, args = get_args(prompt, prompt_type, chat=False, stream_output=stream_output,
+                            max_new_tokens=max_new_tokens, langchain_mode=langchain_mode,
+                            langchain_action=langchain_action, langchain_agents=langchain_agents,
+                            version=version, h2ogpt_key=h2ogpt_key)
+    return run_client_gen(client, prompt, args, kwargs)
+
+
+def run_client_gen(client, prompt, args, kwargs, do_md_to_text=True, verbose=False):
+    res_dict = kwargs
+    res_dict['prompt'] = prompt
+    if not kwargs['stream_output']:
+        res = client.predict(str(dict(kwargs)), api_name='/submit_nochat_api')
+        res_dict.update(ast.literal_eval(res))
+        print(md_to_text(res_dict['response'], do_md_to_text=do_md_to_text))
+        return res_dict, client
+    else:
+        job = client.submit(str(dict(kwargs)), api_name='/submit_nochat_api')
+        while not job.done():
+            outputs_list = job.communicator.job.outputs
+            if outputs_list:
+                res = job.communicator.job.outputs[-1]
+                res_dict = ast.literal_eval(res)
+                print('Stream: %s' % res_dict['response'])
+            time.sleep(0.1)
+        res_list = job.outputs()
+        assert len(res_list) > 0, "No response, check server"
+        res = res_list[-1]
+        res_dict = ast.literal_eval(res)
+        print('Final: %s' % res_dict['response'])
+        return res_dict, client
+
+
+def md_to_text(md, do_md_to_text=True):
+    if not do_md_to_text:
+        return md
+    assert md is not None, "Markdown is None"
+    html = markdown.markdown(md)
+    soup = BeautifulSoup(html, features='html.parser')
+    return soup.get_text()
+
+
+def run_client_many(prompt_type='human_bot', version=None, h2ogpt_key=None):
+    kwargs = dict(prompt_type=prompt_type, version=version, h2ogpt_key=h2ogpt_key)
+    ret1, _ = test_client_chat(**kwargs)
+    ret2, _ = test_client_chat_stream(**kwargs)
+    ret3, _ = test_client_nochat_stream(**kwargs)
+    ret4, _ = test_client_basic(**kwargs)
+    ret5, _ = test_client_basic_api(**kwargs)
+    ret6, _ = test_client_basic_api_lean(**kwargs)
+    ret7, _ = test_client_basic_api_lean_morestuff(**kwargs)
+    return ret1, ret2, ret3, ret4, ret5, ret6, ret7
+
+
+if __name__ == '__main__':
+    run_client_many()

src/create_data.py

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

src/db_utils.py

@@ -0,0 +1,54 @@
+import uuid
+
+from enums import LangChainMode
+
+
+def set_userid(db1s, requests_state1, get_userid_auth):
+    db1 = db1s[LangChainMode.MY_DATA.value]
+    assert db1 is not None and len(db1) == length_db1()
+    if not db1[1]:
+        db1[1] = get_userid_auth(requests_state1)
+    if not db1[2]:
+        username1 = None
+        if 'username' in requests_state1:
+            username1 = requests_state1['username']
+        db1[2] = username1
+
+
+def set_userid_direct(db1s, userid, username):
+    db1 = db1s[LangChainMode.MY_DATA.value]
+    db1[1] = userid
+    db1[2] = username
+
+
+def get_userid_direct(db1s):
+    return db1s[LangChainMode.MY_DATA.value][1] if db1s is not None else ''
+
+
+def get_username_direct(db1s):
+    return db1s[LangChainMode.MY_DATA.value][2] if db1s is not None else ''
+
+
+def get_dbid(db1):
+    return db1[1]
+
+
+def set_dbid(db1):
+    # can only call this after function called so for specific user, not in gr.State() that occurs during app init
+    assert db1 is not None and len(db1) == length_db1()
+    if db1[1] is None:
+        #  uuid in db is used as user ID
+        db1[1] = str(uuid.uuid4())
+
+
+def length_db1():
+    # For MyData:
+    # 0: db
+    # 1: userid and dbid
+    # 2: username
+
+    # For others:
+    # 0: db
+    # 1: dbid
+    # 2: None
+    return 3

src/enums.py

@@ -0,0 +1,225 @@
+from enum import Enum
+
+
+class PromptType(Enum):
+    custom = -1
+    plain = 0
+    instruct = 1
+    quality = 2
+    human_bot = 3
+    dai_faq = 4
+    summarize = 5
+    simple_instruct = 6
+    instruct_vicuna = 7
+    instruct_with_end = 8
+    human_bot_orig = 9
+    prompt_answer = 10
+    open_assistant = 11
+    wizard_lm = 12
+    wizard_mega = 13
+    instruct_vicuna2 = 14
+    instruct_vicuna3 = 15
+    wizard2 = 16
+    wizard3 = 17
+    instruct_simple = 18
+    wizard_vicuna = 19
+    openai = 20
+    openai_chat = 21
+    gptj = 22
+    prompt_answer_openllama = 23
+    vicuna11 = 24
+    mptinstruct = 25
+    mptchat = 26
+    falcon = 27
+    guanaco = 28
+    llama2 = 29
+    beluga = 30
+    wizard3nospace = 31
+    one_shot = 32
+    falcon_chat = 33
+
+
+class DocumentSubset(Enum):
+    Relevant = 0
+    RelSources = 1
+    TopKSources = 2
+
+
+non_query_commands = [
+    DocumentSubset.RelSources.name,
+    DocumentSubset.TopKSources.name
+]
+
+
+class DocumentChoice(Enum):
+    ALL = 'All'
+
+
+class LangChainMode(Enum):
+    """LangChain mode"""
+
+    DISABLED = "Disabled"
+    LLM = "LLM"
+    WIKI = "wiki"
+    WIKI_FULL = "wiki_full"
+    USER_DATA = "UserData"
+    MY_DATA = "MyData"
+    GITHUB_H2OGPT = "github h2oGPT"
+    H2O_DAI_DOCS = "DriverlessAI docs"
+
+
+class LangChainTypes(Enum):
+    SHARED = 'shared'
+    PERSONAL = 'personal'
+    EITHER = 'either'  # used when user did not pass which one, so need to try both
+
+
+# modes should not be removed from visible list or added by name
+langchain_modes_intrinsic = [LangChainMode.DISABLED.value,
+                             LangChainMode.LLM.value,
+                             LangChainMode.MY_DATA.value]
+
+langchain_modes_non_db = [LangChainMode.DISABLED.value,
+                          LangChainMode.LLM.value]
+
+
+class LangChainAction(Enum):
+    """LangChain action"""
+
+    QUERY = "Query"
+    # WIP:
+    # SUMMARIZE_MAP = "Summarize_map_reduce"
+    SUMMARIZE_MAP = "Summarize"
+    SUMMARIZE_ALL = "Summarize_all"
+    SUMMARIZE_REFINE = "Summarize_refine"
+
+
+class LangChainAgent(Enum):
+    """LangChain agents"""
+
+    SEARCH = "Search"
+    COLLECTION = "Collection"
+    PYTHON = "Python"
+    CSV = "CSV"
+    PANDAS = "Pandas"
+    JSON = 'JSON'
+
+
+no_server_str = no_lora_str = no_model_str = '[None/Remove]'
+
+# from site-packages/langchain/llms/openai.py
+# but needed since ChatOpenAI doesn't have this information
+model_token_mapping = {
+    "gpt-4": 8192,
+    "gpt-4-0314": 8192,
+    "gpt-4-32k": 32768,
+    "gpt-4-32k-0314": 32768,
+    "gpt-3.5-turbo": 4096,
+    "gpt-3.5-turbo-16k": 16 * 1024,
+    "gpt-3.5-turbo-0301": 4096,
+    "text-ada-001": 2049,
+    "ada": 2049,
+    "text-babbage-001": 2040,
+    "babbage": 2049,
+    "text-curie-001": 2049,
+    "curie": 2049,
+    "davinci": 2049,
+    "text-davinci-003": 4097,
+    "text-davinci-002": 4097,
+    "code-davinci-002": 8001,
+    "code-davinci-001": 8001,
+    "code-cushman-002": 2048,
+    "code-cushman-001": 2048,
+}
+
+font_size = 2
+head_acc = 40  # 40 for 6-way
+source_prefix = "Sources [Score | Link]:"
+source_postfix = "End Sources<p>"
+
+super_source_prefix = f"""<details><summary><font size="{font_size}">Sources</font></summary><font size="{font_size}"><font size="{font_size}">Sources [Score | Link]:"""
+super_source_postfix = f"""End Sources<p></font></font></details>"""
+
+
+def t5_type(model_name):
+    return 't5' == model_name.lower() or \
+        't5-' in model_name.lower() or \
+        'flan-' in model_name.lower() or \
+        'fastchat-t5' in model_name.lower()
+
+
+def get_langchain_prompts(pre_prompt_query, prompt_query, pre_prompt_summary, prompt_summary,
+                          model_name, inference_server, model_path_llama):
+    if model_name and ('falcon' in model_name or
+                       'Llama-2'.lower() in model_name.lower() or
+                       model_path_llama and 'llama-2' in model_path_llama.lower()) or \
+            model_name in [None, '']:
+        # use when no model, like no --base_model
+        pre_prompt_query1 = "Pay attention and remember the information below, which will help to answer the question or imperative after the context ends.\n"
+        prompt_query1 = "According to only the information in the document sources provided within the context above, "
+    elif inference_server and inference_server.startswith('openai'):
+        pre_prompt_query1 = "Pay attention and remember the information below, which will help to answer the question or imperative after the context ends.  If the answer cannot be primarily obtained from information within the context, then respond that the answer does not appear in the context of the documents.\n"
+        prompt_query1 = "According to (primarily) the information in the document sources provided within context above, "
+    else:
+        pre_prompt_query1 = ""
+        prompt_query1 = ""
+
+    pre_prompt_summary1 = """In order to write a concise single-paragraph or bulleted list summary, pay attention to the following text\n"""
+    prompt_summary1 = "Using only the information in the document sources above, write a condensed and concise summary of key results (preferably as bullet points):\n"
+
+    if pre_prompt_query is None:
+        pre_prompt_query = pre_prompt_query1
+    if prompt_query is None:
+        prompt_query = prompt_query1
+    if pre_prompt_summary is None:
+        pre_prompt_summary = pre_prompt_summary1
+    if prompt_summary is None:
+        prompt_summary = prompt_summary1
+
+    return pre_prompt_query, prompt_query, pre_prompt_summary, prompt_summary
+
+
+def gr_to_lg(image_loaders,
+             pdf_loaders,
+             url_loaders,
+             **kwargs,
+             ):
+    if image_loaders is None:
+        image_loaders = kwargs['image_loaders_options0']
+    if pdf_loaders is None:
+        pdf_loaders = kwargs['pdf_loaders_options0']
+    if url_loaders is None:
+        url_loaders = kwargs['url_loaders_options0']
+    # translate:
+    # 'auto' wouldn't be used here
+    ret = dict(
+        # urls
+        use_unstructured='Unstructured' in url_loaders,
+        use_playwright='PlayWright' in url_loaders,
+        use_selenium='Selenium' in url_loaders,
+
+        # pdfs
+        use_pymupdf='on' if 'PyMuPDF' in pdf_loaders else 'off',
+        use_unstructured_pdf='on' if 'Unstructured' in pdf_loaders else 'off',
+        use_pypdf='on' if 'PyPDF' in pdf_loaders else 'off',
+        enable_pdf_ocr='on' if 'OCR' in pdf_loaders else 'off',
+        enable_pdf_doctr='on' if 'DocTR' in pdf_loaders else 'off',
+        try_pdf_as_html='on' if 'TryHTML' in pdf_loaders else 'off',
+
+        # images
+        enable_ocr='OCR' in image_loaders,
+        enable_doctr='DocTR' in image_loaders,
+        enable_pix2struct='Pix2Struct' in image_loaders,
+        enable_captions='Caption' in image_loaders or 'CaptionBlip2' in image_loaders,
+    )
+    if 'CaptionBlip2' in image_loaders:
+        # just override, don't actually do both even if user chose both
+        captions_model = "Salesforce/blip2-flan-t5-xl"
+    else:
+        captions_model = kwargs['captions_model']
+    return ret, captions_model
+
+
+invalid_key_msg = 'Invalid Access Key, request access key from sales@h2o.ai or jon.mckinney@h2o.ai'
+
+docs_ordering_types = ['best_first', 'best_near_prompt', 'reverse_ucurve_sort']

src/evaluate_params.py

@@ -0,0 +1,71 @@
+input_args_list = ['model_state', 'my_db_state', 'selection_docs_state', 'requests_state']
+
+no_default_param_names = [
+    'instruction',
+    'iinput',
+    'context',
+    'instruction_nochat',
+    'iinput_nochat',
+]
+
+gen_hyper0 = ['num_beams',
+              'max_new_tokens',
+              'min_new_tokens',
+              'early_stopping',
+              'max_time',
+              'repetition_penalty',
+              'num_return_sequences',
+              'do_sample',
+              ]
+gen_hyper = ['temperature',
+             'top_p',
+             'top_k'] + gen_hyper0
+reader_names = ['image_loaders', 'pdf_loaders', 'url_loaders', 'jq_schema']
+
+eval_func_param_names = ['instruction',
+                         'iinput',
+                         'context',
+                         'stream_output',
+                         'prompt_type',
+                         'prompt_dict'] + \
+                        gen_hyper + \
+                        ['chat',
+                         'instruction_nochat',
+                         'iinput_nochat',
+                         'langchain_mode',
+                         'add_chat_history_to_context',
+                         'langchain_action',
+                         'langchain_agents',
+                         'top_k_docs',
+                         'chunk',
+                         'chunk_size',
+                         'document_subset',
+                         'document_choice',
+                         'pre_prompt_query',
+                         'prompt_query',
+                         'pre_prompt_summary',
+                         'prompt_summary',
+                         'system_prompt',
+                         ] + \
+                        reader_names + \
+                        ['visible_models',
+                         'h2ogpt_key',
+                         'add_search_to_context',
+                         'chat_conversation',
+                         'text_context_list',
+                         'docs_ordering_type',
+                         'min_max_new_tokens',
+                         ]
+
+# form evaluate defaults for submit_nochat_api
+eval_func_param_names_defaults = eval_func_param_names.copy()
+for k in no_default_param_names:
+    if k in eval_func_param_names_defaults:
+        eval_func_param_names_defaults.remove(k)
+
+eval_extra_columns = ['prompt', 'response', 'score']
+
+# override default_kwargs if user_kwargs None for args evaluate() uses that are not just in model_state
+# ensure prompt_type consistent with prep_bot(), so nochat API works same way
+# see how default_kwargs is set in gradio_runner.py
+key_overrides = ['prompt_type', 'prompt_dict']

src/gpt4all_llm.py

@@ -0,0 +1,403 @@
+import inspect
+import os
+from typing import Dict, Any, Optional, List, Iterator
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.schema.output import GenerationChunk
+from pydantic import root_validator
+from langchain.llms import gpt4all
+
+from utils import FakeTokenizer, get_ngpus_vis, url_alive, download_simple
+
+
+def get_model_tokenizer_gpt4all(base_model, n_jobs=None, max_seq_len=None, llamacpp_dict=None):
+    assert llamacpp_dict is not None
+    # defaults (some of these are generation parameters, so need to be passed in at generation time)
+    model_name = base_model.lower()
+    model = get_llm_gpt4all(model_name, model=None,
+                            # max_new_tokens=max_new_tokens,
+                            # temperature=temperature,
+                            # repetition_penalty=repetition_penalty,
+                            # top_k=top_k,
+                            # top_p=top_p,
+                            # callbacks=callbacks,
+                            n_jobs=n_jobs,
+                            # verbose=verbose,
+                            # streaming=stream_output,
+                            # prompter=prompter,
+                            # context=context,
+                            # iinput=iinput,
+                            inner_class=True,
+                            max_seq_len=max_seq_len,
+                            llamacpp_dict=llamacpp_dict,
+                            )
+    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(llamacpp_dict, default_kwargs, cls, exclude_list=[]):
+    # default from class
+    model_kwargs = {k: v.default for k, v in dict(inspect.signature(cls).parameters).items() if k not in exclude_list}
+    # from our defaults
+    model_kwargs.update(default_kwargs)
+    # from user defaults
+    model_kwargs.update(llamacpp_dict)
+    # 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}
+    # make int or float if can to satisfy types for class
+    for k, v in model_kwargs.items():
+        try:
+            if float(v) == int(v):
+                model_kwargs[k] = int(v)
+            else:
+                model_kwargs[k] = float(v)
+        except:
+            pass
+    return model_kwargs
+
+
+def get_gpt4all_default_kwargs(max_new_tokens=256,
+                               temperature=0.1,
+                               repetition_penalty=1.0,
+                               top_k=40,
+                               top_p=0.7,
+                               n_jobs=None,
+                               verbose=False,
+                               max_seq_len=None,
+                               ):
+    if n_jobs in [None, -1]:
+        n_jobs = int(os.getenv('OMP_NUM_THREADS', str(os.cpu_count()//2)))
+    n_jobs = max(1, min(20, n_jobs))  # hurts beyond some point
+    n_gpus = get_ngpus_vis()
+    default_kwargs = dict(context_erase=0.5,
+                          n_batch=1,
+                          max_tokens=max_seq_len - max_new_tokens,
+                          n_predict=max_new_tokens,
+                          repeat_last_n=64 if repetition_penalty != 1.0 else 0,
+                          repeat_penalty=repetition_penalty,
+                          temp=temperature,
+                          temperature=temperature,
+                          top_k=top_k,
+                          top_p=top_p,
+                          use_mlock=True,
+                          n_ctx=max_seq_len,
+                          n_threads=n_jobs,
+                          verbose=verbose)
+    if n_gpus != 0:
+        default_kwargs.update(dict(n_gpu_layers=100))
+    return default_kwargs
+
+
+def get_llm_gpt4all(model_name,
+                    model=None,
+                    max_new_tokens=256,
+                    temperature=0.1,
+                    repetition_penalty=1.0,
+                    top_k=40,
+                    top_p=0.7,
+                    streaming=False,
+                    callbacks=None,
+                    prompter=None,
+                    context='',
+                    iinput='',
+                    n_jobs=None,
+                    verbose=False,
+                    inner_class=False,
+                    max_seq_len=None,
+                    llamacpp_dict=None,
+                    ):
+    if not inner_class:
+        assert prompter is not None
+
+    default_kwargs = \
+        get_gpt4all_default_kwargs(max_new_tokens=max_new_tokens,
+                                   temperature=temperature,
+                                   repetition_penalty=repetition_penalty,
+                                   top_k=top_k,
+                                   top_p=top_p,
+                                   n_jobs=n_jobs,
+                                   verbose=verbose,
+                                   max_seq_len=max_seq_len,
+                                   )
+    if model_name == 'llama':
+        cls = H2OLlamaCpp
+        if model is None:
+            llamacpp_dict = llamacpp_dict.copy()
+            model_path = llamacpp_dict.pop('model_path_llama')
+            if os.path.isfile(os.path.basename(model_path)):
+                # e.g. if offline but previously downloaded
+                model_path = os.path.basename(model_path)
+            elif url_alive(model_path):
+                # online
+                ggml_path = os.getenv('GGML_PATH')
+                dest = os.path.join(ggml_path, os.path.basename(model_path)) if ggml_path else None
+                model_path = download_simple(model_path, dest=dest)
+        else:
+            model_path = model
+        model_kwargs = get_model_kwargs(llamacpp_dict, default_kwargs, cls, exclude_list=['lc_kwargs'])
+        model_kwargs.update(dict(model_path=model_path, callbacks=callbacks, streaming=streaming,
+                                 prompter=prompter, context=context, iinput=iinput))
+
+        # migration to  new langchain fix:
+        odd_keys = ['model_kwargs', 'grammar_path', 'grammar']
+        for key in odd_keys:
+            model_kwargs.pop(key, None)
+
+        llm = cls(**model_kwargs)
+        llm.client.verbose = verbose
+        inner_model = llm.client
+    elif model_name == 'gpt4all_llama':
+        cls = H2OGPT4All
+        if model is None:
+            llamacpp_dict = llamacpp_dict.copy()
+            model_path = llamacpp_dict.pop('model_name_gpt4all_llama')
+            if url_alive(model_path):
+                # online
+                ggml_path = os.getenv('GGML_PATH')
+                dest = os.path.join(ggml_path, os.path.basename(model_path)) if ggml_path else None
+                model_path = download_simple(model_path, dest=dest)
+        else:
+            model_path = model
+        model_kwargs = get_model_kwargs(llamacpp_dict, default_kwargs, cls, exclude_list=['lc_kwargs'])
+        model_kwargs.update(
+            dict(model=model_path, backend='llama', callbacks=callbacks, streaming=streaming,
+                 prompter=prompter, context=context, iinput=iinput))
+        llm = cls(**model_kwargs)
+        inner_model = llm.client
+    elif model_name == 'gptj':
+        cls = H2OGPT4All
+        if model is None:
+            llamacpp_dict = llamacpp_dict.copy()
+            model_path = llamacpp_dict.pop('model_name_gptj') if model is None else model
+            if url_alive(model_path):
+                ggml_path = os.getenv('GGML_PATH')
+                dest = os.path.join(ggml_path, os.path.basename(model_path)) if ggml_path else None
+                model_path = download_simple(model_path, dest=dest)
+        else:
+            model_path = model
+        model_kwargs = get_model_kwargs(llamacpp_dict, default_kwargs, cls, exclude_list=['lc_kwargs'])
+        model_kwargs.update(
+            dict(model=model_path, backend='gptj', callbacks=callbacks, streaming=streaming,
+                 prompter=prompter, context=context, iinput=iinput))
+        llm = cls(**model_kwargs)
+        inner_model = llm.client
+    else:
+        raise RuntimeError("No such model_name %s" % model_name)
+    if inner_class:
+        return inner_model
+    else:
+        return llm
+
+
+class H2OGPT4All(gpt4all.GPT4All):
+    model: Any
+    prompter: Any
+    context: Any = ''
+    iinput: Any = ''
+    """Path to the pre-trained GPT4All model file."""
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that the python package exists in the environment."""
+        try:
+            if isinstance(values["model"], str):
+                from gpt4all import GPT4All as GPT4AllModel
+
+                full_path = values["model"]
+                model_path, delimiter, model_name = full_path.rpartition("/")
+                model_path += delimiter
+
+                values["client"] = GPT4AllModel(
+                    model_name=model_name,
+                    model_path=model_path or None,
+                    model_type=values["backend"],
+                    allow_download=True,
+                )
+                if values["n_threads"] is not None:
+                    # set n_threads
+                    values["client"].model.set_thread_count(values["n_threads"])
+            else:
+                values["client"] = values["model"]
+                if values["n_threads"] is not None:
+                    # set n_threads
+                    values["client"].model.set_thread_count(values["n_threads"])
+            try:
+                values["backend"] = values["client"].model_type
+            except AttributeError:
+                # The below is for compatibility with GPT4All Python bindings <= 0.2.3.
+                values["backend"] = values["client"].model.model_type
+
+        except ImportError:
+            raise ValueError(
+                "Could not import gpt4all python package. "
+                "Please install it with `pip install gpt4all`."
+            )
+        return values
+
+    def _call(
+            self,
+            prompt: str,
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs,
+    ) -> str:
+        # Roughly 4 chars per token if natural language
+        n_ctx = 2048
+        prompt = prompt[-self.max_tokens * 4:]
+
+        # use instruct prompting
+        data_point = dict(context=self.context, instruction=prompt, input=self.iinput)
+        prompt = self.prompter.generate_prompt(data_point)
+
+        verbose = False
+        if verbose:
+            print("_call prompt: %s" % prompt, flush=True)
+        # FIXME: GPT4ALl doesn't support yield during generate, so cannot support streaming except via itself to stdout
+        return super()._call(prompt, stop=stop, run_manager=run_manager)
+
+    # FIXME:  Unsure what uses
+    #def get_token_ids(self, text: str) -> List[int]:
+    #    return self.client.tokenize(b" " + text.encode("utf-8"))
+
+
+from langchain.llms import LlamaCpp
+
+
+class H2OLlamaCpp(LlamaCpp):
+    model_path: Any
+    prompter: Any
+    context: Any
+    iinput: Any
+    """Path to the pre-trained GPT4All model file."""
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that llama-cpp-python library is installed."""
+        if isinstance(values["model_path"], str):
+            model_path = values["model_path"]
+            model_param_names = [
+                "lora_path",
+                "lora_base",
+                "n_ctx",
+                "n_parts",
+                "seed",
+                "f16_kv",
+                "logits_all",
+                "vocab_only",
+                "use_mlock",
+                "n_threads",
+                "n_batch",
+                "use_mmap",
+                "last_n_tokens_size",
+            ]
+            model_params = {k: values[k] for k in model_param_names}
+            # For backwards compatibility, only include if non-null.
+            if values["n_gpu_layers"] is not None:
+                model_params["n_gpu_layers"] = values["n_gpu_layers"]
+
+            try:
+                try:
+                    from llama_cpp import Llama
+                except ImportError:
+                    from llama_cpp_cuda import Llama
+
+                values["client"] = Llama(model_path, **model_params)
+            except ImportError:
+                raise ModuleNotFoundError(
+                    "Could not import llama-cpp-python library. "
+                    "Please install the llama-cpp-python library to "
+                    "use this embedding model: pip install llama-cpp-python"
+                )
+            except Exception as e:
+                raise ValueError(
+                    f"Could not load Llama model from path: {model_path}. "
+                    f"Received error {e}"
+                )
+        else:
+            values["client"] = values["model_path"]
+        return values
+
+    def _call(
+            self,
+            prompt: str,
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs,
+    ) -> str:
+        verbose = False
+        # tokenize twice, just to count tokens, since llama cpp python wrapper has no way to truncate
+        # still have to avoid crazy sizes, else hit llama_tokenize: too many tokens -- might still hit, not fatal
+        prompt = prompt[-self.n_ctx * 4:]
+        prompt_tokens = self.client.tokenize(b" " + prompt.encode("utf-8"))
+        num_prompt_tokens = len(prompt_tokens)
+        if num_prompt_tokens > self.n_ctx:
+            # conservative by using int()
+            chars_per_token = int(len(prompt) / num_prompt_tokens)
+            prompt = prompt[-self.n_ctx * chars_per_token:]
+            if verbose:
+                print("reducing tokens, assuming average of %s chars/token: %s" % chars_per_token, flush=True)
+                prompt_tokens2 = self.client.tokenize(b" " + prompt.encode("utf-8"))
+                num_prompt_tokens2 = len(prompt_tokens2)
+                print("reduced tokens from %d -> %d" % (num_prompt_tokens, num_prompt_tokens2), flush=True)
+
+        # use instruct prompting
+        data_point = dict(context=self.context, instruction=prompt, input=self.iinput)
+        prompt = self.prompter.generate_prompt(data_point)
+
+        if verbose:
+            print("_call prompt: %s" % prompt, flush=True)
+
+        if self.streaming:
+            # parent handler of streamer expects to see prompt first else output="" and lose if prompt=None in prompter
+            text = ""
+            for token in self.stream(input=prompt, stop=stop):
+                # for token in self.stream(input=prompt, stop=stop, run_manager=run_manager):
+                text_chunk = token  # ["choices"][0]["text"]
+                # self.stream already calls text_callback
+                # if text_callback:
+                #    text_callback(text_chunk)
+                text += text_chunk
+            # parent handler of streamer expects to see prompt first else output="" and lose if prompt=None in prompter
+            return text[len(prompt):]
+        else:
+            params = self._get_parameters(stop)
+            params = {**params, **kwargs}
+            result = self.client(prompt=prompt, **params)
+            return result["choices"][0]["text"]
+
+    def _stream(
+            self,
+            prompt: str,
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs: Any,
+        ) -> Iterator[GenerationChunk]:
+        # parent handler of streamer expects to see prompt first else output="" and lose if prompt=None in prompter
+        logprobs = 0
+        chunk = GenerationChunk(
+            text=prompt,
+            generation_info={"logprobs": logprobs},
+        )
+        yield chunk
+        if run_manager:
+            run_manager.on_llm_new_token(
+                token=chunk.text, verbose=self.verbose, log_probs=logprobs
+            )
+        # actual new tokens
+        for chunk in super()._stream(prompt, stop=stop, run_manager=run_manager, **kwargs):
+            yield chunk
+
+    def get_token_ids(self, text: str) -> List[int]:
+        return self.client.tokenize(b" " + text.encode("utf-8"))

src/gradio_themes.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+from typing import Iterable
+
+from gradio.themes.soft import Soft
+from gradio.themes import Color, Size
+from gradio.themes.utils import colors, sizes, fonts
+
+h2o_yellow = Color(
+    name="yellow",
+    c50="#fffef2",
+    c100="#fff9e6",
+    c200="#ffecb3",
+    c300="#ffe28c",
+    c400="#ffd659",
+    c500="#fec925",
+    c600="#e6ac00",
+    c700="#bf8f00",
+    c800="#a67c00",
+    c900="#664d00",
+    c950="#403000",
+)
+h2o_gray = Color(
+    name="gray",
+    c50="#f8f8f8",
+    c100="#e5e5e5",
+    c200="#cccccc",
+    c300="#b2b2b2",
+    c400="#999999",
+    c500="#7f7f7f",
+    c600="#666666",
+    c700="#4c4c4c",
+    c800="#333333",
+    c900="#191919",
+    c950="#0d0d0d",
+)
+
+text_xsm = Size(
+    name="text_xsm",
+    xxs="4px",
+    xs="5px",
+    sm="6px",
+    md="7px",
+    lg="8px",
+    xl="10px",
+    xxl="12px",
+)
+
+spacing_xsm = Size(
+    name="spacing_xsm",
+    xxs="1px",
+    xs="1px",
+    sm="1px",
+    md="2px",
+    lg="3px",
+    xl="5px",
+    xxl="7px",
+)
+
+radius_xsm = Size(
+    name="radius_xsm",
+    xxs="1px",
+    xs="1px",
+    sm="1px",
+    md="2px",
+    lg="3px",
+    xl="5px",
+    xxl="7px",
+)
+
+
+class H2oTheme(Soft):
+    def __init__(
+            self,
+            *,
+            primary_hue: colors.Color | str = h2o_yellow,
+            secondary_hue: colors.Color | str = h2o_yellow,
+            neutral_hue: colors.Color | str = h2o_gray,
+            spacing_size: sizes.Size | str = sizes.spacing_md,
+            radius_size: sizes.Size | str = sizes.radius_md,
+            text_size: sizes.Size | str = sizes.text_lg,
+            font: fonts.Font
+                  | str
+                  | Iterable[fonts.Font | str] = (
+                    fonts.GoogleFont("Montserrat"),
+                    "ui-sans-serif",
+                    "system-ui",
+                    "sans-serif",
+            ),
+            font_mono: fonts.Font
+                       | str
+                       | Iterable[fonts.Font | str] = (
+                    fonts.GoogleFont("IBM Plex Mono"),
+                    "ui-monospace",
+                    "Consolas",
+                    "monospace",
+            ),
+    ):
+        super().__init__(
+            primary_hue=primary_hue,
+            secondary_hue=secondary_hue,
+            neutral_hue=neutral_hue,
+            spacing_size=spacing_size,
+            radius_size=radius_size,
+            text_size=text_size,
+            font=font,
+            font_mono=font_mono,
+        )
+        super().set(
+            background_fill_primary_dark="*block_background_fill",
+            block_background_fill_dark="*neutral_950",
+            block_border_width='1px',
+            block_border_width_dark='1px',
+            block_label_background_fill="*primary_300",
+            block_label_background_fill_dark="*primary_600",
+            block_label_text_color="*neutral_950",
+            block_label_text_color_dark="*neutral_950",
+            block_radius="0 0 8px 8px",
+            block_title_text_color="*neutral_950",
+            block_title_text_color_dark="*neutral_950",
+            body_background_fill="*neutral_50",
+            body_background_fill_dark="*neutral_900",
+            border_color_primary="*neutral_100",
+            border_color_primary_dark="*neutral_700",
+            button_border_width="1px",
+            button_border_width_dark="1px",
+            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",
+            button_secondary_background_fill_hover_dark="*primary_700",
+            button_secondary_border_color="*primary_500",
+            button_secondary_border_color_dark="*primary_500",
+            button_secondary_border_color_hover_dark="*primary_700",
+            checkbox_label_text_color_selected_dark='#000000',
+            # checkbox_label_text_size="*text_xs",  # too small for iPhone etc. but good if full large screen zoomed to fit
+            checkbox_label_text_size="*text_sm",
+            # radio_circle="""url("data:image/svg+xml,%3csvg viewBox='0 0 32 32' fill='white' xmlns='http://www.w3.org/2000/svg'%3e%3ccircle cx='32' cy='32' r='1'/%3e%3c/svg%3e")""",
+            # checkbox_border_width=1,
+            # heckbox_border_width_dark=1,
+            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",
+        )
+
+
+class SoftTheme(Soft):
+    def __init__(
+            self,
+            *,
+            primary_hue: colors.Color | str = colors.indigo,
+            secondary_hue: colors.Color | str = colors.indigo,
+            neutral_hue: colors.Color | str = colors.gray,
+            spacing_size: sizes.Size | str = sizes.spacing_md,
+            radius_size: sizes.Size | str = sizes.radius_md,
+            text_size: sizes.Size | str = sizes.text_md,
+            font: fonts.Font
+                  | str
+                  | Iterable[fonts.Font | str] = (
+                    fonts.GoogleFont("Montserrat"),
+                    "ui-sans-serif",
+                    "system-ui",
+                    "sans-serif",
+            ),
+            font_mono: fonts.Font
+                       | str
+                       | Iterable[fonts.Font | str] = (
+                    fonts.GoogleFont("IBM Plex Mono"),
+                    "ui-monospace",
+                    "Consolas",
+                    "monospace",
+            ),
+    ):
+        super().__init__(
+            primary_hue=primary_hue,
+            secondary_hue=secondary_hue,
+            neutral_hue=neutral_hue,
+            spacing_size=spacing_size,
+            radius_size=radius_size,
+            text_size=text_size,
+            font=font,
+            font_mono=font_mono,
+        )
+        super().set(
+            checkbox_label_text_size="*text_sm",
+        )
+
+
+h2o_logo = '<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" width="100%" height="100%"' \
+           ' viewBox="0 0 600.28 600.28"><defs><style>.cls-1{fill:#fec925;}.cls-2{fill:#161616;}.cls-3{fill:' \
+           '#54585a;}</style></defs><g id="Fill-1"><rect class="cls-1" width="600.28" height="600.28" ' \
+           'rx="23.24"/></g><path class="cls-2" d="M174.33,246.06v92.78H152.86v-38H110.71v38H89.24V246.06h21.' \
+           '47v36.58h42.15V246.06Z"/><path class="cls-2" d="M259.81,321.34v17.5H189.7V324.92l35.78-33.8c8.22-7.' \
+           '82,9.68-12.59,9.68-17.09,0-7.29-5-11.53-14.85-11.53-7.95,0-14.71,3-19.21,9.27L185.46,261.7c7.15-10' \
+           '.47,20.14-17.23,36.84-17.23,20.68,0,34.46,10.6,34.46,27.44,0,9-2.52,17.22-15.51,29.29l-21.33,20.14Z"' \
+           '/><path class="cls-2" d="M268.69,292.45c0-27.57,21.47-48,50.76-48s50.76,20.28,50.76,48-21.6,48-50.' \
+           '76,48S268.69,320,268.69,292.45Zm79.78,0c0-17.63-12.46-29.69-29-29.69s-29,12.06-29,29.69,12.46,29.69' \
+           ',29,29.69S348.47,310.08,348.47,292.45Z"/><path class="cls-3" d="M377.23,326.91c0-7.69,5.7-12.73,12.' \
+           '85-12.73s12.86,5,12.86,12.73a12.86,12.86,0,1,1-25.71,0Z"/><path class="cls-3" d="M481.4,298.15v40.' \
+           '69H462.05V330c-3.84,6.49-11.27,9.94-21.74,9.94-16.7,0-26.64-9.28-26.64-21.61,0-12.59,8.88-21.34,30.' \
+           '62-21.34h16.43c0-8.87-5.3-14-16.43-14-7.55,0-15.37,2.51-20.54,6.62l-7.43-14.44c7.82-5.57,19.35-8.' \
+           '62,30.75-8.62C468.81,266.47,481.4,276.54,481.4,298.15Zm-20.68,18.16V309H446.54c-9.67,0-12.72,3.57-' \
+           '12.72,8.35,0,5.16,4.37,8.61,11.66,8.61C452.37,326,458.34,322.8,460.72,316.31Z"/><path class="cls-3"' \
+           ' d="M497.56,246.06c0-6.49,5.17-11.53,12.86-11.53s12.86,4.77,12.86,11.13c0,6.89-5.17,11.93-12.86,' \
+           '11.93S497.56,252.55,497.56,246.06Zm2.52,21.47h20.68v71.31H500.08Z"/></svg>'
+
+
+def get_h2o_title(title, description):
+    # NOTE: Check full width desktop, smallest width browser desktop, iPhone browsers to ensure no overlap etc.
+    return f"""<div style="float:left; justify-content:left; height: 80px; width: 195px; margin-top:0px">
+                    {description}
+                </div>
+                <div style="display:flex; justify-content:center; margin-bottom:30px; margin-right:330px;">
+                    <div style="height: 60px; width: 60px; margin-right:20px;">{h2o_logo}</div>
+                    <h1 style="line-height:60px">{title}</h1>
+                </div>
+                <div style="float:right; height: 80px; width: 80px; margin-top:-100px">
+                    <img src="https://raw.githubusercontent.com/h2oai/h2ogpt/main/docs/h2o-qr.png">
+                </div>
+                """
+
+
+def get_simple_title(title, description):
+    return f"""{description}<h1 align="center"> {title}</h1>"""
+
+
+def get_dark_js() -> str:
+    return """
+        if (document.querySelectorAll('.dark').length) {
+            document.querySelectorAll('.dark').forEach(el => el.classList.remove('dark'));
+        } else {
+            document.querySelector('body').classList.add('dark');
+        }
+    """
+
+
+def get_heap_js(heapAppId: str) -> str:
+    return (
+        """globalThis.window.heap=window.heap||[],heap.load=function(e,t){window.heap.appid=e,window.heap.config=t=t||{};var r=document.createElement("script");r.type="text/javascript",r.async=!0,r.src="https://cdn.heapanalytics.com/js/heap-"+e+".js";var a=document.getElementsByTagName("script")[0];a.parentNode.insertBefore(r,a);for(var n=function(e){return function(){heap.push([e].concat(Array.prototype.slice.call(arguments,0)))}},p=["addEventProperties","addUserProperties","clearEventProperties","identify","resetIdentity","removeEventProperty","setEventProperties","track","unsetEventProperty"],o=0;o<p.length;o++)heap[p[o]]=n(p[o])};"""
+        f"""heap.load("{heapAppId}");""")
+
+
+def wrap_js_to_lambda(num_params: int, *args: str) -> str:
+    """
+    Generates a JS code representing JS lambda that wraps all given '*args' code strings.
+    The lambda function has number of parameters based on 'num_params' and returns them
+    without modification in an array. Lambda with zero parameters returns an empty array.
+    """
+    params = ", ".join([f"p{i}" for i in range(num_params)])
+    newline = "\n"
+    return f"""
+        ({params}) => {{
+            {newline.join([a for a in args if a is not None])}
+            return [{params}];
+        }}
+    """

src/gradio_utils/css.py

@@ -0,0 +1,148 @@
+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 """
+    #col_container {margin-left: auto; margin-right: auto; text-align: left;}
+
+    @import url('https://fonts.googleapis.com/css2?family=Source+Sans+Pro:wght@400;600&display=swap');
+    
+    body.dark{#warning {background-color: #555555};}
+    
+    #sidebar {
+        order: 1;
+        
+        @media (max-width: 463px) {
+          order: 2;
+        }
+    }
+    
+    #col-tabs {
+        order: 2;
+        
+        @media (max-width: 463px) {
+          order: 1;
+        }
+    }
+    
+    #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;
+    }
+    .chatsmall chatbot {font-size: 10px !important}
+    
+    .gradio-container {
+        max-width: none !important;
+    }
+    
+    div.message {
+        padding: var(--text-lg) !important;
+    }
+    
+    div.message.user > div.icon-button {
+        top: unset;
+        bottom: 0;
+    }
+    
+    div.message.bot > div.icon-button {
+        top: unset;
+        bottom: 0;
+    }
+    
+    #prompt-form-row {
+        position: relative;
+    }
+    
+    #attach-button {
+        position: absolute;
+        top: 45px;
+        right: 20px;
+        
+        display: flex;
+        justify-content: center;
+        border: 1px solid var(--primary-500) !important;
+        
+        @media (max-width: 463px) {
+          width: 56px;
+        }
+    }
+    
+    #attach-button > img {
+        margin-right: 0;
+    }
+    
+    #prompt-form > label > textarea {
+        padding-right: 104px;
+        
+        @media (max-width: 463px) {
+          min-height: 94px;
+          padding-right: 70px;
+        }
+    }
+
+    #visible-models > label > div.wrap > div.wrap-inner > div.secondary-wrap > div.remove-all {
+        display: none !important;
+    }
+    
+    #visible-models > label > div.wrap > div.wrap-inner > div.token {
+        display: none !important;
+    }
+    
+    #visible-models > label > div.wrap > div.wrap-inner > div.secondary-wrap::before {
+        content: "Select";
+        padding: 0 4px;
+        margin-right: 2px;
+    }
+
+    #langchain_agents > label > div.wrap > div.wrap-inner > div.secondary-wrap > div.remove-all {
+        display: none !important;
+    }
+
+    #langchain_agents > label > div.wrap > div.wrap-inner > div.token {
+        display: none !important;
+    }
+
+    #langchain_agents > label > div.wrap > div.wrap-inner > div.secondary-wrap::before {
+        content: "Select";
+        padding: 0 4px;
+        margin-right: 2px;
+    }
+    """

src/gradio_utils/grclient.py

@@ -0,0 +1,82 @@
+import traceback
+from typing import Callable
+import os
+
+from gradio_client.client import Job
+
+os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
+
+from gradio_client import Client
+
+
+class GradioClient(Client):
+    """
+    Parent class of gradio client
+    To handle automatically refreshing client if detect gradio server changed
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+        super().__init__(*args, **kwargs)
+        self.server_hash = self.get_server_hash()
+
+    def get_server_hash(self):
+        """
+        Get server hash using super without any refresh action triggered
+        Returns: git hash of gradio server
+        """
+        return super().submit(api_name='/system_hash').result()
+
+    def refresh_client_if_should(self):
+        # get current hash in order to update api_name -> fn_index map in case gradio server changed
+        # FIXME: Could add cli api as hash
+        server_hash = self.get_server_hash()
+        if self.server_hash != server_hash:
+            self.refresh_client()
+            self.server_hash = server_hash
+        else:
+            self.reset_session()
+
+    def refresh_client(self):
+        """
+        Ensure every client call is independent
+        Also ensure map between api_name and fn_index is updated in case server changed (e.g. restarted with new code)
+        Returns:
+        """
+        # need session hash to be new every time, to avoid "generator already executing"
+        self.reset_session()
+
+        client = Client(*self.args, **self.kwargs)
+        for k, v in client.__dict__.items():
+            setattr(self, k, v)
+
+    def submit(
+        self,
+        *args,
+        api_name: str | None = None,
+        fn_index: int | None = None,
+        result_callbacks: Callable | list[Callable] | None = None,
+    ) -> Job:
+        # Note predict calls submit
+        try:
+            self.refresh_client_if_should()
+            job = super().submit(*args, api_name=api_name, fn_index=fn_index)
+        except Exception as e:
+            print("Hit e=%s" % str(e), flush=True)
+            # force reconfig in case only that
+            self.refresh_client()
+            job = super().submit(*args, api_name=api_name, fn_index=fn_index)
+
+        # see if immediately failed
+        e = job.future._exception
+        if e is not None:
+            print("GR job failed: %s %s" % (str(e), ''.join(traceback.format_tb(e.__traceback__))), flush=True)
+            # force reconfig in case only that
+            self.refresh_client()
+            job = super().submit(*args, api_name=api_name, fn_index=fn_index)
+            e2 = job.future._exception
+            if e2 is not None:
+                print("GR job failed again: %s\n%s" % (str(e2), ''.join(traceback.format_tb(e2.__traceback__))), flush=True)
+
+        return job

src/gradio_utils/prompt_form.py

@@ -0,0 +1,108 @@
+import os
+import math
+
+import gradio as gr
+
+
+def make_chatbots(output_label0, output_label0_model2, **kwargs):
+    visible_models = kwargs['visible_models']
+    all_models = kwargs['all_models']
+
+    text_outputs = []
+    chat_kwargs = []
+    for model_state_locki, model_state_lock in enumerate(kwargs['model_states']):
+        if os.environ.get('DEBUG_MODEL_LOCK'):
+            model_name = model_state_lock["base_model"] + " : " + model_state_lock["inference_server"]
+        else:
+            model_name = model_state_lock["base_model"]
+        output_label = f'h2oGPT [{model_name}]'
+        min_width = 250 if kwargs['gradio_size'] in ['small', 'large', 'medium'] else 160
+        chat_kwargs.append(dict(label=output_label, elem_classes='chatsmall',
+                                height=kwargs['height'] or 400, min_width=min_width,
+                                show_copy_button=kwargs['show_copy_button'],
+                                visible=kwargs['model_lock'] and (visible_models is None or
+                                                                  model_state_locki in visible_models or
+                                                                  all_models[model_state_locki] in visible_models
+                                                                  )))
+
+    # base view on initial visible choice
+    if visible_models:
+        len_visible = len(visible_models)
+    else:
+        len_visible = len(kwargs['model_states'])
+    if kwargs['model_lock_columns'] == -1:
+        kwargs['model_lock_columns'] = len_visible
+    if kwargs['model_lock_columns'] is None:
+        kwargs['model_lock_columns'] = 3
+
+    ncols = kwargs['model_lock_columns']
+    if kwargs['model_states'] == 0:
+        nrows = 0
+    else:
+        nrows = math.ceil(len_visible / kwargs['model_lock_columns'])
+
+    if kwargs['model_lock_columns'] == 0:
+        # not using model_lock
+        pass
+    elif nrows <= 1:
+        with gr.Row():
+            for chat_kwargs1, model_state_lock in zip(chat_kwargs, kwargs['model_states']):
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+    elif nrows == kwargs['model_states']:
+        with gr.Row():
+            for chat_kwargs1, model_state_lock in zip(chat_kwargs, kwargs['model_states']):
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+    elif nrows == 2:
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii >= len_visible / 2:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < len_visible / 2:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+    elif nrows == 3:
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii >= 1 * len_visible / 3:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 1 * len_visible / 3 or mii >= 2 * len_visible / 3:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 2 * len_visible / 3:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+    elif nrows >= 4:
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii >= 1 * len_visible / 4:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 1 * len_visible / 4 or mii >= 2 * len_visible / 4:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 2 * len_visible / 4 or mii >= 3 * len_visible / 4:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+        with gr.Row():
+            for mii, (chat_kwargs1, model_state_lock) in enumerate(zip(chat_kwargs, kwargs['model_states'])):
+                if mii < 3 * len_visible / 4:
+                    continue
+                text_outputs.append(gr.Chatbot(**chat_kwargs1))
+
+    with gr.Row():
+        text_output = gr.Chatbot(label=output_label0, visible=not kwargs['model_lock'], height=kwargs['height'] or 400)
+        text_output2 = gr.Chatbot(label=output_label0_model2,
+                                  visible=False and not kwargs['model_lock'], height=kwargs['height'] or 400)
+    return text_output, text_output2, text_outputs