Update with h2oGPT hash 236c95819e80ab122193bfb843b55618ae285c39

client_test.py

@@ -97,6 +97,8 @@ def get_args(prompt, prompt_type, chat=False, stream_output=False,
                          chunk_size=512,
                          document_choice=[DocumentChoices.All_Relevant.name],
                          )
+    from generate import eval_func_param_names
+    assert len(set(eval_func_param_names).difference(set(list(kwargs.keys())))) == 0
     if chat:
         # add chatbot output on end.  Assumes serialize=False
         kwargs.update(dict(chatbot=[]))
@@ -105,8 +107,8 @@ def get_args(prompt, prompt_type, chat=False, stream_output=False,
 
 
 @pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_basic():
-    return run_client_nochat(prompt='Who are you?', prompt_type='human_bot', max_new_tokens=50)
+def test_client_basic(prompt_type='human_bot'):
+    return run_client_nochat(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
 
 
 def run_client_nochat(prompt, prompt_type, max_new_tokens):
@@ -122,12 +124,12 @@ def run_client_nochat(prompt, prompt_type, max_new_tokens):
     res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
                     response=md_to_text(res))
     print(res_dict)
-    return res_dict
+    return res_dict, client
 
 
 @pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_basic_api():
-    return run_client_nochat_api(prompt='Who are you?', prompt_type='human_bot', max_new_tokens=50)
+def test_client_basic_api(prompt_type='human_bot'):
+    return run_client_nochat_api(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
 
 
 def run_client_nochat_api(prompt, prompt_type, max_new_tokens):
@@ -144,12 +146,12 @@ def run_client_nochat_api(prompt, prompt_type, max_new_tokens):
                     response=md_to_text(ast.literal_eval(res)['response']),
                     sources=ast.literal_eval(res)['sources'])
     print(res_dict)
-    return 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():
-    return run_client_nochat_api_lean(prompt='Who are you?', prompt_type='human_bot', max_new_tokens=50)
+def test_client_basic_api_lean(prompt_type='human_bot'):
+    return run_client_nochat_api_lean(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
 
 
 def run_client_nochat_api_lean(prompt, prompt_type, max_new_tokens):
@@ -166,21 +168,21 @@ def run_client_nochat_api_lean(prompt, prompt_type, max_new_tokens):
                     response=md_to_text(ast.literal_eval(res)['response']),
                     sources=ast.literal_eval(res)['sources'])
     print(res_dict)
-    return 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():
-    return run_client_nochat_api_lean_morestuff(prompt='Who are you?', prompt_type='human_bot', max_new_tokens=50)
+def test_client_basic_api_lean_morestuff(prompt_type='human_bot'):
+    return run_client_nochat_api_lean_morestuff(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
 
 
-def run_client_nochat_api_lean_morestuff(prompt, prompt_type, max_new_tokens):
+def run_client_nochat_api_lean_morestuff(prompt, prompt_type='human_bot', max_new_tokens=512):
     kwargs = dict(
         instruction='',
         iinput='',
         context='',
         stream_output=False,
-        prompt_type='human_bot',
+        prompt_type=prompt_type,
         temperature=0.1,
         top_p=0.75,
         top_k=40,
@@ -211,12 +213,19 @@ def run_client_nochat_api_lean_morestuff(prompt, prompt_type, max_new_tokens):
                     response=md_to_text(ast.literal_eval(res)['response']),
                     sources=ast.literal_eval(res)['sources'])
     print(res_dict)
-    return res_dict
+    return res_dict, client
 
 
 @pytest.mark.skip(reason="For manual use against some server, no server launched")
-def test_client_chat():
-    return run_client_chat(prompt='Who are you?', prompt_type='human_bot', stream_output=False, max_new_tokens=50,
+def test_client_chat(prompt_type='human_bot'):
+    return run_client_chat(prompt='Who are you?', prompt_type=prompt_type, stream_output=False, max_new_tokens=50,
+                           langchain_mode='Disabled')
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_chat_stream(prompt_type='human_bot'):
+    return run_client_chat(prompt="Tell a very long kid's story about birds.", prompt_type=prompt_type,
+                           stream_output=True, max_new_tokens=512,
                            langchain_mode='Disabled')
 
 
@@ -229,6 +238,7 @@ def run_client_chat(prompt, prompt_type, stream_output, max_new_tokens, langchai
 
 
 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]]
 
@@ -262,6 +272,46 @@ def run_client(client, prompt, args, kwargs, do_md_to_text=True, verbose=False):
         return res_dict, client
 
 
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_nochat_stream(prompt_type='human_bot'):
+    return run_client_nochat_gen(prompt="Tell a very long kid's story about birds.", prompt_type=prompt_type,
+                                 stream_output=True, max_new_tokens=512,
+                                 langchain_mode='Disabled')
+
+
+def run_client_nochat_gen(prompt, prompt_type, stream_output, max_new_tokens, langchain_mode):
+    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)
+    return run_client_gen(client, prompt, args, kwargs)
+
+
+def run_client_gen(client, prompt, args, kwargs, do_md_to_text=True, verbose=False):
+    res_dict = kwargs
+    res_dict['prompt'] = prompt
+    if not kwargs['stream_output']:
+        res = client.predict(str(dict(kwargs)), api_name='/submit_nochat_api')
+        res_dict['response'] = res[0]
+        print(md_to_text(res_dict['response'], do_md_to_text=do_md_to_text))
+        return res_dict, client
+    else:
+        job = client.submit(str(dict(kwargs)), api_name='/submit_nochat_api')
+        while not job.done():
+            outputs_list = job.communicator.job.outputs
+            if outputs_list:
+                res = job.communicator.job.outputs[-1]
+                res_dict = ast.literal_eval(res)
+                print('Stream: %s' % res_dict['response'])
+            time.sleep(0.1)
+        res_list = job.outputs()
+        assert len(res_list) > 0, "No response, check server"
+        res = res_list[-1]
+        res_dict = ast.literal_eval(res)
+        print('Final: %s' % res_dict['response'])
+        return res_dict, client
+
+
 def md_to_text(md, do_md_to_text=True):
     if not do_md_to_text:
         return md
@@ -271,8 +321,16 @@ def md_to_text(md, do_md_to_text=True):
     return soup.get_text()
 
 
+def run_client_many(prompt_type='human_bot'):
+    ret1, _ = test_client_chat(prompt_type=prompt_type)
+    ret2, _ = test_client_chat_stream(prompt_type=prompt_type)
+    ret3, _ = test_client_nochat_stream(prompt_type=prompt_type)
+    ret4, _ = test_client_basic(prompt_type=prompt_type)
+    ret5, _ = test_client_basic_api(prompt_type=prompt_type)
+    ret6, _ = test_client_basic_api_lean(prompt_type=prompt_type)
+    ret7, _ = test_client_basic_api_lean_morestuff(prompt_type=prompt_type)
+    return ret1, ret2, ret3, ret4, ret5, ret6, ret7
+
+
 if __name__ == '__main__':
-    test_client_basic()
-    test_client_basic_api()
-    test_client_basic_api_lean()
-    test_client_basic_api_lean_morestuff()
+    run_client_many()

create_data.py

@@ -567,7 +567,7 @@ def test_show_prompts():
     from prompter import generate_prompt
     for data_points in file_points:
         for data_point in data_points:
-            print(generate_prompt(data_point, 'plain', '', False, False)[0])
+            print(generate_prompt(data_point, 'plain', '', False, False, False)[0])
 
 
 def test_get_open_datasets():
@@ -1571,7 +1571,7 @@ def test_check_stats_data():
 
     llama_type = False
     tokenizer_base_model = base_model = 'h2oai/h2ogpt-oasst1-512-20b'
-    model_loader, tokenizer_loader = get_loaders(llama_type=llama_type, model_name=base_model, reward_type=False)
+    model_loader, tokenizer_loader = get_loaders(model_name=base_model, reward_type=False, llama_type=llama_type)
     local_files_only = False
     resume_download = True
     use_auth_token = False

enums.py

@@ -22,6 +22,12 @@ class PromptType(Enum):
     wizard2 = 16
     wizard3 = 17
     instruct_simple = 18
+    wizard_vicuna = 19
+    openai = 20
+    openai_chat = 21
+    gptj = 22
+    prompt_answer_openllama = 23
+    vicuna11 = 24
 
 
 class DocumentChoices(Enum):
@@ -44,3 +50,35 @@ class LangChainMode(Enum):
     MY_DATA = "MyData"
     GITHUB_H2OGPT = "github h2oGPT"
     H2O_DAI_DOCS = "DriverlessAI docs"
+
+
+no_server_str = no_lora_str = no_model_str = '[None/Remove]'
+
+
+# from site-packages/langchain/llms/openai.py, but needed since ChatOpenAI doesn't have this information
+model_token_mapping = {
+    "gpt-4": 8192,
+    "gpt-4-0314": 8192,
+    "gpt-4-32k": 32768,
+    "gpt-4-32k-0314": 32768,
+    "gpt-3.5-turbo": 4096,
+    "gpt-3.5-turbo-16k": 16*1024,
+    "gpt-3.5-turbo-0301": 4096,
+    "text-ada-001": 2049,
+    "ada": 2049,
+    "text-babbage-001": 2040,
+    "babbage": 2049,
+    "text-curie-001": 2049,
+    "curie": 2049,
+    "davinci": 2049,
+    "text-davinci-003": 4097,
+    "text-davinci-002": 4097,
+    "code-davinci-002": 8001,
+    "code-davinci-001": 8001,
+    "code-cushman-002": 2048,
+    "code-cushman-001": 2048,
+}
+
+
+source_prefix = "Sources [Score | Link]:"
+source_postfix = "End Sources<p>"

finetune.py

@@ -5,6 +5,9 @@ from typing import List, Union
 import fire
 import numpy as np
 
+if os.path.dirname(os.path.abspath(__file__)) not in sys.path:
+    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
 from loaders import get_loaders, get_tokenizer
 from prompter import generate_prompt, prompt_types, PromptType
 from utils import get_githash, copy_code
@@ -182,7 +185,7 @@ def train(
             log("num_gpus: %d" % gpus)
             log("max mem: %s" % max_memory)
 
-    model_loader, tokenizer_loader = get_loaders(llama_type=llama_type, model_name=base_model, reward_type=False)
+    model_loader, tokenizer_loader = get_loaders(model_name=base_model, reward_type=False, llama_type=llama_type)
 
     model = model_loader.from_pretrained(
         base_model,
@@ -556,13 +559,6 @@ def train(
     )
     model.config.use_cache = False
 
-    old_state_dict = model.state_dict
-    from peft import get_peft_model_state_dict
-
-    model.state_dict = (
-        lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())
-    ).__get__(model, type(model))
-
     if torch.__version__ >= "2" and sys.platform != "win32":
         model = torch.compile(model)
         # WIP (not generally replacing layers until pytorch 2.1)
@@ -621,10 +617,10 @@ def generate_and_tokenize_prompt(data_point, prompt_type=None, train_on_inputs=F
     assert tokenizer is not None
     prompt_dict = ''  # only for custom prompt_type
     assert prompt_type != PromptType.custom.name, "custom not setup for finetune"
-    full_prompt, _, _, _ = generate_prompt(data_point, prompt_type, prompt_dict, False, False)
+    full_prompt, _, _, _, _ = generate_prompt(data_point, prompt_type, prompt_dict, False, False, False)
     tokenized_full_prompt = tokenize(full_prompt, tokenizer, cutoff_len, add_eos_token=add_eos_token)
     if not train_on_inputs:
-        user_prompt, _, _, _ = generate_prompt({**data_point, "output": ""}, prompt_type, prompt_dict, False, False)
+        user_prompt, _, _, _, _ = generate_prompt({**data_point, "output": ""}, prompt_type, prompt_dict, False, False, False)
         tokenized_user_prompt = tokenize(user_prompt, tokenizer, cutoff_len, add_eos_token=add_eos_token)
         user_prompt_len = len(tokenized_user_prompt["input_ids"])
         if add_eos_token:
@@ -643,7 +639,7 @@ def test_debug():
     fire.Fire(train)
 
 
-if __name__ == "__main__":
+def entrypoint_main():
     CONFIG = "NCCL_P2P_LEVEL=LOC WORLD_SIZE=5 torchrun --nnodes=5 --master_addr=10.10.10.2 --master_port=1111 --nproc_per_node=1"
     CMD = "finetune.py --data_path=config.json --num_epochs=1 --base_model=decapoda-research/llama-13b-hf"
     log(f"""
@@ -674,3 +670,7 @@ NCCL_P2P_LEVEL=LOC WORLD_SIZE=7 CUDA_VISIBLE_DEVICES="0,1" torchrun --node_rank
             "CUDA_VISIBLE_DEVICES") is not None, "Run python script using: torchrun finetune.py OR set CUDA_VISIBLE_DEVICES to single GPU"
 
     fire.Fire(train)
+
+
+if __name__ == "__main__":
+    entrypoint_main()

generate.py

@@ -1,27 +1,37 @@
 import ast
+import copy
 import functools
 import glob
 import inspect
 import queue
-import shutil
 import sys
 import os
 import time
 import traceback
+import types
 import typing
 import warnings
 from datetime import datetime
 import filelock
+import requests
 import psutil
+from requests import ConnectTimeout, JSONDecodeError
+from urllib3.exceptions import ConnectTimeoutError, MaxRetryError, ConnectionError
+from requests.exceptions import ConnectionError as ConnectionError2
+from requests.exceptions import ReadTimeout as ReadTimeout2
+
+if os.path.dirname(os.path.abspath(__file__)) not in sys.path:
+    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
 os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
 os.environ['BITSANDBYTES_NOWELCOME'] = '1'
 warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
 
-from enums import DocumentChoices, LangChainMode
+from enums import DocumentChoices, LangChainMode, no_lora_str, model_token_mapping, no_model_str, source_prefix, \
+    source_postfix
 from loaders import get_loaders
 from utils import set_seed, clear_torch_cache, save_generate_output, NullContext, wrapped_partial, EThread, get_githash, \
-    import_matplotlib, get_device, makedirs, get_kwargs, start_faulthandler
+    import_matplotlib, get_device, makedirs, get_kwargs, start_faulthandler, get_hf_server, FakeTokenizer, remove
 
 start_faulthandler()
 import_matplotlib()
@@ -34,9 +44,8 @@ from typing import Union
 import fire
 import torch
 from transformers import GenerationConfig, AutoModel, TextIteratorStreamer
-from accelerate import init_empty_weights, infer_auto_device_map
 
-from prompter import Prompter, inv_prompt_type_to_model_lower, non_hf_types, PromptType, get_prompt
+from prompter import Prompter, inv_prompt_type_to_model_lower, non_hf_types, PromptType, get_prompt, generate_prompt
 from stopping import get_stopping
 
 eval_extra_columns = ['prompt', 'response', 'score']
@@ -56,9 +65,15 @@ def main(
         lora_weights: str = "",
         gpu_id: int = 0,
         compile_model: bool = True,
-
+        use_cache: bool = None,
+        inference_server: str = "",
         prompt_type: Union[int, str] = None,
         prompt_dict: typing.Dict = None,
+
+        model_lock: typing.List[typing.Dict[str, str]] = None,
+        model_lock_columns: int = None,
+        fail_if_cannot_connect: bool = False,
+
         # input to generation
         temperature: float = None,
         top_p: float = None,
@@ -88,14 +103,13 @@ def main(
         cli: bool = False,
         cli_loop: bool = True,
         gradio: bool = True,
-        gradio_avoid_processing_markdown: bool = False,
         gradio_offline_level: int = 0,
         chat: bool = True,
         chat_context: bool = False,
         stream_output: bool = True,
         show_examples: bool = None,
         verbose: bool = False,
-        h2ocolors: bool = False,
+        h2ocolors: bool = True,
         height: int = 600,
         show_lora: bool = True,
         login_mode_if_model0: bool = False,
@@ -104,16 +118,19 @@ def main(
         api_open: bool = False,
         allow_api: bool = True,
         input_lines: int = 1,
+        gradio_size: str = None,
         auth: typing.List[typing.Tuple[str, str]] = None,
+        max_max_time=None,
+        max_max_new_tokens=None,
 
-        sanitize_user_prompt: bool = True,
-        sanitize_bot_response: bool = True,
+        sanitize_user_prompt: bool = False,
+        sanitize_bot_response: bool = False,
 
         extra_model_options: typing.List[str] = [],
         extra_lora_options: typing.List[str] = [],
+        extra_server_options: typing.List[str] = [],
 
         score_model: str = 'OpenAssistant/reward-model-deberta-v3-large-v2',
-        auto_score: bool = True,
 
         eval_filename: str = None,
         eval_prompts_only_num: int = 0,
@@ -121,6 +138,7 @@ def main(
         eval_as_output: bool = False,
 
         langchain_mode: str = 'Disabled',
+        force_langchain_evaluate: bool = False,
         visible_langchain_modes: list = ['UserData', 'MyData'],
         document_choice: list = [DocumentChoices.All_Relevant.name],
         user_path: str = None,
@@ -138,7 +156,10 @@ def main(
         enable_sources_list: bool = True,
         chunk: bool = True,
         chunk_size: int = 512,
-        top_k_docs: int = 3,  # FIXME: Can go back to 4 once https://github.com/h2oai/h2ogpt/issues/192 fixed
+        top_k_docs: int = None,
+        reverse_docs: bool = True,
+        auto_reduce_chunks: bool = True,
+        max_chunks: int = 100,
         n_jobs: int = -1,
         enable_captions: bool = True,
         captions_model: str = "Salesforce/blip-image-captioning-base",
@@ -157,8 +178,31 @@ def main(
     :param lora_weights: LORA weights path/HF link
     :param gpu_id: if infer_devices, then use gpu_id for cuda device ID, or auto mode if gpu_id != -1
     :param compile_model Whether to compile the model
+    :param use_cache: Whether to use caching in model (some models fail when multiple threads use)
+    :param inference_server: Consume base_model as type of model at this address
+                             Address can be text-generation-server hosting that base_model
+                             e.g. python generate.py --inference_server="http://192.168.1.46:6112" --base_model=h2oai/h2ogpt-oasst1-512-12b
+                             Or Address can be "openai_chat" or "openai" for OpenAI API
+                             e.g. python generate.py --inference_server="openai_chat" --base_model=gpt-3.5-turbo
+                             e.g. python generate.py --inference_server="openai" --base_model=text-davinci-003
     :param prompt_type: type of prompt, usually matched to fine-tuned model or plain for foundational model
     :param prompt_dict: If prompt_type=custom, then expects (some) items returned by get_prompt(..., return_dict=True)
+    :param model_lock: Lock models to specific combinations, for ease of use and extending to many models
+           Only used if gradio = True
+           List of dicts, each dict has base_model, tokenizer_base_model, lora_weights, inference_server, prompt_type, and prompt_dict
+           If all models have same prompt_type, and prompt_dict, can still specify that once in CLI outside model_lock as default for dict
+           Can specify model_lock instead of those items on CLI
+           As with CLI itself, base_model can infer prompt_type and prompt_dict if in prompter.py.
+             Also, tokenizer_base_model and lora_weights are optional.
+             Also, inference_server is optional if loading model from local system.
+           All models provided will automatically appear in compare model mode
+           Model loading-unloading and related choices will be disabled.  Model/lora/server adding will be disabled
+    :param model_lock_columns: How many columns to show if locking models (and so showing all at once)
+           If None, then defaults to up to 3
+           if -1, then all goes into 1 row
+           Maximum value is 4 due to non-dynamic gradio rendering elements
+    :param fail_if_cannot_connect: if doing model locking (e.g. with many models), fail if True.  Otherwise ignore.
+           Useful when many endpoints and want to just see what works, but still have to wait for timeout.
     :param temperature: generation temperature
     :param top_p: generation top_p
     :param top_k: generation top_k
@@ -184,13 +228,13 @@ def main(
     :param cli: whether to use CLI (non-gradio) interface.
     :param cli_loop: whether to loop for CLI (False usually only for testing)
     :param gradio: whether to enable gradio, or to enable benchmark mode
-    :param gradio_avoid_processing_markdown:
     :param gradio_offline_level: > 0, then change fonts so full offline
            == 1 means backend won't need internet for fonts, but front-end UI might if font not cached
            == 2 means backend and frontend don't need internet to download any fonts.
            Note: Some things always disabled include HF telemetry, gradio telemetry, chromadb posthog that involve uploading.
            This option further disables google fonts for downloading, which is less intrusive than uploading,
            but still required in air-gapped case.  The fonts don't look as nice as google fonts, but ensure full offline behavior.
+           Also set --share=False to avoid sharing a gradio live link.
     :param chat: whether to enable chat mode with chat history
     :param chat_context: whether to use extra helpful context if human_bot
     :param stream_output: whether to stream output from generate
@@ -205,20 +249,25 @@ def main(
     :param api_open: If False, don't let API calls skip gradio queue
     :param allow_api: whether to allow API calls at all to gradio server
     :param input_lines: how many input lines to show for chat box (>1 forces shift-enter for submit, else enter is submit)
+    :param gradio_size: Overall size of text and spaces: "xsmall", "small", "medium", "large".
+           Small useful for many chatbots in model_lock mode
     :param auth: gradio auth for launcher in form [(user1, pass1), (user2, pass2), ...]
                  e.g. --auth=[('jon','password')] with no spaces
-    :param sanitize_user_prompt: whether to remove profanity from user input
-    :param sanitize_bot_response: whether to remove profanity and repeat lines from bot output
+    :param max_max_time: Maximum max_time for gradio slider
+    :param max_max_new_tokens: Maximum max_new_tokens for gradio slider
+    :param sanitize_user_prompt: whether to remove profanity from user input (slows down input processing)
+    :param sanitize_bot_response: whether to remove profanity and repeat lines from bot output (about 2x slower generation for long streaming cases due to better_profanity being slow)
     :param extra_model_options: extra models to show in list in gradio
     :param extra_lora_options: extra LORA to show in list in gradio
+    :param extra_server_options: extra servers to show in list in gradio
     :param score_model: which model to score responses (None means no scoring)
-    :param auto_score: whether to automatically score responses
     :param eval_filename: json file to use for evaluation, if None is sharegpt
     :param eval_prompts_only_num: for no gradio benchmark, if using eval_filename prompts for eval instead of examples
     :param eval_prompts_only_seed: for no gradio benchmark, seed for eval_filename sampling
     :param eval_as_output: for no gradio benchmark, whether to test eval_filename output itself
     :param langchain_mode: Data source to include.  Choose "UserData" to only consume files from make_db.py.
            WARNING: wiki_full requires extra data processing via read_wiki_full.py and requires really good workstation to generate db, unless already present.
+    :param force_langchain_evaluate: Whether to force langchain LLM use even if not doing langchain, mostly for testing.
     :param user_path: user path to glob from to generate db for vector search, for 'UserData' langchain mode.
            If already have db, any new/changed files are added automatically if path set, does not have to be same path used for prior db sources
     :param detect_user_path_changes_every_query: whether to detect if any files changed or added every similarity search (by file hashes).
@@ -248,12 +297,17 @@ def main(
     :param chunk: Whether to chunk data (True unless know data is already optimally chunked)
     :param chunk_size: Size of chunks, with typically top-4 passed to LLM, so neesd to be in context length
     :param top_k_docs: number of chunks to give LLM
+    :param reverse_docs: whether to reverse docs order so most relevant is closest to question.
+           Best choice for sufficiently smart model, and truncation occurs for oldest context, so best then too.
+           But smaller 6_9 models fail to use newest context and can get stuck on old information.
+    :param auto_reduce_chunks: Whether to automatically reduce top_k_docs to fit context given prompt
+    :param max_chunks: If top_k_docs=-1, maximum number of chunks to allow
     :param n_jobs: Number of processors to use when consuming documents (-1 = all, is default)
     :param enable_captions: Whether to support captions using BLIP for image files as documents, then preloads that model
     :param captions_model: Which model to use for captions.
-           captions_model: int = "Salesforce/blip-image-captioning-base",  # continue capable
+           captions_model: str = "Salesforce/blip-image-captioning-base",  # continue capable
            captions_model: str = "Salesforce/blip2-flan-t5-xl",   # question/answer capable, 16GB state
-           captions_model: int = "Salesforce/blip2-flan-t5-xxl",  # question/answer capable, 60GB state
+           captions_model: str = "Salesforce/blip2-flan-t5-xxl",  # question/answer capable, 60GB state
            Note: opt-based blip2 are not permissive license due to opt and Meta license restrictions
     :param pre_load_caption_model: Whether to preload caption model, or load after forking parallel doc loader
            parallel loading disabled if preload and have images, to prevent deadlocking on cuda context
@@ -262,6 +316,32 @@ def main(
     :param enable_ocr: Whether to support OCR on images
     :return:
     """
+    if base_model is None:
+        base_model = ''
+    if tokenizer_base_model is None:
+        tokenizer_base_model = ''
+    if lora_weights is None:
+        lora_weights = ''
+    if inference_server is None:
+        inference_server = ''
+
+    # listen to env if set
+    model_lock = os.getenv('model_lock', str(model_lock))
+    model_lock = ast.literal_eval(model_lock)
+
+    if model_lock:
+        assert gradio, "model_lock only supported for gradio=True"
+        if len(model_lock) > 1:
+            assert chat, "model_lock only works for multiple models for chat=True"
+        assert not cli, "model_lock only supported for cli=False"
+        assert not (not cli and not gradio), "model_lock only supported for eval (cli=gradio=False)"
+        assert not base_model, "Don't specify model_lock and base_model"
+        assert not tokenizer_base_model, "Don't specify model_lock and tokenizer_base_model"
+        assert not lora_weights, "Don't specify model_lock and lora_weights"
+        assert not inference_server, "Don't specify model_lock and inference_server"
+        # assert not prompt_type, "Don't specify model_lock and prompt_type"
+        # assert not prompt_dict, "Don't specify model_lock and prompt_dict"
+
     is_hf = bool(int(os.getenv("HUGGINGFACE_SPACES", '0')))
     is_gpth2oai = bool(int(os.getenv("GPT_H2O_AI", '0')))
     is_public = is_hf or is_gpth2oai  # multi-user case with fixed model and disclaimer
@@ -276,7 +356,8 @@ def main(
 
     # allow set token directly
     use_auth_token = os.environ.get("HUGGINGFACE_API_TOKEN", use_auth_token)
-    allow_upload_to_user_data = bool(int(os.environ.get("allow_upload_to_user_data", str(int(allow_upload_to_user_data)))))
+    allow_upload_to_user_data = bool(
+        int(os.environ.get("allow_upload_to_user_data", str(int(allow_upload_to_user_data)))))
     allow_upload_to_my_data = bool(int(os.environ.get("allow_upload_to_my_data", str(int(allow_upload_to_my_data)))))
     height = int(os.environ.get("HEIGHT", height))
     h2ocolors = bool(int(os.getenv('h2ocolors', h2ocolors)))
@@ -289,6 +370,12 @@ def main(
     if langchain_mode not in visible_langchain_modes and langchain_mode in langchain_modes:
         visible_langchain_modes += [langchain_mode]
 
+    # if specifically chose not to show My or User Data, disable upload, so gradio elements are simpler
+    if LangChainMode.MY_DATA.value not in visible_langchain_modes:
+        allow_upload_to_my_data = False
+    if LangChainMode.USER_DATA.value not in visible_langchain_modes:
+        allow_upload_to_user_data = False
+
     if is_public:
         allow_upload_to_user_data = False
         input_lines = 1  # ensure set, for ease of use
@@ -297,26 +384,50 @@ def main(
         top_k = 70 if top_k is None else top_k
         if is_hf:
             do_sample = True if do_sample is None else do_sample
+            top_k_docs = 3 if top_k_docs is None else top_k_docs
         else:
             # by default don't sample, too chatty
             do_sample = False if do_sample is None else do_sample
+            top_k_docs = 4 if top_k_docs is None else top_k_docs
 
         if memory_restriction_level == 2:
-            if not base_model:
+            if not base_model and not inference_server:
                 base_model = 'h2oai/h2ogpt-oasst1-512-12b'
                 # don't set load_8bit if passed base_model, doesn't always work so can't just override
                 load_8bit = True
                 load_4bit = False  # FIXME - consider using 4-bit instead of 8-bit
-        else:
-            base_model = 'h2oai/h2ogpt-oasst1-512-20b' if not base_model else base_model
+        elif not inference_server:
+            top_k_docs = 10 if top_k_docs is None else top_k_docs
     if memory_restriction_level >= 2:
         load_8bit = True
         load_4bit = False  # FIXME - consider using 4-bit instead of 8-bit
         if hf_embedding_model is None:
             hf_embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
+        top_k_docs = 3 if top_k_docs is None else top_k_docs
+    if top_k_docs is None:
+        top_k_docs = 3
+    if is_public:
+        if not max_time:
+            max_time = 60 * 2
+        if not max_max_time:
+            max_max_time = max_time
+        if not max_new_tokens:
+            max_new_tokens = 256
+        if not max_max_new_tokens:
+            max_max_new_tokens = 256
+    else:
+        if not max_max_time:
+            max_max_time = 60 * 20
+        if not max_max_new_tokens:
+            max_max_new_tokens = 512
     if is_hf:
         # must override share if in spaces
         share = False
+        if not max_time:
+            max_time = 60 * 1
+        if not max_max_time:
+            max_max_time = max_time
+        # HF accounted for later in get_max_max_new_tokens()
     save_dir = os.getenv('SAVE_DIR', save_dir)
     score_model = os.getenv('SCORE_MODEL', score_model)
     if score_model == 'None' or score_model is None:
@@ -335,7 +446,7 @@ def main(
         torch.backends.cudnn.benchmark = True
         torch.backends.cudnn.enabled = False
         torch.set_default_dtype(torch.float32)
-        if psutil.virtual_memory().available < 94 * 1024 ** 3:
+        if psutil.virtual_memory().available < 94 * 1024 ** 3 and not inference_server:
             # 12B uses ~94GB
             # 6.9B uses ~47GB
             base_model = 'h2oai/h2ogpt-oig-oasst1-512-6_9b' if not base_model else base_model
@@ -360,8 +471,8 @@ def main(
 
     if offload_folder:
         makedirs(offload_folder)
-
-    user_set_max_new_tokens = max_new_tokens is not None
+    if user_path:
+        makedirs(user_path)
 
     placeholder_instruction, placeholder_input, \
         stream_output, show_examples, \
@@ -386,11 +497,12 @@ def main(
                             verbose,
                             )
 
+    git_hash = get_githash()
     locals_dict = locals()
     locals_print = '\n'.join(['%s: %s' % (k, v) for k, v in locals_dict.items()])
     if verbose:
         print(f"Generating model with params:\n{locals_print}", flush=True)
-        print("Command: %s\nHash: %s" % (str(' '.join(sys.argv)), get_githash()), flush=True)
+        print("Command: %s\nHash: %s" % (str(' '.join(sys.argv)), git_hash), flush=True)
 
     if langchain_mode != "Disabled":
         # SECOND PLACE where LangChain referenced, but all imports are kept local so not required
@@ -404,7 +516,7 @@ def main(
                 for gpath1 in glob.glob(os.path.join(scratch_base_dir, 'db_dir_%s*' % langchain_mode1)):
                     if os.path.isdir(gpath1):
                         print("Removing old MyData: %s" % gpath1, flush=True)
-                        shutil.rmtree(gpath1)
+                        remove(gpath1)
                 continue
             if langchain_mode1 in ['All']:
                 # FIXME: All should be avoided until scans over each db, shouldn't be separate db
@@ -430,6 +542,10 @@ def main(
             assert 'gpt_langchain' not in sys.modules, "Dev bug, import of langchain when should not have"
             assert 'langchain' not in sys.modules, "Dev bug, import of langchain when should not have"
 
+    model_state_none = dict(model=None, tokenizer=None, device=None,
+                            base_model=None, tokenizer_base_model=None, lora_weights=None,
+                            inference_server=None, prompt_type=None, prompt_dict=None)
+
     if cli:
         from cli import run_cli
         return run_cli(**get_kwargs(run_cli, exclude_names=['model_state0'], **locals()))
@@ -441,20 +557,68 @@ def main(
         from gradio_runner import go_gradio
 
         # get default model
-        all_kwargs = locals().copy()
-        if all_kwargs.get('base_model') and not all_kwargs['login_mode_if_model0']:
-            model0, tokenizer0, device = get_model(reward_type=False,
-                                                   **get_kwargs(get_model, exclude_names=['reward_type'], **all_kwargs))
-        else:
-            # if empty model, then don't load anything, just get gradio up
-            model0, tokenizer0, device = None, None, None
-        model_state0 = [model0, tokenizer0, device, all_kwargs['base_model']]
+        model_states = []
+        model_list = [dict(base_model=base_model, tokenizer_base_model=tokenizer_base_model, lora_weights=lora_weights,
+                           inference_server=inference_server, prompt_type=prompt_type, prompt_dict=prompt_dict)]
+        model_list0 = copy.deepcopy(model_list)  # just strings, safe to deepcopy
+        model_state0 = model_state_none.copy()
+        assert len(model_state_none) == len(model_state0)
+        if model_lock:
+            model_list = model_lock
+        for model_dict in reversed(model_list):
+            # do reverse, so first is default base_model etc., so some logic works in go_gradio() more easily
+            # handles defaults user didn't have to pass
+            model_dict['base_model'] = base_model = model_dict.get('base_model', '')
+            model_dict['tokenizer_base_model'] = tokenizer_base_model = model_dict.get('tokenizer_base_model', '')
+            model_dict['lora_weights'] = lora_weights = model_dict.get('lora_weights', '')
+            model_dict['inference_server'] = inference_server = model_dict.get('inference_server', '')
+            prompt_type = model_dict.get('prompt_type', model_list0[0]['prompt_type'])  # don't use mutated value
+            # try to infer, ignore empty initial state leading to get_generate_params -> 'plain'
+            if model_dict.get('prompt_type') is None:
+                model_lower = base_model.lower()
+                if model_lower in inv_prompt_type_to_model_lower:
+                    prompt_type = inv_prompt_type_to_model_lower[model_lower]
+                    prompt_dict, error0 = get_prompt(prompt_type, '',
+                                                     chat=False, context='', reduced=False, making_context=False,
+                                                     return_dict=True)
+            model_dict['prompt_type'] = prompt_type
+            model_dict['prompt_dict'] = prompt_dict = model_dict.get('prompt_dict', prompt_dict)
+            all_kwargs = locals().copy()
+            if base_model and not login_mode_if_model0:
+                model0, tokenizer0, device = get_model(reward_type=False,
+                                                       **get_kwargs(get_model, exclude_names=['reward_type'],
+                                                                    **all_kwargs))
+            else:
+                # if empty model, then don't load anything, just get gradio up
+                model0, tokenizer0, device = None, None, None
+            if model0 is None:
+                if fail_if_cannot_connect:
+                    raise RuntimeError("Could not connect, see logs")
+                # skip
+                if isinstance(model_lock, list):
+                    model_lock.remove(model_dict)
+                continue
+            model_state_trial = dict(model=model0, tokenizer=tokenizer0, device=device)
+            model_state_trial.update(model_dict)
+            assert len(model_state_none) == len(model_state_trial)
+            print("Model %s" % model_dict, flush=True)
+            if model_lock:
+                # last in iteration will be first
+                model_states.insert(0, model_state_trial)
+                # fill model_state0 so go_gradio() easier, manage model_states separately
+                model_state0 = model_state_trial.copy()
+            else:
+                model_state0 = model_state_trial.copy()
+            assert len(model_state_none) == len(model_state0)
 
         # get score model
+        all_kwargs = locals().copy()
         smodel, stokenizer, sdevice = get_score_model(reward_type=True,
                                                       **get_kwargs(get_score_model, exclude_names=['reward_type'],
                                                                    **all_kwargs))
-        score_model_state0 = [smodel, stokenizer, sdevice, score_model]
+        score_model_state0 = dict(model=smodel, tokenizer=stokenizer, device=sdevice,
+                                  base_model=score_model, tokenizer_base_model='', lora_weights='',
+                                  inference_server='', prompt_type='', prompt_dict='')
 
         if enable_captions:
             if pre_load_caption_model:
@@ -469,34 +633,33 @@ def main(
         go_gradio(**locals())
 
 
-def get_non_lora_model(base_model, model_loader, load_half, model_kwargs, reward_type,
-                       gpu_id=0,
-                       use_auth_token=False,
-                       trust_remote_code=True,
-                       offload_folder=None,
-                       triton_attn=False,
-                       long_sequence=True,
-                       ):
-    """
-    Ensure model gets on correct device
-    :param base_model:
-    :param model_loader:
-    :param load_half:
-    :param model_kwargs:
-    :param reward_type:
-    :param gpu_id:
-    :param use_auth_token:
-    :param trust_remote_code:
-    :param offload_folder:
-    :param triton_attn:
-    :param long_sequence:
-    :return:
-    """
+def get_config(base_model,
+               use_auth_token=False,
+               trust_remote_code=True,
+               offload_folder=None,
+               triton_attn=False,
+               long_sequence=True,
+               return_model=False,
+               raise_exception=False,
+               ):
+    from accelerate import init_empty_weights
     with init_empty_weights():
         from transformers import AutoConfig
-        config = AutoConfig.from_pretrained(base_model, use_auth_token=use_auth_token,
-                                            trust_remote_code=trust_remote_code,
-                                            offload_folder=offload_folder)
+        try:
+            config = AutoConfig.from_pretrained(base_model, use_auth_token=use_auth_token,
+                                                trust_remote_code=trust_remote_code,
+                                                offload_folder=offload_folder)
+        except OSError as e:
+            if raise_exception:
+                raise
+            if 'not a local folder and is not a valid model identifier listed on' in str(
+                    e) or '404 Client Error' in str(e):
+                # e.g. llama, gpjt, etc.
+                # e.g. HF TGI but not model on HF or private etc.
+                # HF TGI server only should really require prompt_type, not HF model state
+                return None, None
+            else:
+                raise
         if triton_attn and 'mpt-' in base_model.lower():
             config.attn_config['attn_impl'] = 'triton'
         if long_sequence:
@@ -504,18 +667,36 @@ def get_non_lora_model(base_model, model_loader, load_half, model_kwargs, reward
                 config.update({"max_seq_len": 83968})
             if 'mosaicml/mpt-7b-chat' in base_model.lower():
                 config.update({"max_seq_len": 4096})
-        if issubclass(config.__class__, tuple(AutoModel._model_mapping.keys())):
+            if 'mpt-30b' in base_model.lower():
+                config.update({"max_seq_len": 2 * 8192})
+        if return_model and \
+                issubclass(config.__class__, tuple(AutoModel._model_mapping.keys())):
             model = AutoModel.from_config(
                 config,
+                trust_remote_code=trust_remote_code,
             )
         else:
             # can't infer
             model = None
+    if 'falcon' in base_model.lower():
+        config.use_cache = False
+
+    return config, model
+
+
+def get_non_lora_model(base_model, model_loader, load_half, model_kwargs, reward_type,
+                       config, model,
+                       gpu_id=0,
+                       ):
+    """
+    Ensure model gets on correct device
+    """
 
     if model is not None:
         # NOTE: Can specify max_memory={0: max_mem, 1: max_mem}, to shard model
         # NOTE: Some models require avoiding sharding some layers,
         # then would pass no_split_module_classes and give list of those layers.
+        from accelerate import infer_auto_device_map
         device_map = infer_auto_device_map(
             model,
             dtype=torch.float16 if load_half else torch.float32,
@@ -567,12 +748,59 @@ def get_non_lora_model(base_model, model_loader, load_half, model_kwargs, reward
     return model
 
 
+def get_client_from_inference_server(inference_server, raise_connection_exception=False):
+    inference_server, headers = get_hf_server(inference_server)
+    # preload client since slow for gradio case especially
+    from gradio_utils.grclient import GradioClient
+    gr_client = None
+    hf_client = None
+    if headers is None:
+        try:
+            print("GR Client Begin: %s" % inference_server, flush=True)
+            # first do sanity check if alive, else gradio client takes too long by default
+            requests.get(inference_server, timeout=int(os.getenv('REQUEST_TIMEOUT', '30')))
+            gr_client = GradioClient(inference_server)
+            print("GR Client End: %s" % inference_server, flush=True)
+        except (OSError, ValueError) as e:
+            # Occurs when wrong endpoint and should have been HF client, so don't hard raise, just move to HF
+            gr_client = None
+            print("GR Client Failed %s: %s" % (inference_server, str(e)), flush=True)
+        except (ConnectTimeoutError, ConnectTimeout, MaxRetryError, ConnectionError, ConnectionError2,
+                JSONDecodeError, ReadTimeout2, KeyError) as e:
+            t, v, tb = sys.exc_info()
+            ex = ''.join(traceback.format_exception(t, v, tb))
+            print("GR Client Failed %s: %s" % (inference_server, str(ex)), flush=True)
+            if raise_connection_exception:
+                raise
+
+    if gr_client is None:
+        res = None
+        from text_generation import Client as HFClient
+        print("HF Client Begin: %s" % inference_server)
+        try:
+            hf_client = HFClient(inference_server, headers=headers, timeout=int(os.getenv('REQUEST_TIMEOUT', '30')))
+            # quick check valid TGI endpoint
+            res = hf_client.generate('What?', max_new_tokens=1)
+            hf_client = HFClient(inference_server, headers=headers, timeout=300)
+        except (ConnectTimeoutError, ConnectTimeout, MaxRetryError, ConnectionError, ConnectionError2,
+                JSONDecodeError, ReadTimeout2, KeyError) as e:
+            hf_client = None
+            t, v, tb = sys.exc_info()
+            ex = ''.join(traceback.format_exception(t, v, tb))
+            print("HF Client Failed %s: %s" % (inference_server, str(ex)))
+            if raise_connection_exception:
+                raise
+        print("HF Client End: %s %s" % (inference_server, res))
+    return inference_server, gr_client, hf_client
+
+
 def get_model(
         load_8bit: bool = False,
         load_4bit: bool = False,
         load_half: bool = True,
         infer_devices: bool = True,
         base_model: str = '',
+        inference_server: str = "",
         tokenizer_base_model: str = '',
         lora_weights: str = "",
         gpu_id: int = 0,
@@ -596,6 +824,7 @@ def get_model(
            For non-LORA case, False will spread shards across multiple GPUs, but this can lead to cuda:x cuda:y mismatches
            So it is not the default
     :param base_model: name/path of base model
+    :param inference_server: whether base_model is hosted locally ('') or via http (url)
     :param tokenizer_base_model: name/path of tokenizer
     :param lora_weights: name/path
     :param gpu_id: which GPU (0..n_gpus-1) or allow all GPUs if relevant (-1)
@@ -611,11 +840,120 @@ def get_model(
     """
     if verbose:
         print("Get %s model" % base_model, flush=True)
+
+    triton_attn = False
+    long_sequence = True
+    config_kwargs = dict(use_auth_token=use_auth_token,
+                         trust_remote_code=trust_remote_code,
+                         offload_folder=offload_folder,
+                         triton_attn=triton_attn,
+                         long_sequence=long_sequence)
+    config, _ = get_config(base_model, **config_kwargs, raise_exception=False)
+
+    if base_model in non_hf_types:
+        assert config is None, "Expected config None for %s" % base_model
+
+    llama_type_from_config = 'llama' in str(config).lower()
+    llama_type_from_name = "llama" in base_model.lower()
+    llama_type = llama_type_from_config or llama_type_from_name
+    if "xgen" in base_model.lower():
+        llama_type = False
+    if llama_type:
+        if verbose:
+            print("Detected as llama type from"
+                  " config (%s) or name (%s)" % (llama_type_from_config, llama_type_from_name), flush=True)
+
+    model_loader, tokenizer_loader = get_loaders(model_name=base_model, reward_type=reward_type, llama_type=llama_type)
+
+    tokenizer_kwargs = dict(local_files_only=local_files_only,
+                            resume_download=resume_download,
+                            use_auth_token=use_auth_token,
+                            trust_remote_code=trust_remote_code,
+                            offload_folder=offload_folder,
+                            padding_side='left',
+                            config=config,
+                            )
+    if not tokenizer_base_model:
+        tokenizer_base_model = base_model
+
+    if config is not None and tokenizer_loader is not None and not isinstance(tokenizer_loader, str):
+        tokenizer = tokenizer_loader.from_pretrained(tokenizer_base_model, **tokenizer_kwargs)
+        # sets raw (no cushion) limit
+        set_model_max_len(config, tokenizer, verbose=False)
+        # if using fake tokenizer, not really accurate when lots of numbers, give a bit of buffer, else get:
+        # Generation Failed: Input validation error: `inputs` must have less than 2048 tokens. Given: 2233
+        tokenizer.model_max_length = tokenizer.model_max_length - 50
+    else:
+        tokenizer = FakeTokenizer()
+
+    if isinstance(inference_server, str) and inference_server.startswith("http"):
+        inference_server, gr_client, hf_client = get_client_from_inference_server(inference_server)
+        client = gr_client or hf_client
+        # Don't return None, None for model, tokenizer so triggers
+        return client, tokenizer, 'http'
+    if isinstance(inference_server, str) and inference_server.startswith('openai'):
+        assert os.getenv('OPENAI_API_KEY'), "Set environment for OPENAI_API_KEY"
+        # Don't return None, None for model, tokenizer so triggers
+        # include small token cushion
+        tokenizer = FakeTokenizer(model_max_length=model_token_mapping[base_model] - 50)
+        return inference_server, tokenizer, inference_server
+    assert not inference_server, "Malformed inference_server=%s" % inference_server
     if base_model in non_hf_types:
         from gpt4all_llm import get_model_tokenizer_gpt4all
         model, tokenizer, device = get_model_tokenizer_gpt4all(base_model)
         return model, tokenizer, device
 
+    # get local torch-HF model
+    return get_hf_model(load_8bit=load_8bit,
+                        load_4bit=load_4bit,
+                        load_half=load_half,
+                        infer_devices=infer_devices,
+                        base_model=base_model,
+                        tokenizer_base_model=tokenizer_base_model,
+                        lora_weights=lora_weights,
+                        gpu_id=gpu_id,
+
+                        reward_type=reward_type,
+                        local_files_only=local_files_only,
+                        resume_download=resume_download,
+                        use_auth_token=use_auth_token,
+                        trust_remote_code=trust_remote_code,
+                        offload_folder=offload_folder,
+                        compile_model=compile_model,
+
+                        llama_type=llama_type,
+                        config_kwargs=config_kwargs,
+                        tokenizer_kwargs=tokenizer_kwargs,
+
+                        verbose=verbose)
+
+
+def get_hf_model(load_8bit: bool = False,
+                 load_4bit: bool = False,
+                 load_half: bool = True,
+                 infer_devices: bool = True,
+                 base_model: str = '',
+                 tokenizer_base_model: str = '',
+                 lora_weights: str = "",
+                 gpu_id: int = 0,
+
+                 reward_type: bool = None,
+                 local_files_only: bool = False,
+                 resume_download: bool = True,
+                 use_auth_token: Union[str, bool] = False,
+                 trust_remote_code: bool = True,
+                 offload_folder: str = None,
+                 compile_model: bool = True,
+
+                 llama_type: bool = False,
+                 config_kwargs=None,
+                 tokenizer_kwargs=None,
+
+                 verbose: bool = False,
+                 ):
+    assert config_kwargs is not None
+    assert tokenizer_kwargs is not None
+
     if lora_weights is not None and lora_weights.strip():
         if verbose:
             print("Get %s lora weights" % lora_weights, flush=True)
@@ -630,30 +968,13 @@ def get_model(
         "Please choose a base model with --base_model (CLI) or load one from Models Tab (gradio)"
     )
 
-    from transformers import AutoConfig
-    config = AutoConfig.from_pretrained(base_model, use_auth_token=use_auth_token,
-                                        trust_remote_code=trust_remote_code,
-                                        offload_folder=offload_folder)
-    llama_type_from_config = 'llama' in str(config).lower()
-    llama_type_from_name = "llama" in base_model.lower()
-    llama_type = llama_type_from_config or llama_type_from_name
-    if llama_type:
-        if verbose:
-            print("Detected as llama type from"
-                  " config (%s) or name (%s)" % (llama_type_from_config, llama_type_from_name), flush=True)
+    model_loader, tokenizer_loader = get_loaders(model_name=base_model, reward_type=reward_type, llama_type=llama_type)
 
-    model_loader, tokenizer_loader = get_loaders(llama_type=llama_type, model_name=base_model, reward_type=reward_type)
-    if not tokenizer_base_model:
-        tokenizer_base_model = base_model
+    config, _ = get_config(base_model, return_model=False, raise_exception=True, **config_kwargs)
 
     if tokenizer_loader is not None and not isinstance(tokenizer_loader, str):
         tokenizer = tokenizer_loader.from_pretrained(tokenizer_base_model,
-                                                     local_files_only=local_files_only,
-                                                     resume_download=resume_download,
-                                                     use_auth_token=use_auth_token,
-                                                     trust_remote_code=trust_remote_code,
-                                                     offload_folder=offload_folder,
-                                                     )
+                                                     **tokenizer_kwargs)
     else:
         tokenizer = tokenizer_loader
 
@@ -677,7 +998,7 @@ def get_model(
                                      load_in_4bit=load_4bit,
                                      device_map={"": 0} if (load_8bit or load_4bit) and device == 'cuda' else "auto",
                                      ))
-        if 'mpt-' in base_model.lower() and gpu_id >= 0:
+        if 'mpt-' in base_model.lower() and gpu_id is not None and gpu_id >= 0:
             model_kwargs.update(dict(device_map={"": gpu_id} if device == 'cuda' else "cpu"))
 
         if 'OpenAssistant/reward-model'.lower() in base_model.lower():
@@ -688,25 +1009,30 @@ def get_model(
 
         if not lora_weights:
             with torch.device(device):
+
                 if infer_devices:
+                    config, model = get_config(base_model, return_model=True, raise_exception=True, **config_kwargs)
                     model = get_non_lora_model(base_model, model_loader, load_half, model_kwargs, reward_type,
+                                               config, model,
                                                gpu_id=gpu_id,
-                                               use_auth_token=use_auth_token,
-                                               trust_remote_code=trust_remote_code,
-                                               offload_folder=offload_folder,
                                                )
                 else:
+                    config, _ = get_config(base_model, **config_kwargs)
                     if load_half and not (load_8bit or load_4bit):
                         model = model_loader.from_pretrained(
                             base_model,
+                            config=config,
                             **model_kwargs).half()
                     else:
                         model = model_loader.from_pretrained(
                             base_model,
+                            config=config,
                             **model_kwargs)
         elif load_8bit or load_4bit:
+            config, _ = get_config(base_model, **config_kwargs)
             model = model_loader.from_pretrained(
                 base_model,
+                config=config,
                 **model_kwargs
             )
             from peft import PeftModel  # loads cuda, so avoid in global scope
@@ -723,8 +1049,10 @@ def get_model(
             )
         else:
             with torch.device(device):
+                config, _ = get_config(base_model, raise_exception=True, **config_kwargs)
                 model = model_loader.from_pretrained(
                     base_model,
+                    config=config,
                     **model_kwargs
                 )
                 from peft import PeftModel  # loads cuda, so avoid in global scope
@@ -758,6 +1086,15 @@ def get_model(
         if torch.__version__ >= "2" and sys.platform != "win32" and compile_model:
             model = torch.compile(model)
 
+    set_model_max_len(config, tokenizer, verbose=False, reward_type=reward_type)
+
+    return model, tokenizer, device
+
+
+def set_model_max_len(config, tokenizer, verbose=False, reward_type=False):
+    if reward_type:
+        # limit deberta, else uses too much memory and not worth response score
+        tokenizer.model_max_length = 512
     if hasattr(config, 'max_seq_len') and isinstance(config.max_seq_len, int):
         tokenizer.model_max_length = config.max_seq_len
     elif hasattr(config, 'max_position_embeddings') and isinstance(config.max_position_embeddings, int):
@@ -767,8 +1104,9 @@ def get_model(
         if verbose:
             print("Could not determine model_max_length, setting to 2048", flush=True)
         tokenizer.model_max_length = 2048
-
-    return model, tokenizer, device
+    # for bug in HF transformers
+    if tokenizer.model_max_length > 100000000:
+        tokenizer.model_max_length = 2048
 
 
 def pop_unused_model_kwargs(model_kwargs):
@@ -790,6 +1128,7 @@ def get_score_model(score_model: str = None,
                     load_half: bool = True,
                     infer_devices: bool = True,
                     base_model: str = '',
+                    inference_server: str = '',
                     tokenizer_base_model: str = '',
                     lora_weights: str = "",
                     gpu_id: int = 0,
@@ -811,6 +1150,7 @@ def get_score_model(score_model: str = None,
         base_model = score_model.strip()
         tokenizer_base_model = ''
         lora_weights = ''
+        inference_server = ''
         llama_type = False
         compile_model = False
         smodel, stokenizer, sdevice = get_model(reward_type=True,
@@ -877,9 +1217,12 @@ def evaluate_from_str(
         debug=False,
         concurrency_count=None,
         save_dir=None,
-        sanitize_bot_response=True,
+        sanitize_bot_response=False,
         model_state0=None,
         memory_restriction_level=None,
+        max_max_new_tokens=None,
+        is_public=None,
+        max_max_time=None,
         raise_generate_gpu_exceptions=None,
         chat_context=None,
         lora_weights=None,
@@ -890,20 +1233,26 @@ def evaluate_from_str(
         use_openai_embedding=None,
         use_openai_model=None,
         hf_embedding_model=None,
-        chunk=None,
-        chunk_size=None,
         db_type=None,
         n_jobs=None,
         first_para=None,
         text_limit=None,
         verbose=False,
         cli=False,
+        reverse_docs=True,
+        use_cache=None,
+        auto_reduce_chunks=None,
+        max_chunks=None,
+        model_lock=None,
+        force_langchain_evaluate=None,
+        model_state_none=None,
 ):
     if isinstance(user_kwargs, str):
         user_kwargs = ast.literal_eval(user_kwargs)
     # only used for submit_nochat_api
     user_kwargs['chat'] = False
-    user_kwargs['stream_output'] = False
+    if 'stream_output' not in user_kwargs:
+        user_kwargs['stream_output'] = False
     if 'langchain_mode' not in user_kwargs:
         # if user doesn't specify, then assume disabled, not use default
         user_kwargs['langchain_mode'] = 'Disabled'
@@ -926,6 +1275,9 @@ def evaluate_from_str(
         sanitize_bot_response=sanitize_bot_response,
         model_state0=model_state0,
         memory_restriction_level=memory_restriction_level,
+        max_max_new_tokens=max_max_new_tokens,
+        is_public=is_public,
+        max_max_time=max_max_time,
         raise_generate_gpu_exceptions=raise_generate_gpu_exceptions,
         chat_context=chat_context,
         lora_weights=lora_weights,
@@ -942,6 +1294,13 @@ def evaluate_from_str(
         text_limit=text_limit,
         verbose=verbose,
         cli=cli,
+        reverse_docs=reverse_docs,
+        use_cache=use_cache,
+        auto_reduce_chunks=auto_reduce_chunks,
+        max_chunks=max_chunks,
+        model_lock=model_lock,
+        force_langchain_evaluate=force_langchain_evaluate,
+        model_state_none=model_state_none,
     )
     try:
         for ret1 in ret:
@@ -986,9 +1345,12 @@ def evaluate(
         debug=False,
         concurrency_count=None,
         save_dir=None,
-        sanitize_bot_response=True,
+        sanitize_bot_response=False,
         model_state0=None,
         memory_restriction_level=None,
+        max_max_new_tokens=None,
+        is_public=None,
+        max_max_time=None,
         raise_generate_gpu_exceptions=None,
         chat_context=None,
         lora_weights=None,
@@ -1005,6 +1367,13 @@ def evaluate(
         text_limit=None,
         verbose=False,
         cli=False,
+        reverse_docs=True,
+        use_cache=None,
+        auto_reduce_chunks=None,
+        max_chunks=None,
+        model_lock=None,
+        force_langchain_evaluate=None,
+        model_state_none=None,
 ):
     # ensure passed these
     assert concurrency_count is not None
@@ -1025,29 +1394,58 @@ def evaluate(
         locals_dict = locals().copy()
         locals_dict.pop('model_state', None)
         locals_dict.pop('model_state0', None)
+        locals_dict.pop('model_states', None)
         print(locals_dict)
 
-    no_model_msg = "Please choose a base model with --base_model (CLI) or load in Models Tab (gradio).\nThen start New Conversation"
+    no_model_msg = "Please choose a base model with --base_model (CLI) or load in Models Tab (gradio).\n" \
+                   "Then start New Conversation"
 
+    if model_state is None:
+        model_state = model_state_none.copy()
     if model_state0 is None:
         # e.g. for no gradio case, set dummy value, else should be set
-        model_state0 = [None, None, None, None]
-
-    if model_state is not None and len(model_state) == 4 and not isinstance(model_state[0], str):
-        # try to free-up original model (i.e. list was passed as reference)
-        if model_state0 is not None and model_state0[0] is not None:
-            model_state0[0].cpu()
-            model_state0[0] = None
-        # try to free-up original tokenizer (i.e. list was passed as reference)
-        if model_state0 is not None and model_state0[1] is not None:
-            model_state0[1] = None
-        clear_torch_cache()
-        model, tokenizer, device, base_model = model_state
-    elif model_state0 is not None and len(model_state0) == 4 and model_state0[0] is not None:
-        assert isinstance(model_state[0], str)
-        model, tokenizer, device, base_model = model_state0
+        model_state0 = model_state_none.copy()
+
+    # model_state['model] is only 'model' if should use model_state0
+    # model could also be None
+    have_model_lock = model_lock is not None
+    have_fresh_model = model_state['model'] not in [None, 'model', no_model_str]
+    # for gradio UI control, expect model_state and model_state0 to match, so if have_model_lock=True, then should have_fresh_model=True
+    # but gradio API control will only use nochat api etc. and won't use fresh model, so can't assert in general
+    # if have_model_lock:
+    #    assert have_fresh_model, "Expected model_state and model_state0 to match if have_model_lock"
+    have_cli_model = model_state0['model'] not in [None, 'model', no_model_str]
+
+    if have_fresh_model:
+        # USE FRESH MODEL
+        if not have_model_lock:
+            # model_state0 is just one of model_state if model_lock, so don't nuke
+            # try to free-up original model (i.e. list was passed as reference)
+            if model_state0['model'] and hasattr(model_state0['model'], 'cpu'):
+                model_state0['model'].cpu()
+                model_state0['model'] = None
+            # try to free-up original tokenizer (i.e. list was passed as reference)
+            if model_state0['tokenizer']:
+                model_state0['tokenizer'] = None
+            clear_torch_cache()
+        chosen_model_state = model_state
+    elif have_cli_model:
+        # USE MODEL SETUP AT CLI
+        assert isinstance(model_state['model'], str)  # expect no fresh model
+        chosen_model_state = model_state0
     else:
         raise AssertionError(no_model_msg)
+    # get variables
+    model = chosen_model_state['model']
+    tokenizer = chosen_model_state['tokenizer']
+    device = chosen_model_state['device']
+    base_model = chosen_model_state['base_model']
+    tokenizer_base_model = chosen_model_state['tokenizer_base_model']
+    lora_weights = chosen_model_state['lora_weights']
+    inference_server = chosen_model_state['inference_server']
+    # prefer use input from API over model state
+    prompt_type = prompt_type or chosen_model_state['prompt_type']
+    prompt_dict = prompt_dict or chosen_model_state['prompt_dict']
 
     if base_model is None:
         raise AssertionError(no_model_msg)
@@ -1061,10 +1459,48 @@ def evaluate(
         instruction = instruction_nochat
         iinput = iinput_nochat
 
+    # in some cases, like lean nochat API, don't want to force sending prompt_type, allow default choice
+    model_lower = base_model.lower()
+    if not prompt_type and model_lower in inv_prompt_type_to_model_lower:
+        prompt_type = inv_prompt_type_to_model_lower[model_lower]
+        if verbose:
+            print("Auto-selecting prompt_type=%s for %s" % (prompt_type, model_lower), flush=True)
+    assert prompt_type is not None, "prompt_type was None"
+
+    # Control generation hyperparameters
+    # adjust for bad inputs, e.g. in case also come from API that doesn't get constrained by gradio sliders
+    # below is for TGI server, not required for HF transformers
+    # limits are chosen similar to gradio_runner.py sliders/numbers
+    top_p = min(max(1e-3, top_p), 1.0 - 1e-3)
+    top_k = min(max(1, int(top_k)), 100)
+    temperature = min(max(0.01, temperature), 2.0)
+    # FIXME: https://github.com/h2oai/h2ogpt/issues/106
+    num_beams = 1 if stream_output else num_beams  # See max_beams in gradio_runner
+    max_max_new_tokens = get_max_max_new_tokens(chosen_model_state,
+                                                memory_restriction_level=memory_restriction_level,
+                                                max_new_tokens=max_new_tokens,
+                                                max_max_new_tokens=max_max_new_tokens)
+    model_max_length = get_model_max_length(chosen_model_state)
+    max_new_tokens = min(max(1, int(max_new_tokens)), max_max_new_tokens)
+    min_new_tokens = min(max(0, int(min_new_tokens)), max_new_tokens)
+    max_time = min(max(0, max_time), max_max_time)
+    repetition_penalty = min(max(0.01, repetition_penalty), 3.0)
+    num_return_sequences = 1 if chat else min(max(1, int(num_return_sequences)), 10)
+    min_top_k_docs, max_top_k_docs, label_top_k_docs = get_minmax_top_k_docs(is_public)
+    top_k_docs = min(max(min_top_k_docs, int(top_k_docs)), max_top_k_docs)
+    chunk_size = min(max(128, int(chunk_size)), 2048)
     if not context:
         # get hidden context if have one
         context = get_context(chat_context, prompt_type)
 
+    # restrict instruction, typically what has large input
+    from h2oai_pipeline import H2OTextGenerationPipeline
+    instruction, num_prompt_tokens1 = H2OTextGenerationPipeline.limit_prompt(instruction, tokenizer)
+    context, num_prompt_tokens2 = H2OTextGenerationPipeline.limit_prompt(context, tokenizer)
+    iinput, num_prompt_tokens3 = H2OTextGenerationPipeline.limit_prompt(iinput, tokenizer)
+    num_prompt_tokens = (num_prompt_tokens1 or 0) + (num_prompt_tokens2 or 0) + (num_prompt_tokens3 or 0)
+
+    # get prompt
     prompter = Prompter(prompt_type, prompt_dict, debug=debug, chat=chat, stream_output=stream_output)
     data_point = dict(context=context, instruction=instruction, input=iinput)
     prompt = prompter.generate_prompt(data_point)
@@ -1077,13 +1513,30 @@ def evaluate(
         db1 = dbs[langchain_mode]
     else:
         db1 = None
-    if langchain_mode not in [False, 'Disabled', 'ChatLLM', 'LLM'] and db1 is not None or base_model in non_hf_types:
+    do_langchain_path = langchain_mode not in [False, 'Disabled', 'ChatLLM', 'LLM'] and \
+                        db1 is not None or \
+                        base_model in non_hf_types or \
+                        force_langchain_evaluate
+    if do_langchain_path:
         query = instruction if not iinput else "%s\n%s" % (instruction, iinput)
         outr = ""
         # use smaller cut_distanct for wiki_full since so many matches could be obtained, and often irrelevant unless close
         from gpt_langchain import run_qa_db
+        gen_hyper_langchain = dict(do_sample=do_sample,
+                                   temperature=temperature,
+                                   repetition_penalty=repetition_penalty,
+                                   top_k=top_k,
+                                   top_p=top_p,
+                                   num_beams=num_beams,
+                                   min_new_tokens=min_new_tokens,
+                                   max_new_tokens=max_new_tokens,
+                                   early_stopping=early_stopping,
+                                   max_time=max_time,
+                                   num_return_sequences=num_return_sequences,
+                                   )
         for r in run_qa_db(query=query,
                            model_name=base_model, model=model, tokenizer=tokenizer,
+                           inference_server=inference_server,
                            stream_output=stream_output,
                            prompter=prompter,
                            load_db_if_exists=load_db_if_exists,
@@ -1103,29 +1556,31 @@ def evaluate(
                            db_type=db_type,
                            top_k_docs=top_k_docs,
 
-                           # gen_hyper:
-                           do_sample=do_sample,
-                           temperature=temperature,
-                           repetition_penalty=repetition_penalty,
-                           top_k=top_k,
-                           top_p=top_p,
-                           num_beams=num_beams,
-                           min_new_tokens=min_new_tokens,
-                           max_new_tokens=max_new_tokens,
-                           early_stopping=early_stopping,
-                           max_time=max_time,
-                           num_return_sequences=num_return_sequences,
+                           **gen_hyper_langchain,
 
                            prompt_type=prompt_type,
                            prompt_dict=prompt_dict,
                            n_jobs=n_jobs,
                            verbose=verbose,
                            cli=cli,
+                           sanitize_bot_response=sanitize_bot_response,
+                           reverse_docs=reverse_docs,
+
+                           lora_weights=lora_weights,
+
+                           auto_reduce_chunks=auto_reduce_chunks,
+                           max_chunks=max_chunks,
                            ):
             outr, extra = r  # doesn't accumulate, new answer every yield, so only save that full answer
             yield dict(response=outr, sources=extra)
         if save_dir:
-            save_generate_output(output=outr, base_model=base_model, save_dir=save_dir)
+            extra_dict = gen_hyper_langchain.copy()
+            extra_dict.update(prompt_type=prompt_type, inference_server=inference_server,
+                              langchain_mode=langchain_mode, document_choice=document_choice,
+                              num_prompt_tokens=num_prompt_tokens)
+            save_generate_output(prompt=query, output=outr, base_model=base_model, save_dir=save_dir,
+                                 where_from='run_qa_db',
+                                 extra_dict=extra_dict)
             if verbose:
                 print(
                     'Post-Generate Langchain: %s decoded_output: %s' % (str(datetime.now()), len(outr) if outr else -1),
@@ -1138,6 +1593,266 @@ def evaluate(
             clear_torch_cache()
             return
 
+    if inference_server.startswith('openai') or inference_server.startswith('http'):
+        if inference_server.startswith('openai'):
+            import openai
+            where_from = "openai_client"
+
+            openai.api_key = os.getenv("OPENAI_API_KEY")
+            stop_sequences = list(set(prompter.terminate_response + [prompter.PreResponse]))
+            # OpenAI will complain if ask for too many new tokens, takes it as min in some sense, wrongly so.
+            max_new_tokens_openai = min(max_new_tokens, model_max_length - num_prompt_tokens)
+            gen_server_kwargs = dict(temperature=temperature if do_sample else 0,
+                                     max_tokens=max_new_tokens_openai,
+                                     top_p=top_p if do_sample else 1,
+                                     frequency_penalty=0,
+                                     n=num_return_sequences,
+                                     presence_penalty=1.07 - repetition_penalty + 0.6,  # so good default
+                                     )
+            if inference_server == 'openai':
+                response = openai.Completion.create(
+                    model=base_model,
+                    prompt=prompt,
+                    **gen_server_kwargs,
+                    stop=stop_sequences,
+                    stream=stream_output,
+                )
+                if not stream_output:
+                    text = response['choices'][0]['text']
+                    yield dict(response=prompter.get_response(prompt + text, prompt=prompt,
+                                                              sanitize_bot_response=sanitize_bot_response),
+                               sources='')
+                else:
+                    collected_events = []
+                    text = ''
+                    for event in response:
+                        collected_events.append(event)  # save the event response
+                        event_text = event['choices'][0]['text']  # extract the text
+                        text += event_text  # append the text
+                        yield dict(response=prompter.get_response(prompt + text, prompt=prompt,
+                                                                  sanitize_bot_response=sanitize_bot_response),
+                                   sources='')
+            elif inference_server == 'openai_chat':
+                response = openai.ChatCompletion.create(
+                    model=base_model,
+                    messages=[
+                        {"role": "system", "content": "You are a helpful assistant."},
+                        {'role': 'user',
+                         'content': prompt,
+                         }
+                    ],
+                    stream=stream_output,
+                    **gen_server_kwargs,
+                )
+                if not stream_output:
+                    text = response["choices"][0]["message"]["content"]
+                    yield dict(response=prompter.get_response(prompt + text, prompt=prompt,
+                                                              sanitize_bot_response=sanitize_bot_response),
+                               sources='')
+                else:
+                    text = ""
+                    for chunk in response:
+                        delta = chunk["choices"][0]["delta"]
+                        if 'content' in delta:
+                            text += delta['content']
+                            yield dict(response=prompter.get_response(prompt + text, prompt=prompt,
+                                                                      sanitize_bot_response=sanitize_bot_response),
+                                       sources='')
+            else:
+                raise RuntimeError("No such OpenAI mode: %s" % inference_server)
+        elif inference_server.startswith('http'):
+            inference_server, headers = get_hf_server(inference_server)
+            from gradio_utils.grclient import GradioClient
+            from text_generation import Client as HFClient
+            if isinstance(model, GradioClient):
+                gr_client = model
+                hf_client = None
+            elif isinstance(model, HFClient):
+                gr_client = None
+                hf_client = model
+            else:
+                inference_server, gr_client, hf_client = get_client_from_inference_server(inference_server)
+
+            # quick sanity check to avoid long timeouts, just see if can reach server
+            requests.get(inference_server, timeout=int(os.getenv('REQUEST_TIMEOUT_FAST', '10')))
+
+            if gr_client is not None:
+                # Note: h2oGPT gradio server could handle input token size issues for prompt,
+                # but best to handle here so send less data to server
+
+                chat_client = False
+                where_from = "gr_client"
+                client_langchain_mode = 'Disabled'
+                gen_server_kwargs = dict(temperature=temperature,
+                                         top_p=top_p,
+                                         top_k=top_k,
+                                         num_beams=num_beams,
+                                         max_new_tokens=max_new_tokens,
+                                         min_new_tokens=min_new_tokens,
+                                         early_stopping=early_stopping,
+                                         max_time=max_time,
+                                         repetition_penalty=repetition_penalty,
+                                         num_return_sequences=num_return_sequences,
+                                         do_sample=do_sample,
+                                         chat=chat_client,
+                                         )
+                # account for gradio into gradio that handles prompting, avoid duplicating prompter prompt injection
+                if prompt_type in [None, '', PromptType.plain.name, PromptType.plain.value,
+                                   str(PromptType.plain.value)]:
+                    # if our prompt is plain, assume either correct or gradio server knows different prompt type,
+                    # so pass empty prompt_Type
+                    gr_prompt_type = ''
+                    gr_prompt_dict = ''
+                    gr_prompt = prompt  # already prepared prompt
+                    gr_context = ''
+                    gr_iinput = ''
+                else:
+                    # if already have prompt_type that is not plain, None, or '', then already applied some prompting
+                    #  But assume server can handle prompting, and need to avoid double-up.
+                    #  Also assume server can do better job of using stopping.py to stop early, so avoid local prompting, let server handle
+                    #  So avoid "prompt" and let gradio server reconstruct from prompt_type we passed
+                    # Note it's ok that prompter.get_response() has prompt+text, prompt=prompt passed,
+                    #  because just means extra processing and removal of prompt, but that has no human-bot prompting doesn't matter
+                    #  since those won't appear
+                    gr_context = context
+                    gr_prompt = instruction
+                    gr_iinput = iinput
+                    gr_prompt_type = prompt_type
+                    gr_prompt_dict = prompt_dict
+                client_kwargs = dict(instruction=gr_prompt if chat_client else '',  # only for chat=True
+                                     iinput=gr_iinput,  # only for chat=True
+                                     context=gr_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,
+
+                                     **gen_server_kwargs,
+
+                                     prompt_type=gr_prompt_type,
+                                     prompt_dict=gr_prompt_dict,
+
+                                     instruction_nochat=gr_prompt if not chat_client else '',
+                                     iinput_nochat=gr_iinput,  # only for chat=False
+                                     langchain_mode=client_langchain_mode,
+                                     top_k_docs=top_k_docs,
+                                     chunk=chunk,
+                                     chunk_size=chunk_size,
+                                     document_choice=[DocumentChoices.All_Relevant.name],
+                                     )
+                api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
+                if not stream_output:
+                    res = gr_client.predict(str(dict(client_kwargs)), api_name=api_name)
+                    res_dict = ast.literal_eval(res)
+                    text = res_dict['response']
+                    sources = res_dict['sources']
+                    yield dict(response=prompter.get_response(prompt + text, prompt=prompt,
+                                                              sanitize_bot_response=sanitize_bot_response),
+                               sources=sources)
+                else:
+                    job = gr_client.submit(str(dict(client_kwargs)), api_name=api_name)
+                    text = ''
+                    sources = ''
+                    res_dict = dict(response=text, sources=sources)
+                    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)
+                            text = res_dict['response']
+                            sources = res_dict['sources']
+                            if gr_prompt_type == 'plain':
+                                # then gradio server passes back full prompt + text
+                                prompt_and_text = text
+                            else:
+                                prompt_and_text = prompt + text
+                            yield dict(response=prompter.get_response(prompt_and_text, prompt=prompt,
+                                                                      sanitize_bot_response=sanitize_bot_response),
+                                       sources=sources)
+                        time.sleep(0.01)
+                    # ensure get last output to avoid race
+                    res_all = job.outputs()
+                    if len(res_all) > 0:
+                        res = res_all[-1]
+                        res_dict = ast.literal_eval(res)
+                        text = res_dict['response']
+                        sources = res_dict['sources']
+                    else:
+                        # go with old text if last call didn't work
+                        e = job.future._exception
+                        if e is not None:
+                            stre = str(e)
+                            strex = ''.join(traceback.format_tb(e.__traceback__))
+                        else:
+                            stre = ''
+                            strex = ''
+
+                        print("Bad final response: %s %s %s %s %s: %s %s" % (base_model, inference_server,
+                                                                             res_all, prompt, text, stre, strex),
+                              flush=True)
+                    if gr_prompt_type == 'plain':
+                        # then gradio server passes back full prompt + text
+                        prompt_and_text = text
+                    else:
+                        prompt_and_text = prompt + text
+                    yield dict(response=prompter.get_response(prompt_and_text, prompt=prompt,
+                                                              sanitize_bot_response=sanitize_bot_response),
+                               sources=sources)
+            elif hf_client:
+                # HF inference server needs control over input tokens
+                where_from = "hf_client"
+
+                # prompt must include all human-bot like tokens, already added by prompt
+                # https://github.com/huggingface/text-generation-inference/tree/main/clients/python#types
+                stop_sequences = list(set(prompter.terminate_response + [prompter.PreResponse]))
+                gen_server_kwargs = dict(do_sample=do_sample,
+                                         max_new_tokens=max_new_tokens,
+                                         # best_of=None,
+                                         repetition_penalty=repetition_penalty,
+                                         return_full_text=True,
+                                         seed=SEED,
+                                         stop_sequences=stop_sequences,
+                                         temperature=temperature,
+                                         top_k=top_k,
+                                         top_p=top_p,
+                                         # truncate=False,  # behaves oddly
+                                         # typical_p=top_p,
+                                         # watermark=False,
+                                         # decoder_input_details=False,
+                                         )
+                # work-around for timeout at constructor time, will be issue if multi-threading,
+                # so just do something reasonable or max_time if larger
+                # lower bound because client is re-used if multi-threading
+                hf_client.timeout = max(300, max_time)
+                if not stream_output:
+                    text = hf_client.generate(prompt, **gen_server_kwargs).generated_text
+                    yield dict(response=prompter.get_response(text, prompt=prompt,
+                                                              sanitize_bot_response=sanitize_bot_response),
+                               sources='')
+                else:
+                    text = ""
+                    for response in hf_client.generate_stream(prompt, **gen_server_kwargs):
+                        if not response.token.special:
+                            # stop_sequences
+                            text_chunk = response.token.text
+                            text += text_chunk
+                            yield dict(response=prompter.get_response(prompt + text, prompt=prompt,
+                                                                      sanitize_bot_response=sanitize_bot_response),
+                                       sources='')
+            else:
+                raise RuntimeError("Failed to get client: %s" % inference_server)
+        else:
+            raise RuntimeError("No such inference_server  %s" % inference_server)
+
+        if save_dir and text:
+            # save prompt + new text
+            extra_dict = gen_server_kwargs.copy()
+            extra_dict.update(dict(inference_server=inference_server, num_prompt_tokens=num_prompt_tokens))
+            save_generate_output(prompt=prompt, output=text, base_model=base_model, save_dir=save_dir,
+                                 where_from=where_from, extra_dict=extra_dict)
+        return
+    else:
+        assert not inference_server, "inferene_server=%s not supported" % inference_server
+
     if isinstance(tokenizer, str):
         # pipeline
         if tokenizer == "summarization":
@@ -1151,37 +1866,37 @@ def evaluate(
         assert src_lang is not None
         tokenizer.src_lang = languages_covered()[src_lang]
 
-    if chat:
-        # override, ignore user change
-        num_return_sequences = 1
-    stopping_criteria = get_stopping(prompt_type, prompt_dict, tokenizer, device)
-    _, _, max_length_tokenize, max_prompt_length = get_cutoffs(memory_restriction_level,
-                                                               model_max_length=tokenizer.model_max_length)
-    prompt = prompt[-max_prompt_length:]
-    inputs = tokenizer(prompt,
-                       return_tensors="pt",
-                       truncation=True,
-                       max_length=max_length_tokenize)
-    if inputs['input_ids'].shape[1] >= max_length_tokenize - 1:
-        print("Cutting off input: %s %s" % (inputs['input_ids'].shape[1], max_length_tokenize), flush=True)
+    stopping_criteria = get_stopping(prompt_type, prompt_dict, tokenizer, device,
+                                     model_max_length=tokenizer.model_max_length)
+
+    inputs = tokenizer(prompt, return_tensors="pt")
     if debug and len(inputs["input_ids"]) > 0:
         print('input_ids length', len(inputs["input_ids"][0]), flush=True)
     input_ids = inputs["input_ids"].to(device)
     # CRITICAL LIMIT else will fail
     max_max_tokens = tokenizer.model_max_length
-    max_input_tokens = max_max_tokens - max_new_tokens
+    max_input_tokens = max_max_tokens - min_new_tokens
+    # NOTE: Don't limit up front due to max_new_tokens, let go up to max or reach max_max_tokens in stopping.py
     input_ids = input_ids[:, -max_input_tokens:]
-    generation_config = GenerationConfig(
-        temperature=float(temperature),
-        top_p=float(top_p),
-        top_k=top_k,
-        num_beams=num_beams,
-        do_sample=do_sample,
-        repetition_penalty=float(repetition_penalty),
-        num_return_sequences=num_return_sequences,
-        renormalize_logits=True,
-        remove_invalid_values=True,
-    )
+    # required for falcon if multiple threads or asyncio accesses to model during generation
+    if use_cache is None:
+        use_cache = False if 'falcon' in base_model else True
+    gen_config_kwargs = dict(temperature=float(temperature),
+                             top_p=float(top_p),
+                             top_k=top_k,
+                             num_beams=num_beams,
+                             do_sample=do_sample,
+                             repetition_penalty=float(repetition_penalty),
+                             num_return_sequences=num_return_sequences,
+                             renormalize_logits=True,
+                             remove_invalid_values=True,
+                             use_cache=use_cache,
+                             )
+    token_ids = ['eos_token_id', 'pad_token_id', 'bos_token_id', 'cls_token_id', 'sep_token_id']
+    for token_id in token_ids:
+        if hasattr(tokenizer, token_id) and getattr(tokenizer, token_id) is not None:
+            gen_config_kwargs.update({token_id: getattr(tokenizer, token_id)})
+    generation_config = GenerationConfig(**gen_config_kwargs)
 
     gen_kwargs = dict(input_ids=input_ids,
                       generation_config=generation_config,
@@ -1200,7 +1915,10 @@ def evaluate(
         tgt_lang = languages_covered()[tgt_lang]
         gen_kwargs.update(dict(forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang]))
     else:
-        gen_kwargs.update(dict(pad_token_id=tokenizer.eos_token_id))
+        token_ids = ['eos_token_id', 'bos_token_id', 'pad_token_id']
+        for token_id in token_ids:
+            if hasattr(tokenizer, token_id) and getattr(tokenizer, token_id) is not None:
+                gen_kwargs.update({token_id: getattr(tokenizer, token_id)})
 
     decoder_kwargs = dict(skip_special_tokens=True,
                           clean_up_tokenization_spaces=True)
@@ -1216,7 +1934,8 @@ def evaluate(
                                     )
 
     with torch.no_grad():
-        context_class_cast = NullContext if device == 'cpu' or lora_weights else torch.autocast
+        have_lora_weights = lora_weights not in [no_lora_str, '', None]
+        context_class_cast = NullContext if device == 'cpu' or have_lora_weights else torch.autocast
         with context_class_cast(device):
             # protection for gradio not keeping track of closed users,
             # else hit bitsandbytes lack of thread safety:
@@ -1299,7 +2018,11 @@ def evaluate(
                     if outputs and len(outputs) >= 1:
                         decoded_output = prompt + outputs[0]
                 if save_dir and decoded_output:
-                    save_generate_output(output=decoded_output, base_model=base_model, save_dir=save_dir)
+                    extra_dict = gen_config_kwargs.copy()
+                    extra_dict.update(dict(num_prompt_tokens=num_prompt_tokens))
+                    save_generate_output(prompt=prompt, output=decoded_output, base_model=base_model, save_dir=save_dir,
+                                         where_from="evaluate_%s" % str(stream_output),
+                                         extra_dict=gen_config_kwargs)
             if verbose:
                 print('Post-Generate: %s decoded_output: %s' % (
                     str(datetime.now()), len(decoded_output) if decoded_output else -1), flush=True)
@@ -1318,6 +2041,7 @@ def get_cutoffs(memory_restriction_level, for_context=False, model_max_length=20
     if memory_restriction_level > 0:
         max_length_tokenize = 768 - 256 if memory_restriction_level <= 2 else 512 - 256
     else:
+        # at least give room for 1 paragraph output
         max_length_tokenize = model_max_length - 256
     cutoff_len = max_length_tokenize * 4  # if reaches limit, then can't generate new tokens
     output_smallest = 30 * 4
@@ -1422,7 +2146,8 @@ def get_generate_params(model_lower, chat,
     if model_lower:
         print(f"Using Model {model_lower}", flush=True)
     else:
-        print("No model defined yet", flush=True)
+        if verbose:
+            print("No model defined yet", flush=True)
 
     min_new_tokens = min_new_tokens if min_new_tokens is not None else 0
     early_stopping = early_stopping if early_stopping is not None else False
@@ -1478,12 +2203,14 @@ Philipp: ok, ok you can find everything here. https://huggingface.co/blog/the-pa
         use_defaults = True
     else:
         if chat:
-            placeholder_instruction = "Enter a question or imperative."
+            placeholder_instruction = ""
         else:
             placeholder_instruction = "Give detailed answer for whether Einstein or Newton is smarter."
         placeholder_input = ""
-        if model_lower:
-            # default is plain, because might relly upon trust_remote_code to handle prompting
+        if model_lower in inv_prompt_type_to_model_lower:
+            prompt_type = inv_prompt_type_to_model_lower[model_lower]
+        elif model_lower:
+            # default is plain, because might rely upon trust_remote_code to handle prompting
             prompt_type = prompt_type or 'plain'
         else:
             prompt_type = ''
@@ -1591,7 +2318,7 @@ y = np.random.randint(0, 1, 100)
 
     # get prompt_dict from prompt_type, so user can see in UI etc., or for custom do nothing except check format
     prompt_dict, error0 = get_prompt(prompt_type, prompt_dict,
-                                     chat=False, context='', reduced=False, return_dict=True)
+                                     chat=False, context='', reduced=False, making_context=False, return_dict=True)
     if error0:
         raise RuntimeError("Prompt wrong: %s" % error0)
 
@@ -1644,7 +2371,8 @@ def score_qa(smodel, stokenizer, max_length_tokenize, question, answer, cutoff_l
         if 'Expected all tensors to be on the same device' in str(e) or \
                 'expected scalar type Half but found Float' in str(e) or \
                 'probability tensor contains either' in str(e) or \
-                'cublasLt ran into an error!' in str(e):
+                'cublasLt ran into an error!' in str(e) or \
+                'device-side assert triggered' in str(e):
             print("GPU Error: question: %s answer: %s exception: %s" % (question, answer, str(e)),
                   flush=True)
             traceback.print_exc()
@@ -1677,25 +2405,101 @@ def check_locals(**kwargs):
         assert k in kwargs, "Missing %s" % k
 
 
+def get_model_max_length(model_state):
+    if not isinstance(model_state['tokenizer'], (str, types.NoneType)):
+        return model_state['tokenizer'].model_max_length
+    else:
+        return 2048
+
+
 def get_max_max_new_tokens(model_state, **kwargs):
-    if kwargs['max_new_tokens'] and kwargs['user_set_max_new_tokens']:
-        max_max_new_tokens = kwargs['max_new_tokens']
+    if not isinstance(model_state['tokenizer'], (str, types.NoneType)):
+        max_max_new_tokens = model_state['tokenizer'].model_max_length
+    else:
+        max_max_new_tokens = None
+
+    if kwargs['max_max_new_tokens'] is not None and max_max_new_tokens is not None:
+        return min(max_max_new_tokens, kwargs['max_max_new_tokens'])
+    elif kwargs['max_max_new_tokens'] is not None:
+        return kwargs['max_max_new_tokens']
     elif kwargs['memory_restriction_level'] == 1:
-        max_max_new_tokens = 768
+        return 768
     elif kwargs['memory_restriction_level'] == 2:
-        max_max_new_tokens = 512
+        return 512
     elif kwargs['memory_restriction_level'] >= 3:
-        max_max_new_tokens = 256
+        return 256
     else:
-        if not isinstance(model_state[1], str):
-            max_max_new_tokens = model_state[1].model_max_length
-        else:
-            # FIXME: Need to update after new model loaded, so user can control with slider
-            max_max_new_tokens = 2048
-    return max_max_new_tokens
+        # FIXME: Need to update after new model loaded, so user can control with slider
+        return 2048
 
 
-if __name__ == "__main__":
+def get_minmax_top_k_docs(is_public):
+    if is_public:
+        min_top_k_docs = 1
+        max_top_k_docs = 3
+        label_top_k_docs = "Number of document chunks"
+    else:
+        min_top_k_docs = -1
+        max_top_k_docs = 100
+        label_top_k_docs = "Number of document chunks (-1 = auto fill model context)"
+    return min_top_k_docs, max_top_k_docs, label_top_k_docs
+
+
+def history_to_context(history, langchain_mode1, prompt_type1, prompt_dict1, chat1, model_max_length1,
+                       memory_restriction_level1, keep_sources_in_context1):
+    """
+    consumes all history up to (but not including) latest history item that is presumed to be an [instruction, None] pair
+    :param history:
+    :param langchain_mode1:
+    :param prompt_type1:
+    :param prompt_dict1:
+    :param chat1:
+    :param model_max_length1:
+    :param memory_restriction_level1:
+    :param keep_sources_in_context1:
+    :return:
+    """
+    # ensure output will be unique to models
+    _, _, _, max_prompt_length = get_cutoffs(memory_restriction_level1,
+                                             for_context=True, model_max_length=model_max_length1)
+    context1 = ''
+    if max_prompt_length is not None and langchain_mode1 not in ['LLM']:
+        context1 = ''
+        # - 1 below because current instruction already in history from user()
+        for histi in range(0, len(history) - 1):
+            data_point = dict(instruction=history[histi][0], input='', output=history[histi][1])
+            prompt, pre_response, terminate_response, chat_sep, chat_turn_sep = generate_prompt(data_point,
+                                                                                                prompt_type1,
+                                                                                                prompt_dict1,
+                                                                                                chat1,
+                                                                                                reduced=True,
+                                                                                                making_context=True)
+            # md -> back to text, maybe not super important if model trained enough
+            if not keep_sources_in_context1 and langchain_mode1 != 'Disabled' and prompt.find(source_prefix) >= 0:
+                # FIXME: This is relatively slow even for small amount of text, like 0.3s each history item
+                import re
+                prompt = re.sub(f'{re.escape(source_prefix)}.*?{re.escape(source_postfix)}', '', prompt,
+                                flags=re.DOTALL)
+                if prompt.endswith('\n<p>'):
+                    prompt = prompt[:-4]
+            prompt = prompt.replace('<br>', chat_turn_sep)
+            if not prompt.endswith(chat_turn_sep):
+                prompt += chat_turn_sep
+            # most recent first, add older if can
+            # only include desired chat history
+            if len(prompt + context1) > max_prompt_length:
+                break
+            context1 += prompt
+
+        _, pre_response, terminate_response, chat_sep, chat_turn_sep = generate_prompt({}, prompt_type1, prompt_dict1,
+                                                                                       chat1, reduced=True,
+                                                                                       making_context=True)
+        if context1 and not context1.endswith(chat_turn_sep):
+            context1 += chat_turn_sep  # ensure if terminates abruptly, then human continues on next line
+    return context1
+
+
+def entrypoint_main():
     """
     Examples:
 
@@ -1726,3 +2530,7 @@ if __name__ == "__main__":
     python generate.py --base_model=h2oai/h2ogpt-oig-oasst1-512-6_9b
     """
     fire.Fire(main)
+
+
+if __name__ == "__main__":
+    entrypoint_main()

gpt4all_llm.py

@@ -1,24 +1,13 @@
 import inspect
 import os
-import sys
+from functools import partial
 from typing import Dict, Any, Optional, List
 from langchain.callbacks.manager import CallbackManagerForLLMRun
 from pydantic import root_validator
 from langchain.llms import gpt4all
 from dotenv import dotenv_values
 
-
-class FakeTokenizer:
-    model_max_length = 2048
-
-    def encode(self, x, *args, **kwargs):
-        return dict(input_ids=[x])
-
-    def decode(self, x, *args, **kwargs):
-        return x
-
-    def __call__(self, x, *args, **kwargs):
-        return self.encode(x, *args, **kwargs)
+from utils import FakeTokenizer
 
 
 def get_model_tokenizer_gpt4all(base_model, **kwargs):
@@ -74,9 +63,9 @@ class H2OStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
         pass
 
 
-def get_model_kwargs(env_kwargs, default_kwargs, cls):
+def get_model_kwargs(env_kwargs, default_kwargs, cls, exclude_list=[]):
     # default from class
-    model_kwargs = {k: v.default for k, v in dict(inspect.signature(cls).parameters).items()}
+    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
@@ -94,10 +83,14 @@ def get_llm_gpt4all(model_name,
                     repetition_penalty=1.0,
                     top_k=40,
                     top_p=0.7,
-                    verbose=False):
+                    streaming=False,
+                    callbacks=None,
+                    prompter=None,
+                    verbose=False,
+                    ):
+    assert prompter is not None
     env_gpt4all_file = ".env_gpt4all"
     env_kwargs = dotenv_values(env_gpt4all_file)
-    callbacks = [H2OStreamingStdOutCallbackHandler()]
     n_ctx = env_kwargs.pop('n_ctx', 2048 - max_new_tokens)
     default_kwargs = dict(context_erase=0.5,
                           n_batch=1,
@@ -114,21 +107,23 @@ def get_llm_gpt4all(model_name,
     if model_name == 'llama':
         cls = H2OLlamaCpp
         model_path = env_kwargs.pop('model_path_llama') if model is None else model
-        model_kwargs = get_model_kwargs(env_kwargs, default_kwargs, cls)
-        model_kwargs.update(dict(model_path=model_path, callbacks=callbacks))
+        model_kwargs = get_model_kwargs(env_kwargs, default_kwargs, cls, exclude_list=['lc_kwargs'])
+        model_kwargs.update(dict(model_path=model_path, callbacks=callbacks, streaming=streaming, prompter=prompter))
         llm = cls(**model_kwargs)
         llm.client.verbose = verbose
     elif model_name == 'gpt4all_llama':
         cls = H2OGPT4All
         model_path = env_kwargs.pop('model_path_gpt4all_llama') if model is None else model
-        model_kwargs = get_model_kwargs(env_kwargs, default_kwargs, cls)
-        model_kwargs.update(dict(model=model_path, backend='llama', callbacks=callbacks))
+        model_kwargs = get_model_kwargs(env_kwargs, default_kwargs, cls, exclude_list=['lc_kwargs'])
+        model_kwargs.update(
+            dict(model=model_path, backend='llama', callbacks=callbacks, streaming=streaming, prompter=prompter))
         llm = cls(**model_kwargs)
     elif model_name == 'gptj':
         cls = H2OGPT4All
         model_path = env_kwargs.pop('model_path_gptj') if model is None else model
-        model_kwargs = get_model_kwargs(env_kwargs, default_kwargs, cls)
-        model_kwargs.update(dict(model=model_path, backend='gptj', callbacks=callbacks))
+        model_kwargs = get_model_kwargs(env_kwargs, default_kwargs, cls, exclude_list=['lc_kwargs'])
+        model_kwargs.update(
+            dict(model=model_path, backend='gptj', callbacks=callbacks, streaming=streaming, prompter=prompter))
         llm = cls(**model_kwargs)
     else:
         raise RuntimeError("No such model_name %s" % model_name)
@@ -137,6 +132,7 @@ def get_llm_gpt4all(model_name,
 
 class H2OGPT4All(gpt4all.GPT4All):
     model: Any
+    prompter: Any
     """Path to the pre-trained GPT4All model file."""
 
     @root_validator()
@@ -156,9 +152,16 @@ class H2OGPT4All(gpt4all.GPT4All):
                     model_type=values["backend"],
                     allow_download=False,
                 )
+                if values["n_threads"] is not None:
+                    # set n_threads
+                    values["client"].model.set_thread_count(values["n_threads"])
             else:
                 values["client"] = values["model"]
-            values["backend"] = values["client"].model.model_type
+            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(
@@ -172,12 +175,19 @@ class H2OGPT4All(gpt4all.GPT4All):
             prompt: str,
             stop: Optional[List[str]] = None,
             run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs,
     ) -> str:
         # Roughly 4 chars per token if natural language
         prompt = prompt[-self.n_ctx * 4:]
+
+        # use instruct prompting
+        data_point = dict(context='', instruction=prompt, input='')
+        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)
 
 
@@ -186,6 +196,7 @@ from langchain.llms import LlamaCpp
 
 class H2OLlamaCpp(LlamaCpp):
     model_path: Any
+    prompter: Any
     """Path to the pre-trained GPT4All model file."""
 
     @root_validator()
@@ -237,6 +248,7 @@ class H2OLlamaCpp(LlamaCpp):
             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
@@ -253,6 +265,33 @@ class H2OLlamaCpp(LlamaCpp):
                 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='', instruction=prompt, input='')
+        prompt = self.prompter.generate_prompt(data_point)
+
         if verbose:
             print("_call prompt: %s" % prompt, flush=True)
-        return super()._call(prompt, stop=stop, run_manager=run_manager)
+
+        if self.streaming:
+            text_callback = None
+            if run_manager:
+                text_callback = partial(
+                    run_manager.on_llm_new_token, verbose=self.verbose
+                )
+            # parent handler of streamer expects to see prompt first else output="" and lose if prompt=None in prompter
+            if text_callback:
+                text_callback(prompt)
+            text = ""
+            for token in self.stream(prompt=prompt, stop=stop, run_manager=run_manager):
+                text_chunk = token["choices"][0]["text"]
+                # self.stream already calls text_callback
+                # if text_callback:
+                #    text_callback(text_chunk)
+                text += text_chunk
+            return text
+        else:
+            params = self._get_parameters(stop)
+            params = {**params, **kwargs}
+            result = self.client(prompt=prompt, **params)
+            return result["choices"][0]["text"]

gpt_langchain.py

@@ -1,31 +1,34 @@
+import ast
 import glob
 import inspect
 import os
 import pathlib
 import pickle
-import queue
-import random
 import shutil
 import subprocess
-import sys
 import tempfile
+import time
 import traceback
+import types
 import uuid
 import zipfile
 from collections import defaultdict
 from datetime import datetime
 from functools import reduce
 from operator import concat
+import filelock
 
-from joblib import Parallel, delayed
+from joblib import delayed
+from langchain.callbacks import streaming_stdout
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 from tqdm import tqdm
 
-from enums import DocumentChoices
-from generate import gen_hyper
-from prompter import non_hf_types, PromptType
+from enums import DocumentChoices, no_lora_str, model_token_mapping, source_prefix, source_postfix
+from generate import gen_hyper, get_model, SEED
+from prompter import non_hf_types, PromptType, Prompter
 from utils import wrapped_partial, EThread, import_matplotlib, sanitize_filename, makedirs, get_url, flatten_list, \
-    get_device, ProgressParallel, remove, hash_file, clear_torch_cache
+    get_device, ProgressParallel, remove, hash_file, clear_torch_cache, NullContext, get_hf_server, FakeTokenizer
+from utils_langchain import StreamingGradioCallbackHandler
 
 import_matplotlib()
 
@@ -40,11 +43,11 @@ from langchain.chains.qa_with_sources import load_qa_with_sources_chain
 from langchain.document_loaders import PyPDFLoader, TextLoader, CSVLoader, PythonLoader, TomlLoader, \
     UnstructuredURLLoader, UnstructuredHTMLLoader, UnstructuredWordDocumentLoader, UnstructuredMarkdownLoader, \
     EverNoteLoader, UnstructuredEmailLoader, UnstructuredODTLoader, UnstructuredPowerPointLoader, \
-    UnstructuredEPubLoader, UnstructuredImageLoader, UnstructuredRTFLoader, ArxivLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+    UnstructuredEPubLoader, UnstructuredImageLoader, UnstructuredRTFLoader, ArxivLoader, UnstructuredPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter, Language
 from langchain.chains.question_answering import load_qa_chain
 from langchain.docstore.document import Document
-from langchain import PromptTemplate
+from langchain import PromptTemplate, HuggingFaceTextGenInference
 from langchain.vectorstores import Chroma
 
 
@@ -142,7 +145,7 @@ def add_to_db(db, sources, db_type='faiss',
             return db, num_new_sources, []
         db.add_documents(documents=sources)
     elif db_type == 'chroma':
-        collection = db.get()
+        collection = get_documents(db)
         # files we already have:
         metadata_files = set([x['source'] for x in collection['metadatas']])
         if avoid_dup_by_file:
@@ -157,6 +160,9 @@ def add_to_db(db, sources, db_type='faiss',
                 [x['hashid'] for x in collection['metadatas'] if 'hashid' in x and x['hashid'] not in ["None", None]])
             # avoid sources with same hash
             sources = [x for x in sources if x.metadata.get('hashid') not in metadata_hash_ids]
+            num_nohash = len([x for x in sources if not x.metadata.get('hashid')])
+            print("Found %s new sources (%d have no hash in original source,"
+                  " so have to reprocess for migration to sources with hash)" % (len(sources), num_nohash), flush=True)
             # get new file names that match existing file names.  delete existing files we are overridding
             dup_metadata_files = set([x.metadata['source'] for x in sources if x.metadata['source'] in metadata_files])
             print("Removing %s duplicate files from db because ingesting those as new documents" % len(
@@ -233,7 +239,7 @@ def get_embedding(use_openai_embedding, hf_embedding_model="sentence-transformer
     if use_openai_embedding:
         assert os.getenv("OPENAI_API_KEY") is not None, "Set ENV OPENAI_API_KEY"
         from langchain.embeddings import OpenAIEmbeddings
-        embedding = OpenAIEmbeddings()
+        embedding = OpenAIEmbeddings(disallowed_special=())
     else:
         # to ensure can fork without deadlock
         from langchain.embeddings import HuggingFaceEmbeddings
@@ -260,8 +266,315 @@ def get_answer_from_sources(chain, sources, question):
     )["output_text"]
 
 
-def get_llm(use_openai_model=False, model_name=None, model=None,
-            tokenizer=None, stream_output=False,
+"""Wrapper around Huggingface text generation inference API."""
+from functools import partial
+from typing import Any, Dict, List, Optional, Set
+
+from pydantic import Extra, Field, root_validator
+
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+
+"""Wrapper around Huggingface text generation inference API."""
+from functools import partial
+from typing import Any, Dict, List, Optional
+
+from pydantic import Extra, Field, root_validator
+
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.llms.base import LLM
+
+
+class GradioInference(LLM):
+    """
+    Gradio generation inference API.
+    """
+    inference_server_url: str = ""
+
+    temperature: float = 0.8
+    top_p: Optional[float] = 0.95
+    top_k: Optional[int] = None
+    num_beams: Optional[int] = 1
+    max_new_tokens: int = 512
+    min_new_tokens: int = 1
+    early_stopping: bool = False
+    max_time: int = 180
+    repetition_penalty: Optional[float] = None
+    num_return_sequences: Optional[int] = 1
+    do_sample: bool = False
+    chat_client: bool = False
+
+    return_full_text: bool = True
+    stream: bool = False
+    sanitize_bot_response: bool = False
+
+    prompter: Any = None
+    client: Any = None
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that python package exists in environment."""
+
+        try:
+            if values['client'] is None:
+                import gradio_client
+                values["client"] = gradio_client.Client(
+                    values["inference_server_url"]
+                )
+        except ImportError:
+            raise ImportError(
+                "Could not import gradio_client python package. "
+                "Please install it with `pip install gradio_client`."
+            )
+        return values
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "gradio_inference"
+
+    def _call(
+            self,
+            prompt: str,
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs: Any,
+    ) -> str:
+        # NOTE: prompt here has no prompt_type (e.g. human: bot:) prompt injection,
+        # so server should get prompt_type or '', not plain
+        # This is good, so gradio server can also handle stopping.py conditions
+        # this is different than TGI server that uses prompter to inject prompt_type prompting
+        stream_output = self.stream
+        gr_client = self.client
+        client_langchain_mode = 'Disabled'
+        top_k_docs = 1
+        chunk = True
+        chunk_size = 512
+        client_kwargs = dict(instruction=prompt if self.chat_client 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=self.prompter.prompt_type,
+                             prompt_dict='',
+
+                             temperature=self.temperature,
+                             top_p=self.top_p,
+                             top_k=self.top_k,
+                             num_beams=self.num_beams,
+                             max_new_tokens=self.max_new_tokens,
+                             min_new_tokens=self.min_new_tokens,
+                             early_stopping=self.early_stopping,
+                             max_time=self.max_time,
+                             repetition_penalty=self.repetition_penalty,
+                             num_return_sequences=self.num_return_sequences,
+                             do_sample=self.do_sample,
+                             chat=self.chat_client,
+
+                             instruction_nochat=prompt if not self.chat_client else '',
+                             iinput_nochat='',  # only for chat=False
+                             langchain_mode=client_langchain_mode,
+                             top_k_docs=top_k_docs,
+                             chunk=chunk,
+                             chunk_size=chunk_size,
+                             document_choice=[DocumentChoices.All_Relevant.name],
+                             )
+        api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
+        if not stream_output:
+            res = gr_client.predict(str(dict(client_kwargs)), api_name=api_name)
+            res_dict = ast.literal_eval(res)
+            text = res_dict['response']
+            return self.prompter.get_response(prompt + text, prompt=prompt,
+                                              sanitize_bot_response=self.sanitize_bot_response)
+        else:
+            text_callback = None
+            if run_manager:
+                text_callback = partial(
+                    run_manager.on_llm_new_token, verbose=self.verbose
+                )
+
+            job = gr_client.submit(str(dict(client_kwargs)), api_name=api_name)
+            text0 = ''
+            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)
+                    text = res_dict['response']
+                    text = self.prompter.get_response(prompt + text, prompt=prompt,
+                                                      sanitize_bot_response=self.sanitize_bot_response)
+                    # FIXME: derive chunk from full for now
+                    text_chunk = text[len(text0):]
+                    # save old
+                    text0 = text
+
+                    if text_callback:
+                        text_callback(text_chunk)
+
+                time.sleep(0.01)
+
+            # ensure get last output to avoid race
+            res_all = job.outputs()
+            if len(res_all) > 0:
+                res = res_all[-1]
+                res_dict = ast.literal_eval(res)
+                text = res_dict['response']
+                # FIXME: derive chunk from full for now
+            else:
+                # go with old if failure
+                text = text0
+            text_chunk = text[len(text0):]
+            if text_callback:
+                text_callback(text_chunk)
+            return self.prompter.get_response(prompt + text, prompt=prompt,
+                                              sanitize_bot_response=self.sanitize_bot_response)
+
+
+class H2OHuggingFaceTextGenInference(HuggingFaceTextGenInference):
+    max_new_tokens: int = 512
+    do_sample: bool = False
+    top_k: Optional[int] = None
+    top_p: Optional[float] = 0.95
+    typical_p: Optional[float] = 0.95
+    temperature: float = 0.8
+    repetition_penalty: Optional[float] = None
+    return_full_text: bool = False
+    stop_sequences: List[str] = Field(default_factory=list)
+    seed: Optional[int] = None
+    inference_server_url: str = ""
+    timeout: int = 300
+    headers: dict = None
+    stream: bool = False
+    sanitize_bot_response: bool = False
+    prompter: Any = None
+    tokenizer: Any = None
+    client: Any = None
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that python package exists in environment."""
+
+        try:
+            if values['client'] is None:
+                import text_generation
+
+                values["client"] = text_generation.Client(
+                    values["inference_server_url"],
+                    timeout=values["timeout"],
+                    headers=values["headers"],
+                )
+        except ImportError:
+            raise ImportError(
+                "Could not import text_generation python package. "
+                "Please install it with `pip install text_generation`."
+            )
+        return values
+
+    def _call(
+            self,
+            prompt: str,
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs: Any,
+    ) -> str:
+        if stop is None:
+            stop = self.stop_sequences
+        else:
+            stop += self.stop_sequences
+
+        # HF inference server needs control over input tokens
+        assert self.tokenizer is not None
+        from h2oai_pipeline import H2OTextGenerationPipeline
+        prompt, num_prompt_tokens = H2OTextGenerationPipeline.limit_prompt(prompt, self.tokenizer)
+
+        # NOTE: TGI server does not add prompting, so must do here
+        data_point = dict(context='', instruction=prompt, input='')
+        prompt = self.prompter.generate_prompt(data_point)
+
+        gen_server_kwargs = dict(do_sample=self.do_sample,
+                                 stop_sequences=stop,
+                                 max_new_tokens=self.max_new_tokens,
+                                 top_k=self.top_k,
+                                 top_p=self.top_p,
+                                 typical_p=self.typical_p,
+                                 temperature=self.temperature,
+                                 repetition_penalty=self.repetition_penalty,
+                                 return_full_text=self.return_full_text,
+                                 seed=self.seed,
+                                 )
+        gen_server_kwargs.update(kwargs)
+
+        # lower bound because client is re-used if multi-threading
+        self.client.timeout = max(300, self.timeout)
+
+        if not self.stream:
+            res = self.client.generate(
+                prompt,
+                **gen_server_kwargs,
+            )
+            if self.return_full_text:
+                gen_text = res.generated_text[len(prompt):]
+            else:
+                gen_text = res.generated_text
+            # remove stop sequences from the end of the generated text
+            for stop_seq in stop:
+                if stop_seq in gen_text:
+                    gen_text = gen_text[:gen_text.index(stop_seq)]
+            text = prompt + gen_text
+            text = self.prompter.get_response(text, prompt=prompt,
+                                              sanitize_bot_response=self.sanitize_bot_response)
+        else:
+            text_callback = None
+            if run_manager:
+                text_callback = partial(
+                    run_manager.on_llm_new_token, verbose=self.verbose
+                )
+            # parent handler of streamer expects to see prompt first else output="" and lose if prompt=None in prompter
+            if text_callback:
+                text_callback(prompt)
+            text = ""
+            # Note: Streaming ignores return_full_text=True
+            for response in self.client.generate_stream(prompt, **gen_server_kwargs):
+                text_chunk = response.token.text
+                text += text_chunk
+                text = self.prompter.get_response(prompt + text, prompt=prompt,
+                                                  sanitize_bot_response=self.sanitize_bot_response)
+                # stream part
+                is_stop = False
+                for stop_seq in stop:
+                    if stop_seq in response.token.text:
+                        is_stop = True
+                        break
+                if is_stop:
+                    break
+                if not response.token.special:
+                    if text_callback:
+                        text_callback(response.token.text)
+        return text
+
+
+from langchain.chat_models import ChatOpenAI
+
+
+class H2OChatOpenAI(ChatOpenAI):
+    @classmethod
+    def all_required_field_names(cls) -> Set:
+        all_required_field_names = super(ChatOpenAI, cls).all_required_field_names()
+        all_required_field_names.update({'top_p', 'frequency_penalty', 'presence_penalty'})
+        return all_required_field_names
+
+
+def get_llm(use_openai_model=False,
+            model_name=None,
+            model=None,
+            tokenizer=None,
+            inference_server=None,
+            stream_output=False,
             do_sample=False,
             temperature=0.1,
             top_k=40,
@@ -276,47 +589,142 @@ def get_llm(use_openai_model=False, model_name=None, model=None,
             prompt_type=None,
             prompt_dict=None,
             prompter=None,
+            sanitize_bot_response=False,
             verbose=False,
             ):
-    if use_openai_model:
-        from langchain.llms import OpenAI
-        llm = OpenAI(temperature=0)
-        model_name = 'openai'
-        streamer = None
-        prompt_type = 'plain'
+    if use_openai_model or inference_server in ['openai', 'openai_chat']:
+        if use_openai_model and model_name is None:
+            model_name = "gpt-3.5-turbo"
+        if inference_server == 'openai':
+            from langchain.llms import OpenAI
+            cls = OpenAI
+        else:
+            cls = H2OChatOpenAI
+        callbacks = [StreamingGradioCallbackHandler()]
+        llm = cls(model_name=model_name,
+                  temperature=temperature if do_sample else 0,
+                  # FIXME: Need to count tokens and reduce max_new_tokens to fit like in generate.py
+                  max_tokens=max_new_tokens,
+                  top_p=top_p if do_sample else 1,
+                  frequency_penalty=0,
+                  presence_penalty=1.07 - repetition_penalty + 0.6,  # so good default
+                  callbacks=callbacks if stream_output else None,
+                  )
+        streamer = callbacks[0] if stream_output else None
+        if inference_server in ['openai', 'openai_chat']:
+            prompt_type = inference_server
+        else:
+            prompt_type = prompt_type or 'plain'
+    elif inference_server:
+        assert inference_server.startswith(
+            'http'), "Malformed inference_server=%s.  Did you add http:// in front?" % inference_server
+
+        from gradio_utils.grclient import GradioClient
+        from text_generation import Client as HFClient
+        if isinstance(model, GradioClient):
+            gr_client = model
+            hf_client = None
+        else:
+            gr_client = None
+            hf_client = model
+            assert isinstance(hf_client, HFClient)
+
+        inference_server, headers = get_hf_server(inference_server)
+
+        # quick sanity check to avoid long timeouts, just see if can reach server
+        requests.get(inference_server, timeout=int(os.getenv('REQUEST_TIMEOUT_FAST', '10')))
+
+        callbacks = [StreamingGradioCallbackHandler()]
+        assert prompter is not None
+        stop_sequences = list(set(prompter.terminate_response + [prompter.PreResponse]))
+
+        if gr_client:
+            chat_client = False
+            llm = GradioInference(
+                inference_server_url=inference_server,
+                return_full_text=True,
+
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                num_beams=num_beams,
+                max_new_tokens=max_new_tokens,
+                min_new_tokens=min_new_tokens,
+                early_stopping=early_stopping,
+                max_time=max_time,
+                repetition_penalty=repetition_penalty,
+                num_return_sequences=num_return_sequences,
+                do_sample=do_sample,
+                chat_client=chat_client,
+
+                callbacks=callbacks if stream_output else None,
+                stream=stream_output,
+                prompter=prompter,
+                client=gr_client,
+                sanitize_bot_response=sanitize_bot_response,
+            )
+        elif hf_client:
+            llm = H2OHuggingFaceTextGenInference(
+                inference_server_url=inference_server,
+                do_sample=do_sample,
+                max_new_tokens=max_new_tokens,
+                repetition_penalty=repetition_penalty,
+                return_full_text=True,
+                seed=SEED,
+
+                stop_sequences=stop_sequences,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                # typical_p=top_p,
+                callbacks=callbacks if stream_output else None,
+                stream=stream_output,
+                prompter=prompter,
+                tokenizer=tokenizer,
+                client=hf_client,
+                timeout=max_time,
+                sanitize_bot_response=sanitize_bot_response,
+            )
+        else:
+            raise RuntimeError("No defined client")
+        streamer = callbacks[0] if stream_output else None
     elif model_name in non_hf_types:
+        if model_name == 'llama':
+            callbacks = [StreamingGradioCallbackHandler()]
+            streamer = callbacks[0] if stream_output else None
+        else:
+            # stream_output = False
+            # doesn't stream properly as generator, but at least
+            callbacks = [streaming_stdout.StreamingStdOutCallbackHandler()]
+            streamer = None
+        if prompter:
+            prompt_type = prompter.prompt_type
+        else:
+            prompter = Prompter(prompt_type, prompt_dict, debug=False, chat=False, stream_output=stream_output)
+            pass  # assume inputted prompt_type is correct
         from gpt4all_llm import get_llm_gpt4all
         llm = get_llm_gpt4all(model_name, model=model, max_new_tokens=max_new_tokens,
                               temperature=temperature,
                               repetition_penalty=repetition_penalty,
                               top_k=top_k,
                               top_p=top_p,
+                              callbacks=callbacks,
                               verbose=verbose,
+                              streaming=stream_output,
+                              prompter=prompter,
                               )
-        streamer = None
-        prompt_type = 'plain'
     else:
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-
         if model is None:
             # only used if didn't pass model in
             assert tokenizer is None
             prompt_type = 'human_bot'
-            model_name = 'h2oai/h2ogpt-oasst1-512-12b'
-            # model_name = 'h2oai/h2ogpt-oig-oasst1-512-6_9b'
-            # model_name = 'h2oai/h2ogpt-oasst1-512-20b'
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            device, torch_dtype, context_class = get_device_dtype()
-
-            with context_class(device):
-                load_8bit = True
-                # FIXME: for now not to spread across hetero GPUs
-                # device_map={"": 0} if load_8bit and device == 'cuda' else "auto"
-                device_map = {"": 0} if device == 'cuda' else "auto"
-                model = AutoModelForCausalLM.from_pretrained(model_name,
-                                                             device_map=device_map,
-                                                             torch_dtype=torch_dtype,
-                                                             load_in_8bit=load_8bit)
+            if model_name is None:
+                model_name = 'h2oai/h2ogpt-oasst1-512-12b'
+                # model_name = 'h2oai/h2ogpt-oig-oasst1-512-6_9b'
+                # model_name = 'h2oai/h2ogpt-oasst1-512-20b'
+            inference_server = ''
+            model, tokenizer, device = get_model(load_8bit=True, base_model=model_name,
+                                                 inference_server=inference_server, gpu_id=0)
 
         max_max_tokens = tokenizer.model_max_length
         gen_kwargs = dict(do_sample=do_sample,
@@ -331,7 +739,7 @@ def get_llm(use_openai_model=False, model_name=None, model=None,
                           repetition_penalty=repetition_penalty,
                           num_return_sequences=num_return_sequences,
                           return_full_text=True,
-                          handle_long_generation='hole')
+                          handle_long_generation=None)
         assert len(set(gen_hyper).difference(gen_kwargs.keys())) == 0
 
         if stream_output:
@@ -348,10 +756,11 @@ def get_llm(use_openai_model=False, model_name=None, model=None,
                                          prompter=prompter,
                                          prompt_type=prompt_type,
                                          prompt_dict=prompt_dict,
-                                         sanitize_bot_response=True,
+                                         sanitize_bot_response=sanitize_bot_response,
                                          chat=False, stream_output=stream_output,
                                          tokenizer=tokenizer,
-                                         max_input_tokens=max_max_tokens - max_new_tokens,
+                                         # leave some room for 1 paragraph, even if min_new_tokens=0
+                                         max_input_tokens=max_max_tokens - max(min_new_tokens, 256),
                                          **gen_kwargs)
         # pipe.task = "text-generation"
         # below makes it listen only to our prompt removal,
@@ -396,7 +805,7 @@ def get_wiki_data(title, first_paragraph_only, text_limit=None, take_head=True):
         data = json.load(open(filename, "rt"))
     page_content = list(data["query"]["pages"].values())[0]["extract"]
     if take_head is not None and text_limit is not None:
-        page_content = page_content[:text_limit] if take_head else page_content[:-text_limit]
+        page_content = page_content[:text_limit] if take_head else page_content[-text_limit:]
     title_url = str(title).replace(' ', '_')
     return Document(
         page_content=page_content,
@@ -518,6 +927,21 @@ try:
 except (pkg_resources.DistributionNotFound, AssertionError):
     have_pymupdf = False
 
+try:
+    assert pkg_resources.get_distribution('selenium') is not None
+    have_selenium = True
+except (pkg_resources.DistributionNotFound, AssertionError):
+    have_selenium = False
+
+try:
+    assert pkg_resources.get_distribution('playwright') is not None
+    have_playwright = True
+except (pkg_resources.DistributionNotFound, AssertionError):
+    have_playwright = False
+
+# disable, hangs too often
+have_playwright = False
+
 image_types = ["png", "jpg", "jpeg"]
 non_image_types = ["pdf", "txt", "csv", "toml", "py", "rst", "rtf",
                    "md", "html",
@@ -535,7 +959,7 @@ file_types = non_image_types + image_types
 def add_meta(docs1, file):
     file_extension = pathlib.Path(file).suffix
     hashid = hash_file(file)
-    if not isinstance(docs1, list):
+    if not isinstance(docs1, (list, tuple, types.GeneratorType)):
         docs1 = [docs1]
     [x.metadata.update(dict(input_type=file_extension, date=str(datetime.now), hashid=hashid)) for x in docs1]
 
@@ -577,8 +1001,24 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False,
             else:
                 docs1 = []
         else:
+            if not (file.startswith("http://") or file.startswith("file://") or file.startswith("https://")):
+                file = 'http://' + file
             docs1 = UnstructuredURLLoader(urls=[file]).load()
+            if len(docs1) == 0 and have_playwright:
+                # then something went wrong, try another loader:
+                from langchain.document_loaders import PlaywrightURLLoader
+                docs1 = PlaywrightURLLoader(urls=[file]).load()
+            if len(docs1) == 0 and have_selenium:
+                # then something went wrong, try another loader:
+                # but requires Chrome binary, else get: selenium.common.exceptions.WebDriverException: Message: unknown error: cannot find Chrome binary
+                from langchain.document_loaders import SeleniumURLLoader
+                from selenium.common.exceptions import WebDriverException
+                try:
+                    docs1 = SeleniumURLLoader(urls=[file]).load()
+                except WebDriverException as e:
+                    print("No web driver: %s" % str(e), flush=True)
             [x.metadata.update(dict(input_type='url', date=str(datetime.now))) for x in docs1]
+        docs1 = clean_doc(docs1)
         doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif is_txt:
         base_path = "user_paste"
@@ -588,10 +1028,12 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False,
             f.write(file)
         metadata = dict(source=source_file, date=str(datetime.now()), input_type='pasted txt')
         doc1 = Document(page_content=file, metadata=metadata)
+        doc1 = clean_doc(doc1)
     elif file.lower().endswith('.html') or file.lower().endswith('.mhtml'):
         docs1 = UnstructuredHTMLLoader(file_path=file).load()
         add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
+        docs1 = clean_doc(docs1)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size, language=Language.HTML)
     elif (file.lower().endswith('.docx') or file.lower().endswith('.doc')) and have_libreoffice:
         docs1 = UnstructuredWordDocumentLoader(file_path=file).load()
         add_meta(docs1, file)
@@ -603,12 +1045,14 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False,
     elif file.lower().endswith('pptx') or file.lower().endswith('ppt'):
         docs1 = UnstructuredPowerPointLoader(file_path=file).load()
         add_meta(docs1, file)
+        docs1 = clean_doc(docs1)
         doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif file.lower().endswith('.txt'):
         # use UnstructuredFileLoader ?
         docs1 = TextLoader(file, encoding="utf8", autodetect_encoding=True).load()
         # makes just one, but big one
         doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
+        doc1 = clean_doc(doc1)
         add_meta(doc1, file)
     elif file.lower().endswith('.rtf'):
         docs1 = UnstructuredRTFLoader(file).load()
@@ -617,7 +1061,8 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False,
     elif file.lower().endswith('.md'):
         docs1 = UnstructuredMarkdownLoader(file).load()
         add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
+        docs1 = clean_doc(docs1)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size, language=Language.MARKDOWN)
     elif file.lower().endswith('.enex'):
         docs1 = EverNoteLoader(file).load()
         add_meta(doc1, file)
@@ -682,6 +1127,7 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False,
         with open(file, "r") as f:
             doc1 = Document(page_content=f.read(), metadata={"source": file})
         add_meta(doc1, file)
+        doc1 = chunk_sources(doc1, chunk=chunk, chunk_size=chunk_size, language=Language.RST)
     elif file.lower().endswith('.pdf'):
         env_gpt4all_file = ".env_gpt4all"
         from dotenv import dotenv_values
@@ -692,11 +1138,17 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False,
             from langchain.document_loaders import PyMuPDFLoader
             # load() still chunks by pages, but every page has title at start to help
             doc1 = PyMuPDFLoader(file).load()
+            doc1 = clean_doc(doc1)
+        elif pdf_class_name == 'UnstructuredPDFLoader':
+            doc1 = UnstructuredPDFLoader(file).load()
+            # seems to not need cleaning in most cases
         else:
             # open-source fallback
             # load() still chunks by pages, but every page has title at start to help
             doc1 = PyPDFLoader(file).load()
+            doc1 = clean_doc(doc1)
         # Some PDFs return nothing or junk from PDFMinerLoader
+        doc1 = chunk_sources(doc1, chunk=chunk, chunk_size=chunk_size)
         add_meta(doc1, file)
     elif file.lower().endswith('.csv'):
         doc1 = CSVLoader(file).load()
@@ -704,6 +1156,7 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False,
     elif file.lower().endswith('.py'):
         doc1 = PythonLoader(file).load()
         add_meta(doc1, file)
+        doc1 = chunk_sources(doc1, chunk=chunk, chunk_size=chunk_size, language=Language.PYTHON)
     elif file.lower().endswith('.toml'):
         doc1 = TomlLoader(file).load()
         add_meta(doc1, file)
@@ -794,15 +1247,16 @@ def path_to_docs(path_or_paths, verbose=False, fail_any_exception=False, n_jobs=
                  existing_files=[],
                  existing_hash_ids={},
                  ):
+    # path_or_paths could be str, list, tuple, generator
     globs_image_types = []
     globs_non_image_types = []
     if not path_or_paths and not url and not text:
         return []
     elif url:
-        globs_non_image_types = [url]
+        globs_non_image_types = url if isinstance(url, (list, tuple, types.GeneratorType)) else [url]
     elif text:
-        globs_non_image_types = [text]
-    elif isinstance(path_or_paths, str):
+        globs_non_image_types = text if isinstance(text, (list, tuple, types.GeneratorType)) else [text]
+    elif isinstance(path_or_paths, str) and os.path.isdir(path_or_paths):
         # single path, only consume allowed files
         path = path_or_paths
         # Below globs should match patterns in file_to_doc()
@@ -811,8 +1265,11 @@ def path_to_docs(path_or_paths, verbose=False, fail_any_exception=False, n_jobs=
         [globs_non_image_types.extend(glob.glob(os.path.join(path, "./**/*.%s" % ftype), recursive=True))
          for ftype in non_image_types]
     else:
+        if isinstance(path_or_paths, str) and (os.path.isfile(path_or_paths) or os.path.isdir(path_or_paths)):
+            path_or_paths = [path_or_paths]
         # list/tuple of files (consume what can, and exception those that selected but cannot consume so user knows)
-        assert isinstance(path_or_paths, (list, tuple)), "Wrong type for path_or_paths: %s" % type(path_or_paths)
+        assert isinstance(path_or_paths, (list, tuple, types.GeneratorType)), "Wrong type for path_or_paths: %s" % type(
+            path_or_paths)
         # reform out of allowed types
         globs_image_types.extend(flatten_list([[x for x in path_or_paths if x.endswith(y)] for y in image_types]))
         # could do below:
@@ -972,7 +1429,7 @@ def check_update_chroma_embedding(db, use_openai_embedding, hf_embedding_model,
     if load_embed(db) != (use_openai_embedding, hf_embedding_model):
         print("Detected new embedding, updating db: %s" % langchain_mode, flush=True)
         # handle embedding changes
-        db_get = db.get()
+        db_get = get_documents(db)
         sources = [Document(page_content=result[0], metadata=result[1] or {})
                    for result in zip(db_get['documents'], db_get['metadatas'])]
         # delete index, has to be redone
@@ -1023,14 +1480,17 @@ def get_existing_db(db, persist_directory, load_db_if_exists, db_type, use_opena
             if changed_db:
                 db = db_trial
                 # only call persist if really changed db, else takes too long for large db
-                db.persist()
-                clear_embedding(db)
+                if db is not None:
+                    db.persist()
+                    clear_embedding(db)
         save_embed(db, use_openai_embedding, hf_embedding_model)
         return db
     return None
 
 
 def clear_embedding(db):
+    if db is None:
+        return
     # don't keep on GPU, wastes memory, push back onto CPU and only put back on GPU once again embed
     db._embedding_function.client.cpu()
     clear_torch_cache()
@@ -1052,9 +1512,10 @@ def make_db(**langchain_kwargs):
 
 
 def save_embed(db, use_openai_embedding, hf_embedding_model):
-    embed_info_file = os.path.join(db._persist_directory, 'embed_info')
-    with open(embed_info_file, 'wb') as f:
-        pickle.dump((use_openai_embedding, hf_embedding_model), f)
+    if db is not None:
+        embed_info_file = os.path.join(db._persist_directory, 'embed_info')
+        with open(embed_info_file, 'wb') as f:
+            pickle.dump((use_openai_embedding, hf_embedding_model), f)
     return use_openai_embedding, hf_embedding_model
 
 
@@ -1201,28 +1662,90 @@ def _make_db(use_openai_embedding=False,
     return db, len(new_sources_metadata), new_sources_metadata
 
 
+def get_metadatas(db):
+    from langchain.vectorstores import FAISS
+    if isinstance(db, FAISS):
+        metadatas = [v.metadata for k, v in db.docstore._dict.items()]
+    elif isinstance(db, Chroma):
+        metadatas = get_documents(db)['metadatas']
+    else:
+        # FIXME: Hack due to https://github.com/weaviate/weaviate/issues/1947
+        # seems no way to get all metadata, so need to avoid this approach for weaviate
+        metadatas = [x.metadata for x in db.similarity_search("", k=10000)]
+    return metadatas
+
+
+def get_documents(db):
+    if hasattr(db, '_persist_directory'):
+        name_path = os.path.basename(db._persist_directory)
+        base_path = 'locks'
+        makedirs(base_path)
+        with filelock.FileLock(os.path.join(base_path, "getdb_%s.lock" % name_path)):
+            # get segfaults and other errors when multiple threads access this
+            return _get_documents(db)
+    else:
+        return _get_documents(db)
+
+
+def _get_documents(db):
+    from langchain.vectorstores import FAISS
+    if isinstance(db, FAISS):
+        documents = [v for k, v in db.docstore._dict.items()]
+    elif isinstance(db, Chroma):
+        documents = db.get()
+    else:
+        # FIXME: Hack due to https://github.com/weaviate/weaviate/issues/1947
+        # seems no way to get all metadata, so need to avoid this approach for weaviate
+        documents = [x for x in db.similarity_search("", k=10000)]
+    return documents
+
+
+def get_docs_and_meta(db, top_k_docs, filter_kwargs={}):
+    if hasattr(db, '_persist_directory'):
+        name_path = os.path.basename(db._persist_directory)
+        base_path = 'locks'
+        makedirs(base_path)
+        with filelock.FileLock(os.path.join(base_path, "getdb_%s.lock" % name_path)):
+            return _get_docs_and_meta(db, top_k_docs, filter_kwargs=filter_kwargs)
+    else:
+        return _get_docs_and_meta(db, top_k_docs, filter_kwargs=filter_kwargs)
+
+
+def _get_docs_and_meta(db, top_k_docs, filter_kwargs={}):
+    from langchain.vectorstores import FAISS
+    if isinstance(db, Chroma):
+        db_get = db._collection.get(where=filter_kwargs.get('filter'))
+        db_metadatas = db_get['metadatas']
+        db_documents = db_get['documents']
+    elif isinstance(db, FAISS):
+        import itertools
+        db_metadatas = get_metadatas(db)
+        # FIXME: FAISS has no filter
+        # slice dict first
+        db_documents = list(dict(itertools.islice(db.docstore._dict.items(), top_k_docs)).values())
+    else:
+        db_metadatas = get_metadatas(db)
+        db_documents = get_documents(db)
+    return db_documents, db_metadatas
+
+
 def get_existing_files(db):
-    collection = db.get()
-    metadata_sources = set([x['source'] for x in collection['metadatas']])
+    metadatas = get_metadatas(db)
+    metadata_sources = set([x['source'] for x in metadatas])
     return metadata_sources
 
 
 def get_existing_hash_ids(db):
-    collection = db.get()
+    metadatas = get_metadatas(db)
     # assume consistency, that any prior hashed source was single hashed file at the time among all source chunks
-    metadata_hash_ids = {x['source']: x.get('hashid') for x in collection['metadatas']}
+    metadata_hash_ids = {x['source']: x.get('hashid') for x in metadatas}
     return metadata_hash_ids
 
 
-source_prefix = "Sources [Score | Link]:"
-source_postfix = "End Sources<p>"
-
-
 def run_qa_db(**kwargs):
     func_names = list(inspect.signature(_run_qa_db).parameters)
     # hard-coded defaults
     kwargs['answer_with_sources'] = True
-    kwargs['sanitize_bot_response'] = True
     kwargs['show_rank'] = False
     missing_kwargs = [x for x in func_names if x not in kwargs]
     assert not missing_kwargs, "Missing kwargs: %s" % missing_kwargs
@@ -1240,7 +1763,7 @@ def _run_qa_db(query=None,
                user_path=None,
                detect_user_path_changes_every_query=False,
                db_type='faiss',
-               model_name=None, model=None, tokenizer=None,
+               model_name=None, model=None, tokenizer=None, inference_server=None,
                hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
                stream_output=False,
                prompter=None,
@@ -1248,7 +1771,7 @@ def _run_qa_db(query=None,
                prompt_dict=None,
                answer_with_sources=True,
                cut_distanct=1.1,
-               sanitize_bot_response=True,
+               sanitize_bot_response=False,
                show_rank=False,
                load_db_if_exists=False,
                db=None,
@@ -1267,7 +1790,12 @@ def _run_qa_db(query=None,
                document_choice=[DocumentChoices.All_Relevant.name],
                n_jobs=-1,
                verbose=False,
-               cli=False):
+               cli=False,
+               reverse_docs=True,
+               lora_weights='',
+               auto_reduce_chunks=True,
+               max_chunks=100,
+               ):
     """
 
     :param query:
@@ -1286,6 +1814,8 @@ def _run_qa_db(query=None,
     :param answer_with_sources
     :return:
     """
+    if model is not None:
+        assert model_name is not None  # require so can make decisions
     assert query is not None
     assert prompter is not None or prompt_type is not None or model is None  # if model is None, then will generate
     if prompter is not None:
@@ -1299,7 +1829,9 @@ def _run_qa_db(query=None,
             prompt_dict = ''
     assert len(set(gen_hyper).difference(inspect.signature(get_llm).parameters)) == 0
     llm, model_name, streamer, prompt_type_out = get_llm(use_openai_model=use_openai_model, model_name=model_name,
-                                                         model=model, tokenizer=tokenizer,
+                                                         model=model,
+                                                         tokenizer=tokenizer,
+                                                         inference_server=inference_server,
                                                          stream_output=stream_output,
                                                          do_sample=do_sample,
                                                          temperature=temperature,
@@ -1315,13 +1847,10 @@ def _run_qa_db(query=None,
                                                          prompt_type=prompt_type,
                                                          prompt_dict=prompt_dict,
                                                          prompter=prompter,
+                                                         sanitize_bot_response=sanitize_bot_response,
                                                          verbose=verbose,
                                                          )
 
-    if model_name in non_hf_types:
-        # FIXME: for now, streams to stdout/stderr currently
-        stream_output = False
-
     use_context = False
     scores = []
     chain = None
@@ -1349,43 +1878,49 @@ def _run_qa_db(query=None,
         # can only return if HF type
         return
 
-    if stream_output:
-        answer = None
-        assert streamer is not None
-        import queue
-        bucket = queue.Queue()
-        thread = EThread(target=chain, streamer=streamer, bucket=bucket)
-        thread.start()
-        outputs = ""
-        prompt = None  # FIXME
-        try:
-            for new_text in streamer:
-                # print("new_text: %s" % new_text, flush=True)
-                if bucket.qsize() > 0 or thread.exc:
-                    thread.join()
-                outputs += new_text
-                if prompter:  # and False:  # FIXME: pipeline can already use prompter
-                    output1 = prompter.get_response(outputs, prompt=prompt,
-                                                    sanitize_bot_response=sanitize_bot_response)
-                    yield output1, ''
-                else:
-                    yield outputs, ''
-        except BaseException:
-            # if any exception, raise that exception if was from thread, first
-            if thread.exc:
-                raise thread.exc
-            raise
-        finally:
-            # in case no exception and didn't join with thread yet, then join
-            if not thread.exc:
-                answer = thread.join()
-        # in case raise StopIteration or broke queue loop in streamer, but still have exception
-        if thread.exc:
-            raise thread.exc
-        # FIXME: answer is not string outputs from streamer.  How to get actual final output?
-        # answer = outputs
-    else:
-        answer = chain()
+    # context stuff similar to used in evaluate()
+    import torch
+    device, torch_dtype, context_class = get_device_dtype()
+    with torch.no_grad():
+        have_lora_weights = lora_weights not in [no_lora_str, '', None]
+        context_class_cast = NullContext if device == 'cpu' or have_lora_weights else torch.autocast
+        with context_class_cast(device):
+            if stream_output and streamer:
+                answer = None
+                import queue
+                bucket = queue.Queue()
+                thread = EThread(target=chain, streamer=streamer, bucket=bucket)
+                thread.start()
+                outputs = ""
+                prompt = None  # FIXME
+                try:
+                    for new_text in streamer:
+                        # print("new_text: %s" % new_text, flush=True)
+                        if bucket.qsize() > 0 or thread.exc:
+                            thread.join()
+                        outputs += new_text
+                        if prompter:  # and False:  # FIXME: pipeline can already use prompter
+                            output1 = prompter.get_response(outputs, prompt=prompt,
+                                                            sanitize_bot_response=sanitize_bot_response)
+                            yield output1, ''
+                        else:
+                            yield outputs, ''
+                except BaseException:
+                    # if any exception, raise that exception if was from thread, first
+                    if thread.exc:
+                        raise thread.exc
+                    raise
+                finally:
+                    # in case no exception and didn't join with thread yet, then join
+                    if not thread.exc:
+                        answer = thread.join()
+                # in case raise StopIteration or broke queue loop in streamer, but still have exception
+                if thread.exc:
+                    raise thread.exc
+                # FIXME: answer is not string outputs from streamer.  How to get actual final output?
+                # answer = outputs
+            else:
+                answer = chain()
 
     if not use_context:
         ret = answer['output_text']
@@ -1404,6 +1939,7 @@ def get_similarity_chain(query=None,
                          detect_user_path_changes_every_query=False,
                          db_type='faiss',
                          model_name=None,
+                         inference_server='',
                          hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
                          prompt_type=None,
                          prompt_dict=None,
@@ -1415,8 +1951,14 @@ def get_similarity_chain(query=None,
                          n_jobs=-1,
                          # beyond run_db_query:
                          llm=None,
+                         tokenizer=None,
                          verbose=False,
                          cmd=None,
+                         reverse_docs=True,
+
+                         # local
+                         auto_reduce_chunks=True,
+                         max_chunks=100,
                          ):
     # determine whether use of context out of docs is planned
     if not use_openai_model and prompt_type not in ['plain'] or model_name in non_hf_types:
@@ -1431,7 +1973,11 @@ def get_similarity_chain(query=None,
     # FIXME: Seems to way to get size of chroma db to limit top_k_docs to avoid
     # Chroma collection MyData contains fewer than 4 elements.
     # type logger error
-    k_db = 1000 if db_type == 'chroma' else top_k_docs  # top_k_docs=100 works ok too for
+    if top_k_docs == -1:
+        k_db = 1000 if db_type == 'chroma' else 100
+    else:
+        # top_k_docs=100 works ok too
+        k_db = 1000 if db_type == 'chroma' else top_k_docs
 
     # FIXME: For All just go over all dbs instead of a separate db for All
     if not detect_user_path_changes_every_query and db is not None:
@@ -1452,6 +1998,29 @@ def get_similarity_chain(query=None,
                                                         n_jobs=n_jobs,
                                                         verbose=verbose)
 
+    if 'falcon' in model_name:
+        extra = "According to only the information in the document sources provided within the context above, "
+        prefix = "Pay attention and remember information below, which will help to answer the question or imperative after the context ends."
+    elif inference_server in ['openai', 'openai_chat']:
+        extra = "According to (primarily) the information in the document sources provided within context above, "
+        prefix = "Pay attention and remember 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."
+    else:
+        extra = ""
+        prefix = ""
+    if langchain_mode in ['Disabled', 'ChatLLM', 'LLM'] or not use_context:
+        template_if_no_docs = template = """%s{context}{question}""" % prefix
+    else:
+        template = """%s
+\"\"\"
+{context}
+\"\"\"
+%s{question}""" % (prefix, extra)
+        template_if_no_docs = """%s{context}%s{question}""" % (prefix, extra)
+    if not use_openai_model and prompt_type not in ['plain'] or model_name in non_hf_types:
+        use_template = True
+    else:
+        use_template = False
+
     if db and use_context:
         if not isinstance(db, Chroma):
             # only chroma supports filtering
@@ -1472,17 +2041,84 @@ def get_similarity_chain(query=None,
             docs = []
             scores = []
         elif cmd == DocumentChoices.Only_All_Sources.name:
-            if isinstance(db, Chroma):
-                db_get = db._collection.get(where=filter_kwargs.get('filter'))
-            else:
-                db_get = db.get()
+            db_documents, db_metadatas = get_docs_and_meta(db, top_k_docs, filter_kwargs=filter_kwargs)
             # similar to langchain's chroma's _results_to_docs_and_scores
             docs_with_score = [(Document(page_content=result[0], metadata=result[1] or {}), 0)
-                               for result in zip(db_get['documents'], db_get['metadatas'])][:top_k_docs]
+                               for result in zip(db_documents, db_metadatas)][:top_k_docs]
             docs = [x[0] for x in docs_with_score]
             scores = [x[1] for x in docs_with_score]
         else:
-            docs_with_score = db.similarity_search_with_score(query, k=k_db, **filter_kwargs)[:top_k_docs]
+            if top_k_docs == -1 or auto_reduce_chunks:
+                # docs_with_score = db.similarity_search_with_score(query, k=k_db, **filter_kwargs)[:top_k_docs]
+                top_k_docs_tokenize = 100
+                base_path = 'locks'
+                makedirs(base_path)
+                if hasattr(db, '_persist_directory'):
+                    name_path = "sim_%s.lock" % os.path.basename(db._persist_directory)
+                else:
+                    name_path = "sim.lock"
+                with filelock.FileLock(os.path.join(base_path, name_path)):
+                    docs_with_score = db.similarity_search_with_score(query, k=k_db, **filter_kwargs)[
+                                      :top_k_docs_tokenize]
+                if hasattr(llm, 'pipeline') and hasattr(llm.pipeline, 'tokenizer'):
+                    # more accurate
+                    tokens = [len(llm.pipeline.tokenizer(x[0].page_content)['input_ids']) for x in docs_with_score]
+                    template_tokens = len(llm.pipeline.tokenizer(template)['input_ids'])
+                elif inference_server in ['openai', 'openai_chat'] or use_openai_model or db_type in ['faiss',
+                                                                                                      'weaviate']:
+                    # use ticktoken for faiss since embedding called differently
+                    tokens = [llm.get_num_tokens(x[0].page_content) for x in docs_with_score]
+                    template_tokens = llm.get_num_tokens(template)
+                elif isinstance(tokenizer, FakeTokenizer):
+                    tokens = [tokenizer.num_tokens_from_string(x[0].page_content) for x in docs_with_score]
+                    template_tokens = tokenizer.num_tokens_from_string(template)
+                else:
+                    # in case model is not our pipeline with HF tokenizer
+                    tokens = [db._embedding_function.client.tokenize([x[0].page_content])['input_ids'].shape[1] for x in
+                              docs_with_score]
+                    template_tokens = db._embedding_function.client.tokenize([template])['input_ids'].shape[1]
+                tokens_cumsum = np.cumsum(tokens)
+                if hasattr(llm, 'pipeline') and hasattr(llm.pipeline, 'max_input_tokens'):
+                    max_input_tokens = llm.pipeline.max_input_tokens
+                elif inference_server in ['openai']:
+                    max_tokens = llm.modelname_to_contextsize(model_name)
+                    # leave some room for 1 paragraph, even if min_new_tokens=0
+                    max_input_tokens = max_tokens - 256
+                elif inference_server in ['openai_chat']:
+                    max_tokens = model_token_mapping[model_name]
+                    # leave some room for 1 paragraph, even if min_new_tokens=0
+                    max_input_tokens = max_tokens - 256
+                elif isinstance(tokenizer, FakeTokenizer):
+                    max_input_tokens = tokenizer.model_max_length - 256
+                else:
+                    # leave some room for 1 paragraph, even if min_new_tokens=0
+                    max_input_tokens = 2048 - 256
+                max_input_tokens -= template_tokens
+                # FIXME: Doesn't account for query, == context, or new lines between contexts
+                where_res = np.where(tokens_cumsum < max_input_tokens)[0]
+                if where_res.shape[0] == 0:
+                    # then no chunk can fit, still do first one
+                    top_k_docs_trial = 1
+                else:
+                    top_k_docs_trial = 1 + where_res[-1]
+                if 0 < top_k_docs_trial < max_chunks:
+                    # avoid craziness
+                    if top_k_docs == -1:
+                        top_k_docs = top_k_docs_trial
+                    else:
+                        top_k_docs = min(top_k_docs, top_k_docs_trial)
+                if top_k_docs == -1:
+                    # if here, means 0 and just do best with 1 doc
+                    print("Unexpected large chunks and can't add to context, will add 1 anyways", flush=True)
+                    top_k_docs = 1
+                docs_with_score = docs_with_score[:top_k_docs]
+            else:
+                docs_with_score = db.similarity_search_with_score(query, k=k_db, **filter_kwargs)[:top_k_docs]
+            # put most relevant chunks closest to question,
+            # esp. if truncation occurs will be "oldest" or "farthest from response" text that is truncated
+            # BUT: for small models, e.g. 6_9 pythia, if sees some stuff related to h2oGPT first, it can connect that and not listen to rest
+            if reverse_docs:
+                docs_with_score.reverse()
             # cut off so no high distance docs/sources considered
             docs = [x[0] for x in docs_with_score if x[1] < cut_distanct]
             scores = [x[1] for x in docs_with_score if x[1] < cut_distanct]
@@ -1517,19 +2153,11 @@ def get_similarity_chain(query=None,
     if len(docs) == 0:
         # avoid context == in prompt then
         use_context = False
+        template = template_if_no_docs
 
-    if not use_openai_model and prompt_type not in ['plain'] or model_name in non_hf_types:
+    if use_template:
         # instruct-like, rather than few-shot prompt_type='plain' as default
         # but then sources confuse the model with how inserted among rest of text, so avoid
-        prefix = ""
-        if langchain_mode in ['Disabled', 'ChatLLM', 'LLM'] or not use_context:
-            template = """%s{context}{question}""" % prefix
-        else:
-            template = """%s
-==
-{context}
-==
-{question}""" % prefix
         prompt = PromptTemplate(
             # input_variables=["summaries", "question"],
             input_variables=["context", "question"],
@@ -1589,17 +2217,32 @@ def get_sources_answer(query, answer, scores, show_rank, answer_with_sources, ve
     return ret, extra
 
 
-def chunk_sources(sources, chunk=True, chunk_size=512):
+def clean_doc(docs1):
+    if not isinstance(docs1, (list, tuple, types.GeneratorType)):
+        docs1 = [docs1]
+    for doci, doc in enumerate(docs1):
+        docs1[doci].page_content = '\n'.join([x.strip() for x in doc.page_content.split("\n") if x.strip()])
+    return docs1
+
+
+def chunk_sources(sources, chunk=True, chunk_size=512, language=None):
     if not chunk:
         return sources
-    source_chunks = []
-    # Below for known separator
-    # splitter = CharacterTextSplitter(separator=" ", chunk_size=chunk_size, chunk_overlap=0)
-    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=0)
-    for source in sources:
-        # print(source.metadata['source'], flush=True)
-        for chunky in splitter.split_text(source.page_content):
-            source_chunks.append(Document(page_content=chunky, metadata=source.metadata))
+    if not isinstance(sources, (list, tuple, types.GeneratorType)) and not callable(sources):
+        # if just one document
+        sources = [sources]
+    if language and False:
+        # Bug in langchain, keep separator=True not working
+        # https://github.com/hwchase17/langchain/issues/2836
+        # so avoid this for now
+        keep_separator = True
+        separators = RecursiveCharacterTextSplitter.get_separators_for_language(language)
+    else:
+        separators = ["\n\n", "\n", " ", ""]
+        keep_separator = False
+    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=0, keep_separator=keep_separator,
+                                              separators=separators)
+    source_chunks = splitter.split_documents(sources)
     return source_chunks
 
 
@@ -1647,15 +2290,17 @@ def _create_local_weaviate_client():
     WEAVIATE_SCOPE = os.getenv('WEAVIATE_SCOPE', "offline_access")
 
     resource_owner_config = None
-    if WEAVIATE_USERNAME is not None and WEAVIATE_PASSWORD is not None:
-        resource_owner_config = weaviate.AuthClientPassword(
-            username=WEAVIATE_USERNAME,
-            password=WEAVIATE_PASSWORD,
-            scope=WEAVIATE_SCOPE
-        )
-
     try:
+        import weaviate
+        if WEAVIATE_USERNAME is not None and WEAVIATE_PASSWORD is not None:
+            resource_owner_config = weaviate.AuthClientPassword(
+                username=WEAVIATE_USERNAME,
+                password=WEAVIATE_PASSWORD,
+                scope=WEAVIATE_SCOPE
+            )
+
         client = weaviate.Client(WEAVIATE_URL, auth_client_secret=resource_owner_config)
+        return client
     except Exception as e:
         print(f"Failed to create Weaviate client: {e}")
         return None

gradio_themes.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from typing import Iterable
 
 from gradio.themes.soft import Soft
-from gradio.themes import Color
+from gradio.themes import Color, Size
 from gradio.themes.utils import colors, sizes, fonts
 
 h2o_yellow = Color(
@@ -36,6 +36,42 @@ h2o_gray = Color(
 )
 
 
+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,
@@ -158,19 +194,23 @@ h2o_logo = '<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/
            '11.93S497.56,252.55,497.56,246.06Zm2.52,21.47h20.68v71.31H500.08Z"/></svg>'
 
 
-def get_h2o_title(title):
-    return f"""<div style="display:flex; justify-content:center; margin-bottom:30px;">
+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></img>
+                    <img src="https://raw.githubusercontent.com/h2oai/h2ogpt/main/docs/h2o-qr.png">
                 </div>
                 """
 
 
-def get_simple_title(title):
-    return f"""<h1 align="center"> {title}</h1>"""
+def get_simple_title(title, description):
+    return f"""{description}<h1 align="center"> {title}</h1>"""
 
 
 def get_dark_js():

gradio_utils/css.py

@@ -0,0 +1,53 @@
+def get_css(kwargs) -> str:
+    if kwargs['h2ocolors']:
+        css_code = """footer {visibility: hidden;}
+        body{background:linear-gradient(#f5f5f5,#e5e5e5);}
+        body.dark{background:linear-gradient(#000000,#0d0d0d);}
+        """
+    else:
+        css_code = """footer {visibility: hidden}"""
+
+    css_code += make_css_base()
+    return css_code
+
+
+def make_css_base() -> str:
+    return """
+    @import url('https://fonts.googleapis.com/css2?family=Source+Sans+Pro:wght@400;600&display=swap');
+    
+    body.dark{#warning {background-color: #555555};}
+    
+    #small_btn {
+        margin: 0.6em 0em 0.55em 0;
+        max-width: 20em;
+        min-width: 5em !important;
+        height: 5em;
+        font-size: 14px !important;
+    }
+    
+    #prompt-form {
+        border: 1px solid var(--primary-500) !important;
+    }
+    
+    #prompt-form.block {
+        border-radius: var(--block-radius) !important;
+    }
+    
+    #prompt-form textarea {
+        border: 1px solid rgb(209, 213, 219);
+    }
+    
+    #prompt-form label > div {
+        margin-top: 4px;
+    }
+    
+    button.primary:hover {
+        background-color: var(--primary-600) !important;
+        transition: .2s;
+    }
+    
+    #prompt-form-area {
+        margin-bottom: 2.5rem;
+    }
+    .chatsmall chatbot {font-size: 10px !important}
+    """

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,118 @@
+import os
+import math
+
+import gradio as gr
+
+
+def make_chatbots(output_label0, output_label0_model2, **kwargs):
+    text_outputs = []
+    chat_kwargs = []
+    for model_state_lock in 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, visible=kwargs['model_lock'], elem_classes='chatsmall',
+                                height=kwargs['height'] or 400, min_width=min_width))
+
+    if kwargs['model_lock_columns'] == -1:
+        kwargs['model_lock_columns'] = len(kwargs['model_states'])
+    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(kwargs['model_states']) / 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(kwargs['model_states']) / 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(kwargs['model_states']) / 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(kwargs['model_states']) / 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(kwargs['model_states']) / 3 or mii >= 2 * len(kwargs['model_states']) / 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(kwargs['model_states']) / 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(kwargs['model_states']) / 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(kwargs['model_states']) / 4 or mii >= 2 * len(kwargs['model_states']) / 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(kwargs['model_states']) / 4 or mii >= 3 * len(kwargs['model_states']) / 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(kwargs['model_states']) / 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
+
+
+def make_prompt_form(kwargs):
+    if kwargs['input_lines'] > 1:
+        instruction_label = "Shift-Enter to Submit, Enter for more lines"
+    else:
+        instruction_label = "Enter to Submit, Shift-Enter for more lines"
+
+    with gr.Row():#elem_id='prompt-form-area'):
+        with gr.Column(scale=50):
+            instruction = gr.Textbox(
+                lines=kwargs['input_lines'],
+                label='Ask anything',
+                placeholder=instruction_label,
+                info=None,
+                elem_id='prompt-form',
+                container=True,
+            )
+        with gr.Row():
+            submit = gr.Button(value='Submit', variant='primary', scale=0, size='sm')
+            stop_btn = gr.Button(value="Stop", variant='secondary', scale=0, size='sm')
+
+    return instruction, submit, stop_btn

h2oai_pipeline.py

@@ -9,7 +9,7 @@ from prompter import Prompter, PromptType
 
 class H2OTextGenerationPipeline(TextGenerationPipeline):
     def __init__(self, *args, debug=False, chat=False, stream_output=False,
-                 sanitize_bot_response=True,
+                 sanitize_bot_response=False,
                  use_prompter=True, prompter=None,
                  prompt_type=None, prompt_dict=None,
                  max_input_tokens=2048 - 256, **kwargs):
@@ -51,25 +51,37 @@ class H2OTextGenerationPipeline(TextGenerationPipeline):
         self.sanitize_bot_response = sanitize_bot_response
         self.max_input_tokens = max_input_tokens  # not for generate, so ok that not kwargs
 
-    def preprocess(self, prompt_text, prefix="", handle_long_generation=None, **generate_kwargs):
-        if hasattr(self.tokenizer, 'model_max_length'):
+    @staticmethod
+    def limit_prompt(prompt_text, tokenizer, max_prompt_length=None):
+        verbose = bool(int(os.getenv('VERBOSE_PIPELINE', '0')))
+
+        if hasattr(tokenizer, 'model_max_length'):
             # model_max_length only defined for generate.py, not raw use of h2oai_pipeline.py
-            model_max_length = self.tokenizer.model_max_length
+            model_max_length = tokenizer.model_max_length
+            if max_prompt_length is not None:
+                model_max_length = min(model_max_length, max_prompt_length)
+            # cut at some upper likely limit to avoid excessive tokenization etc
+            # upper bound of 10 chars/token, e.g. special chars sometimes are long
+            if len(prompt_text) > model_max_length * 10:
+                len0 = len(prompt_text)
+                prompt_text = prompt_text[-model_max_length * 10:]
+                if verbose:
+                    print("Cut of input: %s -> %s" % (len0, len(prompt_text)), flush=True)
         else:
             # unknown
             model_max_length = None
 
-        verbose = bool(int(os.getenv('VERBOSE_PIPELINE', '0')))
+        num_prompt_tokens = None
         if model_max_length is not None:
-            num_prompt_tokens = None
             # can't wait for "hole" if not plain prompt_type, since would lose prefix like <human>:
             # For https://github.com/h2oai/h2ogpt/issues/192
             for trial in range(0, 3):
-                prompt_tokens = self.tokenizer(prompt_text)['input_ids']
+                prompt_tokens = tokenizer(prompt_text)['input_ids']
                 num_prompt_tokens = len(prompt_tokens)
                 if num_prompt_tokens > model_max_length:
                     # conservative by using int()
                     chars_per_token = int(len(prompt_text) / num_prompt_tokens)
+                    # keep tail, where question is if using langchain
                     prompt_text = prompt_text[-model_max_length * chars_per_token:]
                     if verbose:
                         print("reducing %s tokens, assuming average of %s chars/token for %s characters" % (
@@ -79,20 +91,27 @@ class H2OTextGenerationPipeline(TextGenerationPipeline):
                         print("using %s tokens with %s chars" % (num_prompt_tokens, len(prompt_text)), flush=True)
                     break
 
-            # if input prompt is some number of tokens, despite user request, can't have max_new_tokens more
-            #
-            if self.prompt_type not in [PromptType.plain.name, PromptType.plain.value]:
-                # then give room for prompt
-                fudge = 20
-            else:
-                fudge = 0
-            assert num_prompt_tokens is not None
-            max_new_tokens = max(0, min(generate_kwargs['max_new_tokens'],
-                                        model_max_length - (num_prompt_tokens + fudge)))
-            if max_new_tokens < generate_kwargs['max_new_tokens']:
-                if verbose:
-                    print("Reduced max_new_tokens from %s -> %s" % (generate_kwargs['max_new_tokens'], max_new_tokens))
-                generate_kwargs['max_new_tokens'] = max_new_tokens
+            # Why Below False: don't limit max_new_tokens more, just rely upon stopping to reach limit of model
+            if False:
+                # if input prompt is some number of tokens, despite user request, can't have max_new_tokens more
+                #
+                assert num_prompt_tokens is not None
+                if self.prompt_type not in [PromptType.plain.name, PromptType.plain.value]:
+                    # then give room for prompt
+                    fudge = 20
+                else:
+                    fudge = 0
+                max_new_tokens = max(0, min(generate_kwargs['max_new_tokens'],
+                                            model_max_length - (num_prompt_tokens + fudge)))
+                if max_new_tokens < generate_kwargs['max_new_tokens']:
+                    if verbose:
+                        print("Reduced max_new_tokens from %s -> %s" % (
+                        generate_kwargs['max_new_tokens'], max_new_tokens))
+                    generate_kwargs['max_new_tokens'] = max_new_tokens
+        return prompt_text, num_prompt_tokens
+
+    def preprocess(self, prompt_text, prefix="", handle_long_generation=None, **generate_kwargs):
+        prompt_text, num_prompt_tokens = H2OTextGenerationPipeline.limit_prompt(prompt_text, self.tokenizer)
 
         data_point = dict(context='', instruction=prompt_text, input='')
         if self.prompter is not None:
@@ -100,7 +119,7 @@ class H2OTextGenerationPipeline(TextGenerationPipeline):
         self.prompt_text = prompt_text
         if handle_long_generation is None:
             # forces truncation of inputs to avoid critical failure
-            handle_long_generation = 'hole'
+            handle_long_generation = None  # disable with new approaches
         return super().preprocess(prompt_text, prefix=prefix, handle_long_generation=handle_long_generation,
                                   **generate_kwargs)
 
@@ -113,7 +132,7 @@ class H2OTextGenerationPipeline(TextGenerationPipeline):
                 outputs = self.prompter.get_response(outputs, prompt=self.prompt_text,
                                                      sanitize_bot_response=self.sanitize_bot_response)
             elif self.bot and self.human:
-                outputs = rec['generated_text'].split(self.bot)[1].strip().split(self.human)[0].strip()
+                outputs = rec['generated_text'].split(self.bot)[1].split(self.human)[0]
             else:
                 outputs = rec['generated_text']
             rec['generated_text'] = outputs
@@ -123,7 +142,8 @@ class H2OTextGenerationPipeline(TextGenerationPipeline):
         if self.can_stop:
             stopping_criteria = get_stopping(self.prompt_type, self.prompt_dict,
                                              self.tokenizer, self.device,
-                                             human=self.human, bot=self.bot)
+                                             human=self.human, bot=self.bot,
+                                             model_max_length=self.tokenizer.model_max_length)
             generate_kwargs['stopping_criteria'] = stopping_criteria
         # return super()._forward(model_inputs, **generate_kwargs)
         return self.__forward(model_inputs, **generate_kwargs)

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 acutally 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)

loaders.py

@@ -1,6 +1,8 @@
-def get_loaders(llama_type, model_name, reward_type):
+def get_loaders(model_name, reward_type, llama_type=None):
     # NOTE: Some models need specific new prompt_type
     # E.g. t5_xxl_true_nli_mixture has input format: "premise: PREMISE_TEXT hypothesis: HYPOTHESIS_TEXT".)
+    if llama_type is None:
+        llama_type = "llama" in model_name.lower()
     if llama_type:
         from transformers import LlamaForCausalLM, LlamaTokenizer
         model_loader = LlamaForCausalLM
@@ -39,7 +41,8 @@ def get_tokenizer(tokenizer_loader, tokenizer_base_model, local_files_only, resu
     tokenizer = tokenizer_loader.from_pretrained(tokenizer_base_model,
                                                  local_files_only=local_files_only,
                                                  resume_download=resume_download,
-                                                 use_auth_token=use_auth_token)
+                                                 use_auth_token=use_auth_token,
+                                                 padding_side='left')
 
     tokenizer.pad_token_id = 0  # different from the eos token
     # when generating, we will use the logits of right-most token to predict the next token

prompter.py

@@ -1,10 +1,10 @@
+import os
 import ast
 import time
 from enums import PromptType  # also supports imports from this file from other files
 
 non_hf_types = ['gpt4all_llama', 'llama', 'gptj']
 
-
 prompt_type_to_model_name = {
     'plain': [
         'EleutherAI/gpt-j-6B',
@@ -25,23 +25,29 @@ prompt_type_to_model_name = {
         'mosaicml/mpt-7b-storywriter',
         'mosaicml/mpt-7b-instruct',  # internal code handles instruct
         'mosaicml/mpt-7b-chat',  # NC, internal code handles instruct
-        'gptj',  # internally handles prompting
-        'llama',  # plain, or need to choose prompt_type for given TheBloke model
-        'gpt4all_llama',  # internally handles prompting
+        'mosaicml/mpt-30b-instruct',  # internal code handles instruct
     ],
+    'gptj': ['gptj', 'gpt4all_llama'],
     'prompt_answer': [
         'h2oai/h2ogpt-gm-oasst1-en-1024-20b',
         'h2oai/h2ogpt-gm-oasst1-en-1024-12b',
         'h2oai/h2ogpt-gm-oasst1-multilang-1024-20b',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b-preview-300bt',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b-preview-300bt-v2',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b-preview-700bt',
-        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b',
         'h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b',
         'h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b-v2',
+        'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v3',
         'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b',
         'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v2',
         'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-40b-v1',
+        'h2oai/h2ogpt-gm-oasst1-en-2048-falcon-40b-v2',
+        'h2oai/h2ogpt-gm-oasst1-en-xgen-7b-8k',
+        'h2oai/h2ogpt-gm-oasst1-multilang-xgen-7b-8k',
+    ],
+    'prompt_answer_openllama': [
+        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b-preview-300bt',
+        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b-preview-300bt-v2',
+        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b-preview-700bt',
+        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-7b',
+        'h2oai/h2ogpt-gm-oasst1-en-2048-open-llama-13b',
     ],
     'instruct': [],
     'instruct_with_end': ['databricks/dolly-v2-12b'],
@@ -54,6 +60,7 @@ prompt_type_to_model_name = {
         'h2oai/h2ogpt-oig-oasst1-256-6.9b',  # legacy
         'h2oai/h2ogpt-oig-oasst1-512-6.9b',  # legacy
         'h2oai/h2ogpt-research-oasst1-512-30b',
+        'h2oai/h2ogpt-research-oasst1-llama-65b',
         'h2oai/h2ogpt-oasst1-falcon-40b',
         'h2oai/h2ogpt-oig-oasst1-falcon-40b',
     ],
@@ -66,7 +73,16 @@ prompt_type_to_model_name = {
     "wizard_lm": ['ehartford/WizardLM-7B-Uncensored', 'ehartford/WizardLM-13B-Uncensored'],
     "wizard_mega": ['openaccess-ai-collective/wizard-mega-13b'],
     "instruct_simple": ['JosephusCheung/Guanaco'],
+    "wizard_vicuna": ['ehartford/Wizard-Vicuna-13B-Uncensored'],
+    "wizard2": ['llama', 'mosaicml/mpt-30b-instruct'],
+    "vicuna11": ['lmsys/vicuna-33b-v1.3'],
+    # could be plain, but default is correct prompt_type for default TheBloke model ggml-wizardLM-7B.q4_2.bin
 }
+if os.getenv('OPENAI_API_KEY'):
+    prompt_type_to_model_name.update({
+        "openai": ["text-davinci-003", "text-curie-001", "text-babbage-001", "text-ada-001"],
+        "openai_chat": ["gpt-3.5-turbo", "gpt-3.5-turbo-16k"],
+    })
 
 inv_prompt_type_to_model_name = {v.strip(): k for k, l in prompt_type_to_model_name.items() for v in l}
 inv_prompt_type_to_model_lower = {v.strip().lower(): k for k, l in prompt_type_to_model_name.items() for v in l}
@@ -80,18 +96,29 @@ for p in PromptType:
     prompt_types.extend([p.name, p.value, str(p.value)])
 
 
-def get_prompt(prompt_type, prompt_dict, chat, context, reduced, return_dict=False):
+def get_prompt(prompt_type, prompt_dict, chat, context, reduced, making_context, return_dict=False):
     prompt_dict_error = ''
+    generates_leading_space = False
+
     if prompt_type == PromptType.custom.name and not isinstance(prompt_dict, dict):
         try:
             prompt_dict = ast.literal_eval(prompt_dict)
         except BaseException as e:
             prompt_dict_error = str(e)
-        if prompt_dict_error:
-            return dict(), prompt_dict_error
-
-    if prompt_type in [PromptType.custom.value, str(PromptType.custom.value),
-                       PromptType.custom.name]:
+    if prompt_dict_error:
+        promptA = None
+        promptB = None
+        PreInstruct = None
+        PreInput = ''
+        PreResponse = ''
+        terminate_response = None
+        chat_sep = ''
+        chat_turn_sep = ''
+        humanstr = ''
+        botstr = ''
+        generates_leading_space = False
+    elif prompt_type in [PromptType.custom.value, str(PromptType.custom.value),
+                         PromptType.custom.name]:
         promptA = prompt_dict.get('promptA', '')
         promptB = prompt_dict('promptB', '')
         PreInstruct = prompt_dict.get('PreInstruct', '')
@@ -99,21 +126,23 @@ def get_prompt(prompt_type, prompt_dict, chat, context, reduced, return_dict=Fal
         PreResponse = prompt_dict.get('PreResponse', '')
         terminate_response = prompt_dict.get('terminate_response', None)
         chat_sep = prompt_dict.get('chat_sep', '\n')
+        chat_turn_sep = prompt_dict.get('chat_turn_sep', '\n')
         humanstr = prompt_dict.get('humanstr', '')
         botstr = prompt_dict.get('botstr', '')
     elif prompt_type in [PromptType.plain.value, str(PromptType.plain.value),
                          PromptType.plain.name]:
-        promptA = promptB = PreInstruct = PreInput = PreResponse = ''
+        promptA = promptB = PreInstruct = PreInput = PreResponse = None
         terminate_response = []
-        chat_sep = ''
-        humanstr = ''
-        botstr = ''
+        chat_turn_sep = chat_sep = ''
+        # plain should have None for human/bot, so nothing truncated out, not '' that would truncate after first token
+        humanstr = None
+        botstr = None
     elif prompt_type == 'simple_instruct':
         promptA = promptB = PreInstruct = PreInput = PreResponse = None
         terminate_response = []
-        chat_sep = '\n'
-        humanstr = ''
-        botstr = ''
+        chat_turn_sep = chat_sep = '\n'
+        humanstr = None
+        botstr = None
     elif prompt_type in [PromptType.instruct.value, str(PromptType.instruct.value),
                          PromptType.instruct.name] + [PromptType.instruct_with_end.value,
                                                       str(PromptType.instruct_with_end.value),
@@ -139,7 +168,7 @@ def get_prompt(prompt_type, prompt_dict, chat, context, reduced, return_dict=Fal
             terminate_response = ['### End']
         else:
             terminate_response = None
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct
         botstr = PreResponse
     elif prompt_type in [PromptType.quality.value, str(PromptType.quality.value),
@@ -161,7 +190,7 @@ def get_prompt(prompt_type, prompt_dict, chat, context, reduced, return_dict=Fal
 ### Response:
 """
         terminate_response = None
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct  # first thing human says
         botstr = PreResponse  # first thing bot says
     elif prompt_type in [PromptType.human_bot.value, str(PromptType.human_bot.value),
@@ -183,14 +212,14 @@ Current Time: {}
 
 """
             preprompt = PRE_PROMPT.format(cur_date, cur_time)
-        start = human
-        promptB = promptA = '%s%s ' % (preprompt, start)
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
 
-        PreInstruct = ""
+        PreInstruct = human + ' '
 
         PreInput = None
 
-        if reduced:
+        if making_context:
             # when making context, want it to appear as-if LLM generated, which starts with space after :
             PreResponse = bot + ' '
         else:
@@ -198,10 +227,11 @@ Current Time: {}
             # if add space here, non-unique tokenization will often make LLM produce wrong output
             PreResponse = bot
 
-        terminate_response = [start, PreResponse]
-        chat_sep = '\n'
+        terminate_response = ['\n' + human, '\n' + bot, human, bot, PreResponse]
+        chat_turn_sep = chat_sep = '\n'
         humanstr = human  # tag before human talks
         botstr = bot  # tag before bot talks
+        generates_leading_space = True
     elif prompt_type in [PromptType.dai_faq.value, str(PromptType.dai_faq.value),
                          PromptType.dai_faq.name]:
         promptA = ''
@@ -217,7 +247,7 @@ Current Time: {}
 ### Driverless AI documentation answer:
 """
         terminate_response = ['\n\n']
-        chat_sep = terminate_response
+        chat_turn_sep = chat_sep = terminate_response
         humanstr = PreInstruct
         botstr = PreResponse
     elif prompt_type in [PromptType.summarize.value, str(PromptType.summarize.value),
@@ -226,7 +256,7 @@ Current Time: {}
         PreInstruct = '## Main Text\n\n'
         PreResponse = '\n\n## Summary\n\n'
         terminate_response = None
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct
         botstr = PreResponse
     elif prompt_type in [PromptType.instruct_vicuna.value, str(PromptType.instruct_vicuna.value),
@@ -246,7 +276,7 @@ Current Time: {}
 """
         terminate_response = [
             '### Human:']  # but only allow terminate after prompt is found correctly, else can't terminate
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct
         botstr = PreResponse
     elif prompt_type in [PromptType.prompt_answer.value, str(PromptType.prompt_answer.value),
@@ -254,33 +284,50 @@ Current Time: {}
         preprompt = ''
         prompt_tokens = "<|prompt|>"
         answer_tokens = "<|answer|>"
-        start = prompt_tokens
+        start = ''
         promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = ""
+        PreInstruct = prompt_tokens
         PreInput = None
         PreResponse = answer_tokens
         eos = '<|endoftext|>'  # neox eos
-        terminate_response = [start, PreResponse, eos]
-        chat_sep = eos
         humanstr = prompt_tokens
         botstr = answer_tokens
+        terminate_response = [humanstr, PreResponse, eos]
+        chat_sep = ''
+        chat_turn_sep = eos
+    elif prompt_type in [PromptType.prompt_answer_openllama.value, str(PromptType.prompt_answer_openllama.value),
+                         PromptType.prompt_answer_openllama.name]:
+        preprompt = ''
+        prompt_tokens = "<|prompt|>"
+        answer_tokens = "<|answer|>"
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        PreInstruct = prompt_tokens
+        PreInput = None
+        PreResponse = answer_tokens
+        eos = '</s>'  # llama eos
+        humanstr = prompt_tokens
+        botstr = answer_tokens
+        terminate_response = [humanstr, PreResponse, eos]
+        chat_sep = ''
+        chat_turn_sep = eos
     elif prompt_type in [PromptType.open_assistant.value, str(PromptType.open_assistant.value),
                          PromptType.open_assistant.name]:
         # From added_tokens.json
         preprompt = ''
         prompt_tokens = "<|prompter|>"
         answer_tokens = "<|assistant|>"
-        start = prompt_tokens
+        start = ''
         promptB = promptA = '%s%s' % (preprompt, start)
-        PreInstruct = ""
+        PreInstruct = prompt_tokens
         PreInput = None
         PreResponse = answer_tokens
         pend = "<|prefix_end|>"
         eos = "</s>"
-        terminate_response = [start, PreResponse, pend, eos]
-        chat_sep = eos
         humanstr = prompt_tokens
         botstr = answer_tokens
+        terminate_response = [humanstr, PreResponse, pend, eos]
+        chat_turn_sep = chat_sep = eos
     elif prompt_type in [PromptType.wizard_lm.value, str(PromptType.wizard_lm.value),
                          PromptType.wizard_lm.name]:
         # https://github.com/ehartford/WizardLM/blob/main/src/train_freeform.py
@@ -292,7 +339,7 @@ Current Time: {}
         PreResponse = "\n\n### Response\n"
         eos = "</s>"
         terminate_response = [PreResponse, eos]
-        chat_sep = eos
+        chat_turn_sep = chat_sep = eos
         humanstr = promptA
         botstr = PreResponse
     elif prompt_type in [PromptType.wizard_mega.value, str(PromptType.wizard_mega.value),
@@ -308,13 +355,12 @@ Current Time: {}
 ### Assistant:
 """
         terminate_response = [PreResponse]
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct
         botstr = PreResponse
     elif prompt_type in [PromptType.instruct_vicuna2.value, str(PromptType.instruct_vicuna2.value),
                          PromptType.instruct_vicuna2.name]:
-        promptA = promptB = "" if not (
-                chat and reduced) else ''
+        promptA = promptB = "" if not (chat and reduced) else ''
 
         PreInstruct = """
 HUMAN:
@@ -327,13 +373,12 @@ ASSISTANT:
 """
         terminate_response = [
             'HUMAN:']  # but only allow terminate after prompt is found correctly, else can't terminate
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct
         botstr = PreResponse
     elif prompt_type in [PromptType.instruct_vicuna3.value, str(PromptType.instruct_vicuna3.value),
                          PromptType.instruct_vicuna3.name]:
-        promptA = promptB = "" if not (
-                chat and reduced) else ''
+        promptA = promptB = "" if not (chat and reduced) else ''
 
         PreInstruct = """
 ### User:
@@ -346,13 +391,14 @@ ASSISTANT:
 """
         terminate_response = [
             '### User:']  # but only allow terminate after prompt is found correctly, else can't terminate
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct
         botstr = PreResponse
     elif prompt_type in [PromptType.wizard2.value, str(PromptType.wizard2.value),
                          PromptType.wizard2.name]:
         # https://huggingface.co/TheBloke/WizardLM-7B-uncensored-GGML
-        preprompt = """Below is an instruction that describes a task. Write a response that appropriately completes the request."""
+        preprompt = """Below is an instruction that describes a task. Write a response that appropriately completes the request.""" if not (
+                chat and reduced) else ''
         start = ''
         promptB = promptA = '%s%s' % (preprompt, start)
         PreInstruct = """
@@ -363,27 +409,39 @@ ASSISTANT:
 ### Response:
 """
         terminate_response = [PreResponse]
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct
         botstr = PreResponse
     elif prompt_type in [PromptType.wizard3.value, str(PromptType.wizard3.value),
                          PromptType.wizard3.name]:
         # https://huggingface.co/TheBloke/wizardLM-13B-1.0-GGML
-        preprompt = """A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions."""
+        preprompt = """A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.""" if not (
+                chat and reduced) else ''
         start = ''
         promptB = promptA = '%s%s' % (preprompt, start)
         PreInstruct = """USER: """
         PreInput = None
         PreResponse = """ASSISTANT: """
         terminate_response = [PreResponse]
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
+        humanstr = PreInstruct
+        botstr = PreResponse
+    elif prompt_type in [PromptType.wizard_vicuna.value, str(PromptType.wizard_vicuna.value),
+                         PromptType.wizard_vicuna.name]:
+        preprompt = ''
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        PreInstruct = """USER: """
+        PreInput = None
+        PreResponse = """ASSISTANT: """
+        terminate_response = [PreResponse]
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct
         botstr = PreResponse
 
     elif prompt_type in [PromptType.instruct_simple.value, str(PromptType.instruct_simple.value),
                          PromptType.instruct_simple.name]:
-        promptA = '' if not (chat and reduced) else ''
-        promptB = '' if not (chat and reduced) else ''
+        promptB = promptA = '' if not (chat and reduced) else ''
 
         PreInstruct = """
 ### Instruction:
@@ -397,21 +455,90 @@ ASSISTANT:
 ### Response:
 """
         terminate_response = None
-        chat_sep = '\n'
+        chat_turn_sep = chat_sep = '\n'
         humanstr = PreInstruct
         botstr = PreResponse
+    elif prompt_type in [PromptType.openai.value, str(PromptType.openai.value),
+                         PromptType.openai.name]:
+        preprompt = """The following is a conversation with an AI assistant. The assistant is helpful, creative, clever, and very friendly.""" if not (
+                chat and reduced) else ''
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        PreInstruct = "\nHuman: "
+        PreInput = None
+        PreResponse = "\nAI:"
+        terminate_response = [PreResponse] + [" Human:", " AI:"]
+        chat_turn_sep = chat_sep = '\n'
+        humanstr = PreInstruct
+        botstr = PreResponse
+    elif prompt_type in [PromptType.gptj.value, str(PromptType.gptj.value),
+                         PromptType.gptj.name]:
+        preprompt = "### Instruction:\n The prompt below is a question to answer, a task to complete, or a conversation to respond to; decide which and write an appropriate response." if not (
+                chat and reduced) else ''
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        PreInstruct = "\n### Prompt: "
+        PreInput = None
+        PreResponse = "\n### Response: "
+        terminate_response = [PreResponse] + ["Prompt:", "Response:"]
+        chat_turn_sep = chat_sep = '\n'
+        humanstr = PreInstruct
+        botstr = PreResponse
+    elif prompt_type in [PromptType.openai_chat.value, str(PromptType.openai_chat.value),
+                         PromptType.openai_chat.name]:
+        # prompting and termination all handled by endpoint
+        preprompt = """"""
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        PreInstruct = ""
+        PreInput = None
+        PreResponse = ""
+        terminate_response = []
+        chat_turn_sep = chat_sep = '\n'
+        humanstr = None
+        botstr = None
+    elif prompt_type in [PromptType.vicuna11.value, str(PromptType.vicuna11.value),
+                         PromptType.vicuna11.name]:
+        preprompt = """A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. """ if not (
+                chat and reduced) else ''
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        eos = '</s>'
+        PreInstruct = """USER: """
+        PreInput = None
+        PreResponse = """ASSISTANT:"""
+        terminate_response = [PreResponse]
+        chat_sep = ' '
+        chat_turn_sep = eos
+        humanstr = PreInstruct
+        botstr = PreResponse
+
+        if making_context:
+            # when making context, want it to appear as-if LLM generated, which starts with space after :
+            PreResponse = PreResponse + ' '
+        else:
+            # normally LLM adds space after this, because was how trained.
+            # if add space here, non-unique tokenization will often make LLM produce wrong output
+            PreResponse = PreResponse
     else:
         raise RuntimeError("No such prompt_type=%s" % prompt_type)
 
-    if return_dict:
-        return dict(promptA=promptA, promptB=promptB, PreInstruct=PreInstruct, PreInput=PreInput,
+    if isinstance(terminate_response, (tuple, list)):
+        assert '' not in terminate_response, "Bad terminate_response"
+
+    ret_dict = dict(promptA=promptA, promptB=promptB, PreInstruct=PreInstruct, PreInput=PreInput,
                     PreResponse=PreResponse, terminate_response=terminate_response, chat_sep=chat_sep,
-                    humanstr=humanstr, botstr=botstr), ''
+                    chat_turn_sep=chat_turn_sep,
+                    humanstr=humanstr, botstr=botstr,
+                    generates_leading_space=generates_leading_space)
+
+    if return_dict:
+        return ret_dict, prompt_dict_error
     else:
-        return promptA, promptB, PreInstruct, PreInput, PreResponse, terminate_response, chat_sep, humanstr, botstr
+        return tuple(list(ret_dict.values()))
 
 
-def generate_prompt(data_point, prompt_type, prompt_dict, chat, reduced):
+def generate_prompt(data_point, prompt_type, prompt_dict, chat, reduced, making_context):
     context = data_point.get('context')
     if context is None:
         context = ''
@@ -422,9 +549,12 @@ def generate_prompt(data_point, prompt_type, prompt_dict, chat, reduced):
     prompt_dict = data_point.get('prompt_dict', prompt_dict)
     assert prompt_type in prompt_types, "Bad prompt type: %s" % prompt_type
     promptA, promptB, PreInstruct, PreInput, PreResponse, \
-        terminate_response, chat_sep, humanstr, botstr = get_prompt(prompt_type, prompt_dict, chat, context, reduced)
+        terminate_response, chat_sep, chat_turn_sep, humanstr, botstr, \
+        generates_leading_space = get_prompt(prompt_type, prompt_dict, chat,
+                                             context, reduced, making_context)
 
-    prompt = context if not reduced else ''
+    # could avoid if reduce=True, but too complex for parent functions to handle
+    prompt = context
 
     if input and promptA:
         prompt += f"""{promptA}"""
@@ -433,37 +563,37 @@ def generate_prompt(data_point, prompt_type, prompt_dict, chat, reduced):
 
     if instruction and PreInstruct is not None and input and PreInput is not None:
         prompt += f"""{PreInstruct}{instruction}{PreInput}{input}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
     elif instruction and input and PreInstruct is None and PreInput is not None:
         prompt += f"""{PreInput}{instruction}
 {input}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
     elif input and instruction and PreInput is None and PreInstruct is not None:
         prompt += f"""{PreInstruct}{instruction}
 {input}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
     elif instruction and PreInstruct is not None:
         prompt += f"""{PreInstruct}{instruction}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
     elif input and PreInput is not None:
         prompt += f"""{PreInput}{input}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
     elif input and instruction and PreInput is not None:
         prompt += f"""{PreInput}{instruction}{input}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
     elif input and instruction and PreInstruct is not None:
         prompt += f"""{PreInstruct}{instruction}{input}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
     elif input and instruction:
         # i.e. for simple_instruct
         prompt += f"""{instruction}: {input}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
     elif input:
         prompt += f"""{input}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
     elif instruction:
         prompt += f"""{instruction}"""
-        prompt = inject_newline(prompt_type, prompt)
+        prompt = inject_chatsep(prompt_type, prompt, chat_sep=chat_sep)
 
     if PreResponse is not None:
         prompt += f"""{PreResponse}"""
@@ -474,13 +604,13 @@ def generate_prompt(data_point, prompt_type, prompt_dict, chat, reduced):
     if output:
         prompt += f"""{output}"""
 
-    return prompt, pre_response, terminate_response, chat_sep
+    return prompt, pre_response, terminate_response, chat_sep, chat_turn_sep
 
 
-def inject_newline(prompt_type, prompt):
-    if prompt_type not in [-1, '-1', 'plain', 'simple_instruct']:
+def inject_chatsep(prompt_type, prompt, chat_sep=None):
+    if chat_sep:
         # only add new line if structured prompt, while 'plain' is just generation of next tokens from input
-        prompt += '\n'
+        prompt += chat_sep
     return prompt
 
 
@@ -489,9 +619,6 @@ class Prompter(object):
                  allowed_repeat_line_length=10):
         self.prompt_type = prompt_type
         self.prompt_dict = prompt_dict
-        data_point = dict(instruction='', input='', output='')
-        _, self.pre_response, self.terminate_response, self.chat_sep = \
-            generate_prompt(data_point, self.prompt_type, self.prompt_dict, chat, False)
         self.debug = debug
         self.chat = chat
         self.stream_output = stream_output
@@ -500,23 +627,41 @@ class Prompter(object):
         self.prompt = None
         context = ""  # not for chat context
         reduced = False  # not for chat context
+        making_context = False  # not for chat context
         self.promptA, self.promptB, self.PreInstruct, self.PreInput, self.PreResponse, \
-            self.terminate_response, self.chat_sep, self.humanstr, self.botstr = \
-            get_prompt(self.prompt_type, self.prompt_dict, chat, context, reduced)
-
-    def generate_prompt(self, data_point):
-        reduced = False
-        prompt, _, _, _ = generate_prompt(data_point, self.prompt_type, self.prompt_dict, self.chat, reduced)
+            self.terminate_response, self.chat_sep, self.chat_turn_sep, self.humanstr, self.botstr, \
+            self.generates_leading_space = \
+            get_prompt(self.prompt_type, self.prompt_dict, chat, context, reduced, making_context)
+        self.pre_response = self.PreResponse
+
+    def generate_prompt(self, data_point, reduced=None):
+        """
+        data_point['context'] is assumed to be like a system prompt or pre-conversation, not inserted after user prompt
+        :param data_point:
+        :param reduced:
+        :return:
+        """
+        reduced = data_point.get('context') not in ['', None] if reduced is None else reduced
+        making_context = False  # whether really making final prompt or just generating context
+        prompt, _, _, _, _ = generate_prompt(data_point, self.prompt_type, self.prompt_dict, self.chat, reduced,
+                                             making_context)
         if self.debug:
-            print("prompt: ", prompt, flush=True)
+            print("prompt: %s" % prompt, flush=True)
+        # if have context, should have always reduced and only preappend promptA/B here
+        if data_point.get('context'):
+            if data_point.get('input') and self.promptA:
+                prompt = self.promptA + prompt
+            elif self.promptB:
+                prompt = self.promptB + prompt
+
         self.prompt = prompt
         return prompt
 
-    def get_response(self, outputs, prompt=None, sanitize_bot_response=True):
+    def get_response(self, outputs, prompt=None, sanitize_bot_response=False):
         if isinstance(outputs, str):
             outputs = [outputs]
         if self.debug:
-            print("output:\n", '\n\n'.join(outputs), flush=True)
+            print("output:\n%s" % '\n\n'.join(outputs), flush=True)
         if prompt is not None:
             self.prompt = prompt
 
@@ -527,7 +672,8 @@ class Prompter(object):
             if sanitize_bot_response:
                 from better_profanity import profanity
                 response = profanity.censor(response)
-            response = response.strip("\n")
+            if self.generates_leading_space and isinstance(response, str) and len(response) > 0 and response[0] == ' ':
+                response = response[1:]
             return response
 
         def clean_repeats(response):
@@ -549,12 +695,12 @@ class Prompter(object):
                 # then use most basic parsing like pipeline
                 if self.botstr in output:
                     if self.humanstr:
-                        output = clean_response(output.split(self.botstr)[1].strip().split(self.humanstr)[0].strip())
+                        output = clean_response(output.split(self.botstr)[1].split(self.humanstr)[0])
                     else:
                         # i.e. use after bot but only up to next bot
-                        output = clean_response(output.split(self.botstr)[1].strip().split(self.botstr)[0].strip())
+                        output = clean_response(output.split(self.botstr)[1].split(self.botstr)[0])
                 else:
-                    # output = clean_response(output.strip())
+                    # output = clean_response(output)
                     # assume just not printed yet
                     output = ""
             else:
@@ -581,9 +727,9 @@ class Prompter(object):
                     allow_terminate = True
                     output = output[len(prompt):]
                 # clean after subtract prompt out, so correct removal of pre_response
-                output = clean_response(output).strip()
+                output = clean_response(output)
                 if self.repeat_penalty:
-                    output = clean_repeats(output).strip()
+                    output = clean_repeats(output)
                 if self.terminate_response and allow_terminate:
                     finds = []
                     for term in self.terminate_response:
@@ -591,11 +737,9 @@ class Prompter(object):
                     finds = [x for x in finds if x >= 0]
                     if len(finds) > 0:
                         termi = finds[0]
-                        output = output[:termi].strip()
+                        output = output[:termi]
                     else:
-                        output = output.strip()
-                else:
-                    output = output.strip()
+                        output = output
             if multi_output:
                 # prefix with output counter
                 output = "\n=========== Output %d\n\n" % (1 + oi) + output
@@ -606,5 +750,5 @@ class Prompter(object):
         # join all outputs, only one extra new line between outputs
         output = '\n'.join(outputs)
         if self.debug:
-            print("outputclean:\n", '\n\n'.join(outputs), flush=True)
+            print("outputclean:\n%s" % '\n\n'.join(outputs), flush=True)
         return output

requirements.txt

@@ -1,50 +1,50 @@
 # for generate (gradio server) and finetune
-datasets==2.12.0
-sentencepiece==0.1.97
-gradio==3.34.0
-huggingface_hub==0.14.1
+datasets==2.13.0
+sentencepiece==0.1.99
+gradio==3.35.2
+huggingface_hub==0.15.1
 appdirs==1.4.4
 fire==0.5.0
-docutils==0.19
+docutils==0.20.1
 torch==2.0.1
 evaluate==0.4.0
 rouge_score==0.1.2
 sacrebleu==2.3.1
 scikit-learn==1.2.2
 alt-profanity-check==1.2.2
-better-profanity==0.6.1
-numpy==1.24.2
-pandas==2.0.0
+better-profanity==0.7.0
+numpy==1.24.3
+pandas==2.0.2
 matplotlib==3.7.1
 loralib==0.1.1
 bitsandbytes==0.39.0
-accelerate==0.19.0
-git+https://github.com/huggingface/peft.git@3714aa2fff158fdfa637b2b65952580801d890b2
-transformers==4.28.1
+accelerate==0.20.3
+git+https://github.com/huggingface/peft.git@0b62b4378b4ce9367932c73540349da9a41bdea8
+transformers==4.30.2
 tokenizers==0.13.3
 APScheduler==3.10.1
 
 # optional for generate
 pynvml==11.5.0
-psutil==5.9.4
+psutil==5.9.5
 boto3==1.26.101
 botocore==1.29.101
 
 # optional for finetune
-tensorboard==2.12.1
-neptune==1.1.1
+tensorboard==2.13.0
+neptune==1.2.0
 
 # for gradio client
-gradio_client==0.2.6
+gradio_client==0.2.7
 beautifulsoup4==4.12.2
-markdown==3.4.1
+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
+# pandoc==2.3
 #pypandoc==1.11
 pypandoc_binary==1.11
 openpyxl==3.1.2
@@ -57,17 +57,20 @@ 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
+# optional: for OpenAI endpoint or embeddings (requires key)
+openai==0.27.8
 # optional for chat with PDF
-langchain==0.0.193
-pypdf==3.8.1
-tiktoken==0.3.3
+langchain==0.0.202
+pypdf==3.9.1
 # avoid textract, requires old six
 #textract==1.6.5
 
 # for HF embeddings
 sentence_transformers==2.2.2
-# for OpenAI embeddings (requires key)
-openai==0.27.6
 
 # local vector db
 chromadb==0.3.25
@@ -79,14 +82,14 @@ chromadb==0.3.25
 
 # strong support for images
 # Requires on Ubuntu: sudo apt-get install libmagic-dev poppler-utils tesseract-ocr libreoffice
-unstructured[local-inference]==0.6.6
+unstructured[local-inference]==0.7.4
 #pdf2image==1.16.3
 #pytesseract==0.3.10
 pillow
 
 pdfminer.six==20221105
-urllib3==1.26.6
-requests_file==1.5.1
+urllib3
+requests_file
 
 #pdf2image==1.16.3
 #pytesseract==0.3.10
@@ -101,18 +104,15 @@ tabulate==0.9.0
 pip-licenses==4.3.0
 
 # weaviate vector db
-weaviate-client==3.19.2
+weaviate-client==3.20.0
 # optional for chat with PDF
-langchain==0.0.193
-pypdf==3.8.1
-tiktoken==0.3.3
+langchain==0.0.202
+pypdf==3.9.1
 # avoid textract, requires old six
 #textract==1.6.5
 
 # for HF embeddings
 sentence_transformers==2.2.2
-# for OpenAI embeddings (requires key)
-openai==0.27.6
 
 # local vector db
 chromadb==0.3.25
@@ -124,14 +124,14 @@ chromadb==0.3.25
 
 # strong support for images
 # Requires on Ubuntu: sudo apt-get install libmagic-dev poppler-utils tesseract-ocr libreoffice
-unstructured[local-inference]==0.6.6
+unstructured[local-inference]==0.7.4
 #pdf2image==1.16.3
 #pytesseract==0.3.10
 pillow
 
 pdfminer.six==20221105
-urllib3==1.26.6
-requests_file==1.5.1
+urllib3
+requests_file
 
 #pdf2image==1.16.3
 #pytesseract==0.3.10
@@ -146,7 +146,7 @@ tabulate==0.9.0
 pip-licenses==4.3.0
 
 # weaviate vector db
-weaviate-client==3.19.2
+weaviate-client==3.20.0
 faiss-gpu==1.7.2
 arxiv==1.4.7
 pymupdf==1.22.3 # AGPL license

stopping.py

@@ -1,17 +1,18 @@
 import torch
 from transformers import StoppingCriteria, StoppingCriteriaList
 
-from prompter import PromptType
+from enums import PromptType
 
 
 class StoppingCriteriaSub(StoppingCriteria):
 
-    def __init__(self, stops=[], encounters=[], device="cuda"):
+    def __init__(self, stops=[], encounters=[], device="cuda", model_max_length=None):
         super().__init__()
         assert len(stops) % len(encounters) == 0, "Number of stops and encounters must match"
         self.encounters = encounters
         self.stops = [stop.to(device) for stop in stops]
         self.num_stops = [0] * len(stops)
+        self.model_max_length = model_max_length
 
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
         for stopi, stop in enumerate(self.stops):
@@ -20,12 +21,15 @@ class StoppingCriteriaSub(StoppingCriteria):
                 if self.num_stops[stopi] >= self.encounters[stopi % len(self.encounters)]:
                     # print("Stopped", flush=True)
                     return True
+        if self.model_max_length is not None and input_ids[0].shape[0] >= self.model_max_length:
+            # critical limit
+            return True
         # print("Tokens: %s" % input_ids[0].cpu().numpy(), flush=True)
         # print("Stop Tokens: %s" % [x.cpu().numpy() for x in self.stops], flush=True)
         return False
 
 
-def get_stopping(prompt_type, prompt_dict, tokenizer, device, human='<human>:', bot="<bot>:"):
+def get_stopping(prompt_type, prompt_dict, tokenizer, device, human='<human>:', bot="<bot>:", model_max_length=None):
     # FIXME: prompt_dict unused currently
     if prompt_type in [PromptType.human_bot.name, PromptType.instruct_vicuna.name, PromptType.instruct_with_end.name]:
         if prompt_type == PromptType.human_bot.name:
@@ -67,7 +71,8 @@ def get_stopping(prompt_type, prompt_dict, tokenizer, device, human='<human>:',
         stop_words_ids = [x[1:] if y[0] == '\n' else x for x, y in zip(stop_words_ids, stop_words)]
         # build stopper
         stopping_criteria = StoppingCriteriaList(
-            [StoppingCriteriaSub(stops=stop_words_ids, encounters=encounters, device=device)])
+            [StoppingCriteriaSub(stops=stop_words_ids, encounters=encounters, device=device,
+                                 model_max_length=model_max_length)])
     else:
         stopping_criteria = StoppingCriteriaList()
     return stopping_criteria

utils.py

@@ -69,6 +69,25 @@ def ping():
         pass
 
 
+def ping_gpu():
+    try:
+        print('Ping_GPU: %s %s' % (str(datetime.now()), system_info()), flush=True)
+    except AttributeError:
+        # some programs wrap print and will fail with flush passed
+        pass
+    try:
+        ping_gpu_memory()
+    except Exception as e:
+        print('Ping_GPU memory failure: %s' % str(e), flush=True)
+
+
+def ping_gpu_memory():
+    from models.gpu_mem_track import MemTracker
+    gpu_tracker = MemTracker()  # define a GPU tracker
+    from torch.cuda import memory_summary
+    gpu_tracker.track()
+
+
 def get_torch_allocated():
     import torch
     return torch.cuda.memory_allocated()
@@ -98,27 +117,29 @@ def system_info():
             system['CPU_C/%s' % k] = v
 
     # https://github.com/gpuopenanalytics/pynvml/blob/master/help_query_gpu.txt
-    from pynvml.smi import nvidia_smi
-    nvsmi = nvidia_smi.getInstance()
-
-    gpu_power_dict = {'W_gpu%d' % i: x['power_readings']['power_draw'] for i, x in
-                      enumerate(nvsmi.DeviceQuery('power.draw')['gpu'])}
-    for k, v in gpu_power_dict.items():
-        system['GPU_W/%s' % k] = v
-
-    gpu_temp_dict = {'C_gpu%d' % i: x['temperature']['gpu_temp'] for i, x in
-                     enumerate(nvsmi.DeviceQuery('temperature.gpu')['gpu'])}
-    for k, v in gpu_temp_dict.items():
-        system['GPU_C/%s' % k] = v
-
-    gpu_memory_free_dict = {'MiB_gpu%d' % i: x['fb_memory_usage']['free'] for i, x in
-                            enumerate(nvsmi.DeviceQuery('memory.free')['gpu'])}
-    gpu_memory_total_dict = {'MiB_gpu%d' % i: x['fb_memory_usage']['total'] for i, x in
-                             enumerate(nvsmi.DeviceQuery('memory.total')['gpu'])}
-    gpu_memory_frac_dict = {k: gpu_memory_free_dict[k] / gpu_memory_total_dict[k] for k in gpu_memory_total_dict}
-    for k, v in gpu_memory_frac_dict.items():
-        system[f'GPU_M/%s' % k] = v
-
+    try:
+        from pynvml.smi import nvidia_smi
+        nvsmi = nvidia_smi.getInstance()
+
+        gpu_power_dict = {'W_gpu%d' % i: x['power_readings']['power_draw'] for i, x in
+                          enumerate(nvsmi.DeviceQuery('power.draw')['gpu'])}
+        for k, v in gpu_power_dict.items():
+            system['GPU_W/%s' % k] = v
+
+        gpu_temp_dict = {'C_gpu%d' % i: x['temperature']['gpu_temp'] for i, x in
+                         enumerate(nvsmi.DeviceQuery('temperature.gpu')['gpu'])}
+        for k, v in gpu_temp_dict.items():
+            system['GPU_C/%s' % k] = v
+
+        gpu_memory_free_dict = {'MiB_gpu%d' % i: x['fb_memory_usage']['free'] for i, x in
+                                enumerate(nvsmi.DeviceQuery('memory.free')['gpu'])}
+        gpu_memory_total_dict = {'MiB_gpu%d' % i: x['fb_memory_usage']['total'] for i, x in
+                                 enumerate(nvsmi.DeviceQuery('memory.total')['gpu'])}
+        gpu_memory_frac_dict = {k: gpu_memory_free_dict[k] / gpu_memory_total_dict[k] for k in gpu_memory_total_dict}
+        for k, v in gpu_memory_frac_dict.items():
+            system[f'GPU_M/%s' % k] = v
+    except ModuleNotFoundError:
+        pass
     system['hash'] = get_githash()
 
     return system
@@ -167,35 +188,39 @@ def _zip_data(root_dirs=None, zip_file=None, base_dir='./'):
     return zip_file, zip_file
 
 
-def save_generate_output(output=None, base_model=None, save_dir=None):
+def save_generate_output(prompt=None, output=None, base_model=None, save_dir=None, where_from='unknown where from',
+                         extra_dict={}):
     try:
-        return _save_generate_output(output=output, base_model=base_model, save_dir=save_dir)
+        return _save_generate_output(prompt=prompt, output=output, base_model=base_model, save_dir=save_dir,
+                                     where_from=where_from, extra_dict=extra_dict)
     except Exception as e:
         traceback.print_exc()
         print('Exception in saving: %s' % str(e))
 
 
-def _save_generate_output(output=None, base_model=None, save_dir=None):
+def _save_generate_output(prompt=None, output=None, base_model=None, save_dir=None, where_from='unknown where from',
+                          extra_dict={}):
     """
     Save conversation to .json, row by row.
     json_file_path is path to final JSON file. If not in ., then will attempt to make directories.
     Appends if file exists
     """
+    prompt = '<not set>' if prompt is None else prompt
+    output = '<not set>' if output is None else output
     assert save_dir, "save_dir must be provided"
     if os.path.exists(save_dir) and not os.path.isdir(save_dir):
         raise RuntimeError("save_dir already exists and is not a directory!")
     os.makedirs(save_dir, exist_ok=True)
     import json
-    if output[-10:] == '\n\n<human>:':
-        # remove trailing <human>:
-        output = output[:-10]
+    dict_to_save = dict(prompt=prompt, text=output, time=time.ctime(), base_model=base_model, where_from=where_from)
+    dict_to_save.update(extra_dict)
     with filelock.FileLock("save_dir.lock"):
         # lock logging in case have concurrency
         with open(os.path.join(save_dir, "history.json"), "a") as f:
             # just add [ at start, and ] at end, and have proper JSON dataset
             f.write(
                 "  " + json.dumps(
-                    dict(text=output, time=time.ctime(), base_model=base_model)
+                    dict_to_save
                 ) + ",\n"
             )
 
@@ -801,6 +826,7 @@ def get_kwargs(func, exclude_names=None, **kwargs):
 
 
 import pkg_resources
+
 have_faiss = False
 
 try:
@@ -828,7 +854,7 @@ def hash_file(file):
         BUF_SIZE = 65536  # lets read stuff in 64kb chunks!
 
         md5 = hashlib.md5()
-        #sha1 = hashlib.sha1()
+        # sha1 = hashlib.sha1()
 
         with open(file, 'rb') as f:
             while True:
@@ -836,7 +862,7 @@ def hash_file(file):
                 if not data:
                     break
                 md5.update(data)
-                #sha1.update(data)
+                # sha1.update(data)
     except BaseException as e:
         print("Cannot hash %s due to %s" % (file, str(e)))
         traceback.print_exc()
@@ -848,8 +874,55 @@ def start_faulthandler():
     # If hit server or any subprocess with signal SIGUSR1, it'll print out all threads stack trace, but wont't quit or coredump
     # If more than one fork tries to write at same time, then looks corrupted.
     import faulthandler
-    import signal
 
     # SIGUSR1 in h2oai/__init__.py as well
     faulthandler.enable()
-    faulthandler.register(signal.SIGUSR1)
+    if hasattr(faulthandler, 'register'):
+        # windows/mac
+        import signal
+        faulthandler.register(signal.SIGUSR1)
+
+
+def get_hf_server(inference_server):
+    inf_split = inference_server.split("    ")
+    assert len(inf_split) == 1 or len(inf_split) == 3
+    inference_server = inf_split[0]
+    if len(inf_split) == 3:
+        headers = {"authorization": "%s %s" % (inf_split[1], inf_split[2])}
+    else:
+        headers = None
+    return inference_server, headers
+
+
+class FakeTokenizer:
+    """
+    1) For keeping track of model_max_length
+    2) For when model doesn't directly expose tokenizer but need to count tokens
+    """
+
+    def __init__(self, model_max_length=2048, encoding_name="cl100k_base"):
+        # dont' push limit, since if using fake tokenizer, only estimate, and seen underestimates by order 250
+        self.model_max_length = model_max_length - 250
+        self.encoding_name = encoding_name
+        # The first time this runs, it will require an internet connection to download. Later runs won't need an internet connection.
+        import tiktoken
+        self.encoding = tiktoken.get_encoding(self.encoding_name)
+
+    def encode(self, x, *args, return_tensors="pt", **kwargs):
+        input_ids = self.encoding.encode(x, disallowed_special=())
+        if return_tensors == 'pt' and isinstance(input_ids, list):
+            import torch
+            input_ids = torch.tensor(input_ids)
+        return dict(input_ids=input_ids)
+
+    def decode(self, x, *args, **kwargs):
+        # input is input_ids[0] form
+        return self.encoding.decode(x)
+
+    def num_tokens_from_string(self, prompt: str) -> int:
+        """Returns the number of tokens in a text string."""
+        num_tokens = len(self.encoding.encode(prompt))
+        return num_tokens
+
+    def __call__(self, x, *args, **kwargs):
+        return self.encode(x, *args, **kwargs)

utils_langchain.py

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