Update with h2oGPT hash 221daabcabfa7f54b732394c15934a347da01079

client_test.py

@@ -48,6 +48,8 @@ import markdown  # pip install markdown
 import pytest
 from bs4 import BeautifulSoup  # pip install beautifulsoup4
 
+from enums import DocumentChoices
+
 debug = False
 
 os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
@@ -62,7 +64,10 @@ def get_client(serialize=True):
     return client
 
 
-def get_args(prompt, prompt_type, chat=False, stream_output=False, max_new_tokens=50, langchain_mode='Disabled'):
+def get_args(prompt, prompt_type, chat=False, stream_output=False,
+             max_new_tokens=50,
+             top_k_docs=3,
+             langchain_mode='Disabled'):
     from collections import OrderedDict
     kwargs = OrderedDict(instruction=prompt if chat else '',  # only for chat=True
                          iinput='',  # only for chat=True
@@ -71,6 +76,7 @@ def get_args(prompt, prompt_type, chat=False, stream_output=False, max_new_token
                          # but leave stream_output=False for simple input/output mode
                          stream_output=stream_output,
                          prompt_type=prompt_type,
+                         prompt_dict='',
                          temperature=0.1,
                          top_p=0.75,
                          top_k=40,
@@ -86,9 +92,13 @@ def get_args(prompt, prompt_type, chat=False, stream_output=False, max_new_token
                          instruction_nochat=prompt if not chat else '',
                          iinput_nochat='',  # only for chat=False
                          langchain_mode=langchain_mode,
-                         top_k_docs=4,
-                         document_choice=['All'],
+                         top_k_docs=top_k_docs,
+                         chunk=True,
+                         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=[]))
@@ -97,8 +107,8 @@ def get_args(prompt, prompt_type, chat=False, stream_output=False, max_new_token
 
 
 @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):
@@ -112,15 +122,110 @@ def run_client_nochat(prompt, prompt_type, max_new_tokens):
     )
     print("Raw client result: %s" % res, flush=True)
     res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
+                    response=md_to_text(res))
+    print(res_dict)
+    return res_dict, client
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic_api(prompt_type='human_bot'):
+    return run_client_nochat_api(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
+
+
+def run_client_nochat_api(prompt, prompt_type, max_new_tokens):
+    kwargs, args = get_args(prompt, prompt_type, chat=False, max_new_tokens=max_new_tokens)
+
+    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
+    client = get_client(serialize=True)
+    res = client.predict(
+        str(dict(kwargs)),
+        api_name=api_name,
+    )
+    print("Raw client result: %s" % res, flush=True)
+    res_dict = dict(prompt=kwargs['instruction_nochat'], iinput=kwargs['iinput_nochat'],
+                    response=md_to_text(ast.literal_eval(res)['response']),
+                    sources=ast.literal_eval(res)['sources'])
+    print(res_dict)
+    return res_dict, client
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic_api_lean(prompt_type='human_bot'):
+    return run_client_nochat_api_lean(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
+
+
+def run_client_nochat_api_lean(prompt, prompt_type, max_new_tokens):
+    kwargs = dict(instruction_nochat=prompt)
+
+    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
+    client = get_client(serialize=True)
+    res = client.predict(
+        str(dict(kwargs)),
+        api_name=api_name,
+    )
+    print("Raw client result: %s" % res, flush=True)
+    res_dict = dict(prompt=kwargs['instruction_nochat'],
+                    response=md_to_text(ast.literal_eval(res)['response']),
+                    sources=ast.literal_eval(res)['sources'])
+    print(res_dict)
+    return res_dict, client
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_basic_api_lean_morestuff(prompt_type='human_bot'):
+    return run_client_nochat_api_lean_morestuff(prompt='Who are you?', prompt_type=prompt_type, max_new_tokens=50)
+
+
+def run_client_nochat_api_lean_morestuff(prompt, prompt_type='human_bot', max_new_tokens=512):
+    kwargs = dict(
+        instruction='',
+        iinput='',
+        context='',
+        stream_output=False,
+        prompt_type=prompt_type,
+        temperature=0.1,
+        top_p=0.75,
+        top_k=40,
+        num_beams=1,
+        max_new_tokens=256,
+        min_new_tokens=0,
+        early_stopping=False,
+        max_time=20,
+        repetition_penalty=1.0,
+        num_return_sequences=1,
+        do_sample=True,
+        chat=False,
+        instruction_nochat=prompt,
+        iinput_nochat='',
+        langchain_mode='Disabled',
+        top_k_docs=4,
+        document_choice=['All'],
+    )
+
+    api_name = '/submit_nochat_api'  # NOTE: like submit_nochat but stable API for string dict passing
+    client = get_client(serialize=True)
+    res = client.predict(
+        str(dict(kwargs)),
+        api_name=api_name,
+    )
+    print("Raw client result: %s" % res, flush=True)
+    res_dict = dict(prompt=kwargs['instruction_nochat'],
                     response=md_to_text(ast.literal_eval(res)['response']),
                     sources=ast.literal_eval(res)['sources'])
     print(res_dict)
-    return res_dict
+    return res_dict, client
+
+
+@pytest.mark.skip(reason="For manual use against some server, no server launched")
+def test_client_chat(prompt_type='human_bot'):
+    return run_client_chat(prompt='Who are you?', prompt_type=prompt_type, stream_output=False, max_new_tokens=50,
+                           langchain_mode='Disabled')
 
 
 @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_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')
 
 
@@ -133,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]]
 
@@ -166,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
@@ -175,5 +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()
+    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

@@ -0,0 +1,84 @@
+from enum import Enum
+
+
+class PromptType(Enum):
+    custom = -1
+    plain = 0
+    instruct = 1
+    quality = 2
+    human_bot = 3
+    dai_faq = 4
+    summarize = 5
+    simple_instruct = 6
+    instruct_vicuna = 7
+    instruct_with_end = 8
+    human_bot_orig = 9
+    prompt_answer = 10
+    open_assistant = 11
+    wizard_lm = 12
+    wizard_mega = 13
+    instruct_vicuna2 = 14
+    instruct_vicuna3 = 15
+    wizard2 = 16
+    wizard3 = 17
+    instruct_simple = 18
+    wizard_vicuna = 19
+    openai = 20
+    openai_chat = 21
+    gptj = 22
+    prompt_answer_openllama = 23
+    vicuna11 = 24
+
+
+class DocumentChoices(Enum):
+    All_Relevant = 0
+    All_Relevant_Only_Sources = 1
+    Only_All_Sources = 2
+    Just_LLM = 3
+
+
+class LangChainMode(Enum):
+    """LangChain mode"""
+
+    DISABLED = "Disabled"
+    CHAT_LLM = "ChatLLM"
+    LLM = "LLM"
+    ALL = "All"
+    WIKI = "wiki"
+    WIKI_FULL = "wiki_full"
+    USER_DATA = "UserData"
+    MY_DATA = "MyData"
+    GITHUB_H2OGPT = "github h2oGPT"
+    H2O_DAI_DOCS = "DriverlessAI docs"
+
+
+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,8 +5,11 @@ 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
+from prompter import generate_prompt, prompt_types, PromptType
 from utils import get_githash, copy_code
 import torch
 
@@ -104,7 +107,6 @@ def train(
         save_total_limit: int = 3,
         add_eos_token: bool = False,
 ):
-
     if llama_flash_attn:
         # Need to call this before importing transformers.
         from llama_flash_attn_monkey_patch import replace_llama_attn_with_flash_attn
@@ -129,10 +131,12 @@ def train(
     if not output_dir:
         output_dir = f"{base_model.split('/')[-1]}.{data_path.replace('/', '')}.{num_epochs}_epochs.{get_githash() or 'nogit'}.{run_id}"
         if os.path.exists(output_dir) and not resume_from_checkpoint:
-            raise FileExistsError(f"output_dir {output_dir} based on run_id {run_id} already exists. Please pick a different run_id.")
+            raise FileExistsError(
+                f"output_dir {output_dir} based on run_id {run_id} already exists. Please pick a different run_id.")
     else:
         if os.path.exists(output_dir) and not resume_from_checkpoint:
-            raise FileExistsError(f"output_dir {output_dir} already exists. Please pick a different output_dir, or specify a run_id instead.")
+            raise FileExistsError(
+                f"output_dir {output_dir} already exists. Please pick a different output_dir, or specify a run_id instead.")
     device_map = "auto"
 
     if save_code:
@@ -181,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,
@@ -398,7 +402,8 @@ def train(
     if train_data_mix_in:
         train_data = concatenate_datasets([train_data, train_data_mix_in])
     log("Tokenizing %s training rows" % train_data.num_rows)
-    train_data = train_data.shuffle().map(generate_and_tokenize_prompt_fun, num_proc=os.cpu_count() // torch.cuda.device_count())
+    train_data = train_data.shuffle().map(generate_and_tokenize_prompt_fun,
+                                          num_proc=os.cpu_count() // torch.cuda.device_count())
     if drop_truncations:
         log("avoid keeping truncated cases to avoid contaminating model with truncation cases.  Original size: %s" % train_data.num_rows)
         prune_long_sequences_func = partial(prune_long_sequences, cutoff_len=cutoff_len)
@@ -413,7 +418,8 @@ def train(
 
     if valid_data:
         log("Tokenizing %s validation rows" % valid_data.num_rows)
-        valid_data = valid_data.shuffle().map(generate_and_tokenize_prompt_fun, num_proc=os.cpu_count() // torch.cuda.device_count())
+        valid_data = valid_data.shuffle().map(generate_and_tokenize_prompt_fun,
+                                              num_proc=os.cpu_count() // torch.cuda.device_count())
         val_set_size = len(valid_data)
     else:
         val_set_size = 0
@@ -468,7 +474,7 @@ def train(
     elif save_steps > eval_steps:
         # save steps must be round multiple of eval_steps
         save_steps0 = save_steps
-        save_steps = max(1, (save_steps//eval_steps)) * eval_steps
+        save_steps = max(1, (save_steps // eval_steps)) * eval_steps
         if save_steps0 != save_steps:
             log("Auto converted save_steps from %s to %s" % (save_steps0, save_steps))
 
@@ -478,21 +484,21 @@ def train(
         label_ids = eval_preds.label_ids
         predictions = eval_preds.predictions
 
-        #inputs = np.where(inputs != -100, inputs, tokenizer.pad_token_id)
-        #decoded_inputs = tokenizer.batch_decode(inputs, skip_special_tokens=True)
-        #decoded_inputs = [pred.strip() for pred in decoded_inputs]
+        # inputs = np.where(inputs != -100, inputs, tokenizer.pad_token_id)
+        # decoded_inputs = tokenizer.batch_decode(inputs, skip_special_tokens=True)
+        # decoded_inputs = [pred.strip() for pred in decoded_inputs]
 
         label_ids = np.where(label_ids != -100, label_ids, tokenizer.pad_token_id)
         # tokenizer behavior like generate time
         decoded_labels = tokenizer.batch_decode(label_ids, skip_special_tokens=True,
-                                                           clean_up_tokenization_spaces=True)
+                                                clean_up_tokenization_spaces=True)
         decoded_labels = [pred.strip() for pred in decoded_labels]
 
         predictions = np.argmax(predictions, -1)
         predictions = np.where(predictions != -100, predictions, tokenizer.pad_token_id)
         # tokenizer behavior like generate time
         decoded_predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True,
-                                                                  clean_up_tokenization_spaces=True)
+                                                     clean_up_tokenization_spaces=True)
         decoded_predictions = [pred.strip() for pred in decoded_predictions]
 
         result = {}
@@ -541,8 +547,8 @@ def train(
             load_best_model_at_end=True if val_set_size > 0 else False,
             ddp_find_unused_parameters=False if ddp else None,
             group_by_length=group_by_length,
-            #fsdp="shard_grad_op auto_wrap" if gpus > 1 and not ddp else None,
-            #fsdp_min_num_params=20000 if gpus > 1 and not ddp else None,
+            # fsdp="shard_grad_op auto_wrap" if gpus > 1 and not ddp else None,
+            # fsdp_min_num_params=20000 if gpus > 1 and not ddp else None,
             report_to='tensorboard' if not neptune_run else 'neptune',
         ),
         data_collator=transformers.DataCollatorForSeq2Seq(
@@ -553,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)
@@ -616,10 +615,12 @@ def generate_and_tokenize_prompt(data_point, prompt_type=None, train_on_inputs=F
     assert prompt_type is not None
     assert cutoff_len is not None
     assert tokenizer is not None
-    full_prompt, _, _, _ = generate_prompt(data_point, prompt_type, False, False)
+    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, False)
     tokenized_full_prompt = tokenize(full_prompt, tokenizer, cutoff_len, add_eos_token=add_eos_token)
     if not train_on_inputs:
-        user_prompt, _, _, _ = generate_prompt({**data_point, "output": ""}, prompt_type, False, False)
+        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:
@@ -638,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"""
@@ -665,6 +666,11 @@ NCCL_P2P_LEVEL=LOC WORLD_SIZE=7 CUDA_VISIBLE_DEVICES="0,1" torchrun --node_rank
 
     if os.environ.get("LOCAL_RANK") is None:
         # then not using torchrun, so can't do distributed, ensure CVD set
-        assert os.environ.get("CUDA_VISIBLE_DEVICES") is not None, "Run python script using: torchrun finetune.py OR set CUDA_VISIBLE_DEVICES to single GPU"
+        assert os.environ.get(
+            "CUDA_VISIBLE_DEVICES") is not None, "Run python script using: torchrun finetune.py OR set CUDA_VISIBLE_DEVICES to single GPU"
 
     fire.Fire(train)
+
+
+if __name__ == "__main__":
+    entrypoint_main()

gpt4all_llm.py

@@ -1,23 +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:
-
-    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):
@@ -73,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
@@ -93,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,
@@ -113,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)
@@ -136,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()
@@ -155,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(
@@ -171,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)
 
 
@@ -185,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()
@@ -236,9 +248,12 @@ 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
+        # still have to avoid crazy sizes, else hit llama_tokenize: too many tokens -- might still hit, not fatal
+        prompt = prompt[-self.n_ctx * 4:]
         prompt_tokens = self.client.tokenize(b" " + prompt.encode("utf-8"))
         num_prompt_tokens = len(prompt_tokens)
         if num_prompt_tokens > self.n_ctx:
@@ -250,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,27 +1,34 @@
+import ast
 import glob
 import inspect
 import os
 import pathlib
 import pickle
-import queue
 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 prompter import non_hf_types
+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
+    get_device, ProgressParallel, remove, hash_file, clear_torch_cache, NullContext, get_hf_server, FakeTokenizer
+from utils_langchain import StreamingGradioCallbackHandler
 
 import_matplotlib()
 
@@ -36,19 +43,22 @@ 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
 
 
-def get_db(sources, use_openai_embedding=False, db_type='faiss', persist_directory="db_dir", langchain_mode='notset',
+def get_db(sources, use_openai_embedding=False, db_type='faiss',
+           persist_directory="db_dir", load_db_if_exists=True,
+           langchain_mode='notset',
            collection_name=None,
            hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2"):
     if not sources:
         return None
+
     # get embedding model
     embedding = get_embedding(use_openai_embedding, hf_embedding_model=hf_embedding_model)
     assert collection_name is not None or langchain_mode != 'notset'
@@ -59,29 +69,41 @@ def get_db(sources, use_openai_embedding=False, db_type='faiss', persist_directo
     if db_type == 'faiss':
         from langchain.vectorstores import FAISS
         db = FAISS.from_documents(sources, embedding)
-
     elif db_type == 'weaviate':
         import weaviate
         from weaviate.embedded import EmbeddedOptions
         from langchain.vectorstores import Weaviate
 
-        # TODO: add support for connecting via docker compose
-        client = weaviate.Client(
-            embedded_options=EmbeddedOptions()
-        )
+        if os.getenv('WEAVIATE_URL', None):
+            client = _create_local_weaviate_client()
+        else:
+            client = weaviate.Client(
+                embedded_options=EmbeddedOptions()
+            )
         index_name = collection_name.capitalize()
         db = Weaviate.from_documents(documents=sources, embedding=embedding, client=client, by_text=False,
                                      index_name=index_name)
-
     elif db_type == 'chroma':
         assert persist_directory is not None
         os.makedirs(persist_directory, exist_ok=True)
-        db = Chroma.from_documents(documents=sources,
-                                   embedding=embedding,
-                                   persist_directory=persist_directory,
-                                   collection_name=collection_name,
-                                   anonymized_telemetry=False)
-        db.persist()
+
+        # see if already actually have persistent db, and deal with possible changes in embedding
+        db = get_existing_db(None, persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
+                             hf_embedding_model, verbose=False)
+        if db is None:
+            db = Chroma.from_documents(documents=sources,
+                                       embedding=embedding,
+                                       persist_directory=persist_directory,
+                                       collection_name=collection_name,
+                                       anonymized_telemetry=False)
+            db.persist()
+            clear_embedding(db)
+            save_embed(db, use_openai_embedding, hf_embedding_model)
+        else:
+            # then just add
+            db, num_new_sources, new_sources_metadata = add_to_db(db, sources, db_type=db_type,
+                                                                  use_openai_embedding=use_openai_embedding,
+                                                                  hf_embedding_model=hf_embedding_model)
     else:
         raise RuntimeError("No such db_type=%s" % db_type)
 
@@ -104,7 +126,10 @@ def _get_unique_sources_in_weaviate(db):
 
 def add_to_db(db, sources, db_type='faiss',
               avoid_dup_by_file=False,
-              avoid_dup_by_content=True):
+              avoid_dup_by_content=True,
+              use_openai_embedding=False,
+              hf_embedding_model=None):
+    assert hf_embedding_model is not None
     num_new_sources = len(sources)
     if not sources:
         return db, num_new_sources, []
@@ -120,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:
@@ -135,11 +160,15 @@ 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(
                 dup_metadata_files), flush=True)
-            client_collection = db._client.get_collection(name=db._collection.name)
+            client_collection = db._client.get_collection(name=db._collection.name,
+                                                          embedding_function=db._collection._embedding_function)
             for dup_file in dup_metadata_files:
                 dup_file_meta = dict(source=dup_file)
                 try:
@@ -151,6 +180,8 @@ def add_to_db(db, sources, db_type='faiss',
             return db, num_new_sources, []
         db.add_documents(documents=sources)
         db.persist()
+        clear_embedding(db)
+        save_embed(db, use_openai_embedding, hf_embedding_model)
     else:
         raise RuntimeError("No such db_type=%s" % db_type)
 
@@ -165,10 +196,13 @@ def create_or_update_db(db_type, persist_directory, collection_name,
         import weaviate
         from weaviate.embedded import EmbeddedOptions
 
-        # TODO: add support for connecting via docker compose
-        client = weaviate.Client(
-            embedded_options=EmbeddedOptions()
-        )
+        if os.getenv('WEAVIATE_URL', None):
+            client = _create_local_weaviate_client()
+        else:
+            client = weaviate.Client(
+                embedded_options=EmbeddedOptions()
+            )
+
         index_name = collection_name.replace(' ', '_').capitalize()
         if client.schema.exists(index_name) and not add_if_exists:
             client.schema.delete_class(index_name)
@@ -205,14 +239,20 @@ 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
 
         device, torch_dtype, context_class = get_device_dtype()
         model_kwargs = dict(device=device)
-        embedding = HuggingFaceEmbeddings(model_name=hf_embedding_model, model_kwargs=model_kwargs)
+        if 'instructor' in hf_embedding_model:
+            encode_kwargs = {'normalize_embeddings': True}
+            embedding = HuggingFaceInstructEmbeddings(model_name=hf_embedding_model,
+                                                      model_kwargs=model_kwargs,
+                                                      encode_kwargs=encode_kwargs)
+        else:
+            embedding = HuggingFaceEmbeddings(model_name=hf_embedding_model, model_kwargs=model_kwargs)
     return embedding
 
 
@@ -226,63 +266,481 @@ 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,
-            max_new_tokens=256,
+"""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,
-            repetition_penalty=1.0,
             top_k=40,
             top_p=0.7,
+            num_beams=1,
+            max_new_tokens=256,
+            min_new_tokens=1,
+            early_stopping=False,
+            max_time=180,
+            repetition_penalty=1.0,
+            num_return_sequences=1,
             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 model_name is None
             assert tokenizer is None
             prompt_type = 'human_bot'
-            model_name = 'h2oai/h2ogpt-oasst1-512-12b'
-            # model_name = 'h2oai/h2ogpt-oig-oasst1-512-6_9b'
-            # model_name = 'h2oai/h2ogpt-oasst1-512-20b'
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            device, torch_dtype, context_class = get_device_dtype()
-
-            with context_class(device):
-                load_8bit = True
-                # FIXME: for now not to spread across hetero GPUs
-                # device_map={"": 0} if load_8bit and device == 'cuda' else "auto"
-                device_map = {"": 0} if device == 'cuda' else "auto"
-                model = AutoModelForCausalLM.from_pretrained(model_name,
-                                                             device_map=device_map,
-                                                             torch_dtype=torch_dtype,
-                                                             load_in_8bit=load_8bit)
+            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(max_new_tokens=max_new_tokens,
+        gen_kwargs = dict(do_sample=do_sample,
+                          temperature=temperature,
+                          top_k=top_k,
+                          top_p=top_p,
+                          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,
                           return_full_text=True,
-                          early_stopping=False,
-                          handle_long_generation='hole')
+                          handle_long_generation=None)
+        assert len(set(gen_hyper).difference(gen_kwargs.keys())) == 0
 
         if stream_output:
             skip_prompt = False
@@ -297,10 +755,12 @@ def get_llm(use_openai_model=False, model_name=None, model=None,
         pipe = H2OTextGenerationPipeline(model=model, use_prompter=True,
                                          prompter=prompter,
                                          prompt_type=prompt_type,
-                                         sanitize_bot_response=True,
+                                         prompt_dict=prompt_dict,
+                                         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,
@@ -345,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,
@@ -467,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",
@@ -484,12 +959,13 @@ 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]
 
 
-def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, chunk=True, chunk_size=512,
+def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False,
+                chunk=True, chunk_size=512,
                 is_url=False, is_txt=False,
                 enable_captions=True,
                 captions_model=None,
@@ -525,9 +1001,25 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
             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]
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+        docs1 = clean_doc(docs1)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif is_txt:
         base_path = "user_paste"
         source_file = os.path.join(base_path, "_%s" % str(uuid.uuid4())[:10])
@@ -536,44 +1028,49 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
             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_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)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif file.lower().endswith('.odt'):
         docs1 = UnstructuredODTLoader(file_path=file).load()
         add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif file.lower().endswith('pptx') or file.lower().endswith('ppt'):
         docs1 = UnstructuredPowerPointLoader(file_path=file).load()
         add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+        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_size=chunk_size)
+        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()
         add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif file.lower().endswith('.md'):
         docs1 = UnstructuredMarkdownLoader(file).load()
         add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_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)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif file.lower().endswith('.epub'):
         docs1 = UnstructuredEPubLoader(file).load()
         add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif file.lower().endswith('.jpeg') or file.lower().endswith('.jpg') or file.lower().endswith('.png'):
         docs1 = []
         if have_tesseract and enable_ocr:
@@ -603,7 +1100,7 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
                 doci.metadata['source'] = doci.metadata['image_path']
                 doci.metadata['hash'] = hash_file(doci.metadata['source'])
             if docs1:
-                doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+                doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif file.lower().endswith('.msg'):
         raise RuntimeError("Not supported, GPL3 license")
         # docs1 = OutlookMessageLoader(file).load()
@@ -612,14 +1109,14 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
         try:
             docs1 = UnstructuredEmailLoader(file).load()
             add_meta(docs1, file)
-            doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+            doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
         except ValueError as e:
             if 'text/html content not found in email' in str(e):
                 # e.g. plain/text dict key exists, but not
                 # doc1 = TextLoader(file, encoding="utf8").load()
                 docs1 = UnstructuredEmailLoader(file, content_source="text/plain").load()
                 add_meta(docs1, file)
-                doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+                doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
             else:
                 raise
     # elif file.lower().endswith('.gcsdir'):
@@ -630,6 +1127,7 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
         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
@@ -638,11 +1136,19 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
         if have_pymupdf and pdf_class_name == 'PyMuPDFParser':
             # GPL, only use if installed
             from langchain.document_loaders import PyMuPDFLoader
-            doc1 = PyMuPDFLoader(file).load_and_split()
+            # 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
-            doc1 = PyPDFLoader(file).load_and_split()
+            # 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()
@@ -650,6 +1156,7 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
     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)
@@ -657,7 +1164,7 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
         with open(file, "r") as f:
             docs1 = UnstructuredURLLoader(urls=f.readlines()).load()
         add_meta(docs1, file)
-        doc1 = chunk_sources(docs1, chunk_size=chunk_size)
+        doc1 = chunk_sources(docs1, chunk=chunk, chunk_size=chunk_size)
     elif file.lower().endswith('.zip'):
         with zipfile.ZipFile(file, 'r') as zip_ref:
             # don't put into temporary path, since want to keep references to docs inside zip
@@ -672,12 +1179,12 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
     # if list of length one, don't trust and chunk it
     if not isinstance(doc1, list):
         if chunk:
-            docs = chunk_sources([doc1], chunk_size=chunk_size)
+            docs = chunk_sources([doc1], chunk=chunk, chunk_size=chunk_size)
         else:
             docs = [doc1]
     elif isinstance(doc1, list) and len(doc1) == 1:
         if chunk:
-            docs = chunk_sources(doc1, chunk_size=chunk_size)
+            docs = chunk_sources(doc1, chunk=chunk, chunk_size=chunk_size)
         else:
             docs = doc1
     else:
@@ -687,7 +1194,8 @@ def file_to_doc(file, base_path=None, verbose=False, fail_any_exception=False, c
     return docs
 
 
-def path_to_doc1(file, verbose=False, fail_any_exception=False, return_file=True, chunk=True, chunk_size=512,
+def path_to_doc1(file, verbose=False, fail_any_exception=False, return_file=True,
+                 chunk=True, chunk_size=512,
                  is_url=False, is_txt=False,
                  enable_captions=True,
                  captions_model=None,
@@ -739,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()
@@ -756,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:
@@ -861,12 +1373,12 @@ def prep_langchain(persist_directory,
 
     if db_dir_exists and user_path is None:
         print("Prep: persist_directory=%s exists, using" % persist_directory, flush=True)
-        db = get_existing_db(persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
+        db = get_existing_db(None, persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
                              hf_embedding_model)
     else:
         if db_dir_exists and user_path is not None:
             print("Prep: persist_directory=%s exists, user_path=%s passed, adding any changed or new documents" % (
-            persist_directory, user_path), flush=True)
+                persist_directory, user_path), flush=True)
         elif not db_dir_exists:
             print("Prep: persist_directory=%s does not exist, regenerating" % persist_directory, flush=True)
         db = None
@@ -912,24 +1424,78 @@ class FakeConsumer(object):
 posthog.Consumer = FakeConsumer
 
 
-def get_existing_db(persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
-                    hf_embedding_model):
+def check_update_chroma_embedding(db, use_openai_embedding, hf_embedding_model, langchain_mode):
+    changed_db = False
+    if load_embed(db) != (use_openai_embedding, hf_embedding_model):
+        print("Detected new embedding, updating db: %s" % langchain_mode, flush=True)
+        # handle embedding changes
+        db_get = 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
+        persist_directory = db._persist_directory
+        shutil.move(persist_directory, persist_directory + "_" + str(uuid.uuid4()) + ".bak")
+        db_type = 'chroma'
+        load_db_if_exists = False
+        db = get_db(sources, use_openai_embedding=use_openai_embedding, db_type=db_type,
+                    persist_directory=persist_directory, load_db_if_exists=load_db_if_exists,
+                    langchain_mode=langchain_mode,
+                    collection_name=None,
+                    hf_embedding_model=hf_embedding_model)
+        if False:
+            # below doesn't work if db already in memory, so have to switch to new db as above
+            # upsert does new embedding, but if index already in memory, complains about size mismatch etc.
+            client_collection = db._client.get_collection(name=db._collection.name,
+                                                          embedding_function=db._collection._embedding_function)
+            client_collection.upsert(ids=db_get['ids'], metadatas=db_get['metadatas'], documents=db_get['documents'])
+        changed_db = True
+        print("Done updating db for new embedding: %s" % langchain_mode, flush=True)
+
+    return db, changed_db
+
+
+def get_existing_db(db, persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
+                    hf_embedding_model, verbose=False, check_embedding=True):
     if load_db_if_exists and db_type == 'chroma' and os.path.isdir(persist_directory) and os.path.isdir(
             os.path.join(persist_directory, 'index')):
-        print("DO Loading db: %s" % langchain_mode, flush=True)
-        embedding = get_embedding(use_openai_embedding, hf_embedding_model=hf_embedding_model)
-        from chromadb.config import Settings
-        client_settings = Settings(anonymized_telemetry=False,
-                                   chroma_db_impl="duckdb+parquet",
-                                   persist_directory=persist_directory)
-        db = Chroma(persist_directory=persist_directory, embedding_function=embedding,
-                    collection_name=langchain_mode.replace(' ', '_'),
-                    client_settings=client_settings)
-        print("DONE Loading db: %s" % langchain_mode, flush=True)
+        if db is None:
+            if verbose:
+                print("DO Loading db: %s" % langchain_mode, flush=True)
+            embedding = get_embedding(use_openai_embedding, hf_embedding_model=hf_embedding_model)
+            from chromadb.config import Settings
+            client_settings = Settings(anonymized_telemetry=False,
+                                       chroma_db_impl="duckdb+parquet",
+                                       persist_directory=persist_directory)
+            db = Chroma(persist_directory=persist_directory, embedding_function=embedding,
+                        collection_name=langchain_mode.replace(' ', '_'),
+                        client_settings=client_settings)
+            if verbose:
+                print("DONE Loading db: %s" % langchain_mode, flush=True)
+        else:
+            if verbose:
+                print("USING already-loaded db: %s" % langchain_mode, flush=True)
+        if check_embedding:
+            db_trial, changed_db = check_update_chroma_embedding(db, use_openai_embedding, hf_embedding_model,
+                                                                 langchain_mode)
+            if changed_db:
+                db = db_trial
+                # only call persist if really changed db, else takes too long for large db
+                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()
+
+
 def make_db(**langchain_kwargs):
     func_names = list(inspect.signature(_make_db).parameters)
     missing_kwargs = [x for x in func_names if x not in langchain_kwargs]
@@ -945,9 +1511,33 @@ def make_db(**langchain_kwargs):
     return _make_db(**langchain_kwargs)
 
 
+def save_embed(db, use_openai_embedding, hf_embedding_model):
+    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
+
+
+def load_embed(db):
+    embed_info_file = os.path.join(db._persist_directory, 'embed_info')
+    if os.path.isfile(embed_info_file):
+        with open(embed_info_file, 'rb') as f:
+            use_openai_embedding, hf_embedding_model = pickle.load(f)
+    else:
+        # migration, assume defaults
+        use_openai_embedding, hf_embedding_model = False, "sentence-transformers/all-MiniLM-L6-v2"
+    return use_openai_embedding, hf_embedding_model
+
+
+def get_persist_directory(langchain_mode):
+    return 'db_dir_%s' % langchain_mode  # single place, no special names for each case
+
+
 def _make_db(use_openai_embedding=False,
              hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
-             first_para=False, text_limit=None, chunk=False, chunk_size=1024,
+             first_para=False, text_limit=None,
+             chunk=True, chunk_size=512,
              langchain_mode=None,
              user_path=None,
              db_type='faiss',
@@ -955,19 +1545,13 @@ def _make_db(use_openai_embedding=False,
              db=None,
              n_jobs=-1,
              verbose=False):
-    persist_directory = 'db_dir_%s' % langchain_mode  # single place, no special names for each case
-    if not db and load_db_if_exists and db_type == 'chroma' and os.path.isdir(persist_directory) and os.path.isdir(
-            os.path.join(persist_directory, 'index')):
-        assert langchain_mode not in ['MyData'], "Should not load MyData db this way"
-        print("Loading existing db", flush=True)
-        embedding = get_embedding(use_openai_embedding, hf_embedding_model=hf_embedding_model)
-        from chromadb.config import Settings
-        client_settings = Settings(anonymized_telemetry=False,
-                                   chroma_db_impl="duckdb+parquet",
-                                   persist_directory=persist_directory)
-        db = Chroma(persist_directory=persist_directory, embedding_function=embedding,
-                    collection_name=langchain_mode.replace(' ', '_'),
-                    client_settings=client_settings)
+    persist_directory = get_persist_directory(langchain_mode)
+    # see if can get persistent chroma db
+    db_trial = get_existing_db(db, persist_directory, load_db_if_exists, db_type, use_openai_embedding, langchain_mode,
+                               hf_embedding_model, verbose=verbose)
+    if db_trial is not None:
+        db = db_trial
+
     sources = []
     if not db and langchain_mode not in ['MyData'] or \
             user_path is not None and \
@@ -992,24 +1576,24 @@ def _make_db(use_openai_embedding=False,
             sources1 = get_all_documents(small_test=small_test, n_jobs=os.cpu_count() // 2)
             print("Got new wiki", flush=True)
             if chunk:
-                sources1 = chunk_sources(sources1, chunk_size=chunk_size)
+                sources1 = chunk_sources(sources1, chunk=chunk, chunk_size=chunk_size)
                 print("Chunked new wiki", flush=True)
             sources.extend(sources1)
         if langchain_mode in ['wiki', 'All', "'All'"]:
             sources1 = get_wiki_sources(first_para=first_para, text_limit=text_limit)
             if chunk:
-                sources1 = chunk_sources(sources1, chunk_size=chunk_size)
+                sources1 = chunk_sources(sources1, chunk=chunk, chunk_size=chunk_size)
             sources.extend(sources1)
         if langchain_mode in ['github h2oGPT', 'All', "'All'"]:
             # sources = get_github_docs("dagster-io", "dagster")
             sources1 = get_github_docs("h2oai", "h2ogpt")
             # FIXME: always chunk for now
-            sources1 = chunk_sources(sources1, chunk_size=chunk_size)
+            sources1 = chunk_sources(sources1, chunk=chunk, chunk_size=chunk_size)
             sources.extend(sources1)
         if langchain_mode in ['DriverlessAI docs', 'All', "'All'"]:
             sources1 = get_dai_docs(from_hf=True)
             if chunk and False:  # FIXME: DAI docs are already chunked well, should only chunk more if over limit
-                sources1 = chunk_sources(sources1, chunk_size=chunk_size)
+                sources1 = chunk_sources(sources1, chunk=chunk, chunk_size=chunk_size)
             sources.extend(sources1)
         if langchain_mode in ['All', 'UserData']:
             if user_path:
@@ -1023,6 +1607,8 @@ def _make_db(use_openai_embedding=False,
                     existing_files = []
                     existing_hash_ids = []
                 # chunk internally for speed over multiple docs
+                # FIXME: If first had old Hash=None and switch embeddings,
+                #  then re-embed, and then hit here and reload so have hash, and then re-embed.
                 sources1 = path_to_docs(user_path, n_jobs=n_jobs, chunk=chunk, chunk_size=chunk_size,
                                         existing_files=existing_files, existing_hash_ids=existing_hash_ids)
                 new_metadata_sources = set([x.metadata['source'] for x in sources1])
@@ -1066,7 +1652,9 @@ def _make_db(use_openai_embedding=False,
         new_sources_metadata = [x.metadata for x in sources]
     elif user_path is not None and langchain_mode in ['UserData']:
         print("Existing db, potentially adding %s sources from user_path=%s" % (len(sources), user_path), flush=True)
-        db, num_new_sources, new_sources_metadata = add_to_db(db, sources, db_type=db_type)
+        db, num_new_sources, new_sources_metadata = add_to_db(db, sources, db_type=db_type,
+                                                              use_openai_embedding=use_openai_embedding,
+                                                              hf_embedding_model=hf_embedding_model)
         print("Existing db, added %s new sources from user_path=%s" % (num_new_sources, user_path), flush=True)
     else:
         new_sources_metadata = [x.metadata for x in sources]
@@ -1074,63 +1662,140 @@ 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
     # only keep actual used
     kwargs = {k: v for k, v in kwargs.items() if k in func_names}
-    return _run_qa_db(**kwargs)
+    try:
+        return _run_qa_db(**kwargs)
+    finally:
+        clear_torch_cache()
 
 
 def _run_qa_db(query=None,
                use_openai_model=False, use_openai_embedding=False,
-               first_para=False, text_limit=None, k=4, chunk=False, chunk_size=1024,
+               first_para=False, text_limit=None, top_k_docs=4, chunk=True, chunk_size=512,
                user_path=None,
                detect_user_path_changes_every_query=False,
                db_type='faiss',
-               model_name=None, model=None, tokenizer=None,
+               model_name=None, model=None, tokenizer=None, inference_server=None,
                hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
                stream_output=False,
                prompter=None,
                prompt_type=None,
+               prompt_dict=None,
                answer_with_sources=True,
                cut_distanct=1.1,
-               sanitize_bot_response=True,
+               sanitize_bot_response=False,
                show_rank=False,
                load_db_if_exists=False,
                db=None,
-               max_new_tokens=256,
+               do_sample=False,
                temperature=0.1,
-               repetition_penalty=1.0,
                top_k=40,
                top_p=0.7,
+               num_beams=1,
+               max_new_tokens=256,
+               min_new_tokens=1,
+               early_stopping=False,
+               max_time=180,
+               repetition_penalty=1.0,
+               num_return_sequences=1,
                langchain_mode=None,
-               document_choice=['All'],
+               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:
@@ -1149,39 +1814,63 @@ 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:
         prompt_type = prompter.prompt_type
+        prompt_dict = prompter.prompt_dict
     if model is not None:
         assert prompt_type is not None
+        if prompt_type == PromptType.custom.name:
+            assert prompt_dict is not None  # should at least be {} or ''
+        else:
+            prompt_dict = ''
+    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,
-                                                         max_new_tokens=max_new_tokens,
+                                                         do_sample=do_sample,
                                                          temperature=temperature,
-                                                         repetition_penalty=repetition_penalty,
                                                          top_k=top_k,
                                                          top_p=top_p,
+                                                         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,
                                                          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
 
+    if isinstance(document_choice, str):
+        # support string as well
+        document_choice = [document_choice]
+    # get first DocumentChoices as command to use, ignore others
+    doc_choices_set = set([x.name for x in list(DocumentChoices)])
+    cmd = [x for x in document_choice if x in doc_choices_set]
+    cmd = None if len(cmd) == 0 else cmd[0]
+    # now have cmd, filter out for only docs
+    document_choice = [x for x in document_choice if x not in doc_choices_set]
+
     func_names = list(inspect.signature(get_similarity_chain).parameters)
     sim_kwargs = {k: v for k, v in locals().items() if k in func_names}
     missing_kwargs = [x for x in func_names if x not in sim_kwargs]
     assert not missing_kwargs, "Missing: %s" % missing_kwargs
     docs, chain, scores, use_context = get_similarity_chain(**sim_kwargs)
-    if len(document_choice) > 0 and document_choice[0] == 'Only':
+    if cmd in [DocumentChoices.All_Relevant_Only_Sources.name, DocumentChoices.Only_All_Sources.name]:
         formatted_doc_chunks = '\n\n'.join([get_url(x) + '\n\n' + x.page_content for x in docs])
         yield formatted_doc_chunks, ''
         return
@@ -1189,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']
@@ -1239,22 +1934,31 @@ def _run_qa_db(query=None,
 
 def get_similarity_chain(query=None,
                          use_openai_model=False, use_openai_embedding=False,
-                         first_para=False, text_limit=None, k=4, chunk=False, chunk_size=1024,
+                         first_para=False, text_limit=None, top_k_docs=4, chunk=True, chunk_size=512,
                          user_path=None,
                          detect_user_path_changes_every_query=False,
                          db_type='faiss',
                          model_name=None,
+                         inference_server='',
                          hf_embedding_model="sentence-transformers/all-MiniLM-L6-v2",
                          prompt_type=None,
+                         prompt_dict=None,
                          cut_distanct=1.1,
                          load_db_if_exists=False,
                          db=None,
                          langchain_mode=None,
-                         document_choice=['All'],
+                         document_choice=[DocumentChoices.All_Relevant.name],
                          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:
@@ -1266,10 +1970,14 @@ def get_similarity_chain(query=None,
         use_context = True
 
     # https://github.com/hwchase17/langchain/issues/1946
-    # FIXME: Seems to way to get size of chroma db to limit k to avoid
+    # 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 k  # k=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:
@@ -1279,7 +1987,8 @@ def get_similarity_chain(query=None,
         user_path = None
     db, num_new_sources, new_sources_metadata = make_db(use_openai_embedding=use_openai_embedding,
                                                         hf_embedding_model=hf_embedding_model,
-                                                        first_para=first_para, text_limit=text_limit, chunk=chunk,
+                                                        first_para=first_para, text_limit=text_limit,
+                                                        chunk=chunk,
                                                         chunk_size=chunk_size,
                                                         langchain_mode=langchain_mode,
                                                         user_path=user_path,
@@ -1289,37 +1998,133 @@ 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 isinstance(document_choice, str):
-            # support string as well
-            document_choice = [document_choice]
-        if not isinstance(db, Chroma) or \
-                len(document_choice) == 0 or \
-                len(document_choice) <= 1 and document_choice[0] == 'All':
-            # treat empty list as All for now, not 'None'
-            filter_kwargs = {}
-        elif len(document_choice) > 0 and document_choice[0] == 'Only':
-            # Only means All docs, but only will return sources, not LLM response
+        if not isinstance(db, Chroma):
+            # only chroma supports filtering
             filter_kwargs = {}
         else:
+            # if here then some cmd + documents selected or just documents selected
             if len(document_choice) >= 2:
                 or_filter = [{"source": {"$eq": x}} for x in document_choice]
                 filter_kwargs = dict(filter={"$or": or_filter})
-            elif len(document_choice) > 0:
+            elif len(document_choice) == 1:
+                # degenerate UX bug in chroma
                 one_filter = [{"source": {"$eq": x}} for x in document_choice][0]
                 filter_kwargs = dict(filter=one_filter)
             else:
+                # shouldn't reach
                 filter_kwargs = {}
-            if len(document_choice) == 1 and document_choice[0] == 'None':
-                k_db = 1
-                k = 0
-        docs_with_score = db.similarity_search_with_score(query, k=k_db, **filter_kwargs)[:k]
-        # cut off so no high distance docs/sources considered
-        docs = [x[0] for x in docs_with_score if x[1] < cut_distanct]
-        scores = [x[1] for x in docs_with_score if x[1] < cut_distanct]
-        if len(scores) > 0 and verbose:
-            print("Distance: min: %s max: %s mean: %s median: %s" %
-                  (scores[0], scores[-1], np.mean(scores), np.median(scores)), flush=True)
+        if cmd == DocumentChoices.Just_LLM.name:
+            docs = []
+            scores = []
+        elif cmd == DocumentChoices.Only_All_Sources.name:
+            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_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:
+            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]
+            if len(scores) > 0 and verbose:
+                print("Distance: min: %s max: %s mean: %s median: %s" %
+                      (scores[0], scores[-1], np.mean(scores), np.median(scores)), flush=True)
     else:
         docs = []
         scores = []
@@ -1328,7 +2133,7 @@ def get_similarity_chain(query=None,
         # if HF type and have no docs, can bail out
         return docs, None, [], False
 
-    if len(document_choice) > 0 and document_choice[0] == 'Only':
+    if cmd in [DocumentChoices.All_Relevant_Only_Sources.name, DocumentChoices.Only_All_Sources.name]:
         # no LLM use
         return docs, None, [], False
 
@@ -1348,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"],
@@ -1420,15 +2217,32 @@ def get_sources_answer(query, answer, scores, show_rank, answer_with_sources, ve
     return ret, extra
 
 
-def chunk_sources(sources, chunk_size=1024):
-    source_chunks = []
-    # Below for known separator
-    # splitter = CharacterTextSplitter(separator=" ", chunk_size=chunk_size, chunk_overlap=0)
-    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=0)
-    for source in sources:
-        # print(source.metadata['source'], flush=True)
-        for chunky in splitter.split_text(source.page_content):
-            source_chunks.append(Document(page_content=chunky, metadata=source.metadata))
+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
+    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
 
 
@@ -1439,6 +2253,8 @@ def get_db_from_hf(dest=".", db_dir='db_dir_DriverlessAI_docs.zip'):
     path_to_zip_file = hf_hub_download('h2oai/db_dirs', db_dir, token=token, repo_type='dataset')
     import zipfile
     with zipfile.ZipFile(path_to_zip_file, 'r') as zip_ref:
+        persist_directory = os.path.dirname(zip_ref.namelist()[0])
+        remove(persist_directory)
         zip_ref.extractall(dest)
     return path_to_zip_file
 
@@ -1467,5 +2283,28 @@ def get_some_dbs_from_hf(dest='.', db_zips=None):
             assert os.path.isdir(os.path.join(dest, dir_expected, 'index')), "Missing index in %s" % dir_expected
 
 
+def _create_local_weaviate_client():
+    WEAVIATE_URL = os.getenv('WEAVIATE_URL', "http://localhost:8080")
+    WEAVIATE_USERNAME = os.getenv('WEAVIATE_USERNAME')
+    WEAVIATE_PASSWORD = os.getenv('WEAVIATE_PASSWORD')
+    WEAVIATE_SCOPE = os.getenv('WEAVIATE_SCOPE', "offline_access")
+
+    resource_owner_config = None
+    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
+
+
 if __name__ == '__main__':
     pass

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

@@ -1,14 +1,17 @@
+import os
+
 from transformers import TextGenerationPipeline
 from transformers.pipelines.text_generation import ReturnType
 
 from stopping import get_stopping
-from prompter import Prompter
+from prompter import Prompter, PromptType
 
 
 class H2OTextGenerationPipeline(TextGenerationPipeline):
     def __init__(self, *args, debug=False, chat=False, stream_output=False,
-                 sanitize_bot_response=True,
-                 use_prompter=True, prompter=None, prompt_type=None,
+                 sanitize_bot_response=False,
+                 use_prompter=True, prompter=None,
+                 prompt_type=None, prompt_dict=None,
                  max_input_tokens=2048 - 256, **kwargs):
         """
         HF-like pipeline, but handle instruction prompting and stopping (for some models)
@@ -21,6 +24,7 @@ class H2OTextGenerationPipeline(TextGenerationPipeline):
         :param prompter: prompter, can pass if have already
         :param prompt_type: prompt_type, e.g. human_bot.  See prompt_type to model mapping in from prompter.py.
                             If use_prompter, then will make prompter and use it.
+        :param prompt_dict: dict of get_prompt(, return_dict=True) for prompt_type=custom
         :param max_input_tokens:
         :param kwargs:
         """
@@ -28,12 +32,14 @@ class H2OTextGenerationPipeline(TextGenerationPipeline):
         self.prompt_text = None
         self.use_prompter = use_prompter
         self.prompt_type = prompt_type
+        self.prompt_dict = prompt_dict
         self.prompter = prompter
         if self.use_prompter:
             if self.prompter is not None:
                 assert self.prompter.prompt_type is not None
             else:
-                self.prompter = Prompter(self.prompt_type, debug=debug, chat=chat, stream_output=stream_output)
+                self.prompter = Prompter(self.prompt_type, self.prompt_dict, debug=debug, chat=chat,
+                                         stream_output=stream_output)
             self.human = self.prompter.humanstr
             self.bot = self.prompter.botstr
             self.can_stop = True
@@ -45,14 +51,75 @@ 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
 
+    @staticmethod
+    def limit_prompt(prompt_text, tokenizer, max_prompt_length=None):
+        verbose = bool(int(os.getenv('VERBOSE_PIPELINE', '0')))
+
+        if hasattr(tokenizer, 'model_max_length'):
+            # model_max_length only defined for generate.py, not raw use of h2oai_pipeline.py
+            model_max_length = tokenizer.model_max_length
+            if max_prompt_length is not None:
+                model_max_length = min(model_max_length, max_prompt_length)
+            # cut at some upper likely limit to avoid excessive tokenization etc
+            # upper bound of 10 chars/token, e.g. special chars sometimes are long
+            if len(prompt_text) > model_max_length * 10:
+                len0 = len(prompt_text)
+                prompt_text = prompt_text[-model_max_length * 10:]
+                if verbose:
+                    print("Cut of input: %s -> %s" % (len0, len(prompt_text)), flush=True)
+        else:
+            # unknown
+            model_max_length = None
+
+        num_prompt_tokens = None
+        if model_max_length is not None:
+            # can't wait for "hole" if not plain prompt_type, since would lose prefix like <human>:
+            # For https://github.com/h2oai/h2ogpt/issues/192
+            for trial in range(0, 3):
+                prompt_tokens = tokenizer(prompt_text)['input_ids']
+                num_prompt_tokens = len(prompt_tokens)
+                if num_prompt_tokens > model_max_length:
+                    # conservative by using int()
+                    chars_per_token = int(len(prompt_text) / num_prompt_tokens)
+                    # keep tail, where question is if using langchain
+                    prompt_text = prompt_text[-model_max_length * chars_per_token:]
+                    if verbose:
+                        print("reducing %s tokens, assuming average of %s chars/token for %s characters" % (
+                            num_prompt_tokens, chars_per_token, len(prompt_text)), flush=True)
+                else:
+                    if verbose:
+                        print("using %s tokens with %s chars" % (num_prompt_tokens, len(prompt_text)), flush=True)
+                    break
+
+            # Why Below False: don't limit max_new_tokens more, just rely upon stopping to reach limit of model
+            if False:
+                # if input prompt is some number of tokens, despite user request, can't have max_new_tokens more
+                #
+                assert num_prompt_tokens is not None
+                if self.prompt_type not in [PromptType.plain.name, PromptType.plain.value]:
+                    # then give room for prompt
+                    fudge = 20
+                else:
+                    fudge = 0
+                max_new_tokens = max(0, min(generate_kwargs['max_new_tokens'],
+                                            model_max_length - (num_prompt_tokens + fudge)))
+                if max_new_tokens < generate_kwargs['max_new_tokens']:
+                    if verbose:
+                        print("Reduced max_new_tokens from %s -> %s" % (
+                        generate_kwargs['max_new_tokens'], max_new_tokens))
+                    generate_kwargs['max_new_tokens'] = max_new_tokens
+        return prompt_text, num_prompt_tokens
+
     def preprocess(self, prompt_text, prefix="", handle_long_generation=None, **generate_kwargs):
+        prompt_text, num_prompt_tokens = H2OTextGenerationPipeline.limit_prompt(prompt_text, self.tokenizer)
+
         data_point = dict(context='', instruction=prompt_text, input='')
         if self.prompter is not None:
             prompt_text = self.prompter.generate_prompt(data_point)
         self.prompt_text = prompt_text
         if handle_long_generation is None:
             # forces truncation of inputs to avoid critical failure
-            handle_long_generation = 'hole'
+            handle_long_generation = None  # disable with new approaches
         return super().preprocess(prompt_text, prefix=prefix, handle_long_generation=handle_long_generation,
                                   **generate_kwargs)
 
@@ -65,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
@@ -73,8 +140,10 @@ class H2OTextGenerationPipeline(TextGenerationPipeline):
 
     def _forward(self, model_inputs, **generate_kwargs):
         if self.can_stop:
-            stopping_criteria = get_stopping(self.prompt_type, self.tokenizer, self.device, human=self.human,
-                                             bot=self.bot)
+            stopping_criteria = get_stopping(self.prompt_type, self.prompt_dict,
+                                             self.tokenizer, self.device,
+                                             human=self.human, bot=self.bot,
+                                             model_max_length=self.tokenizer.model_max_length)
             generate_kwargs['stopping_criteria'] = stopping_criteria
         # return super()._forward(model_inputs, **generate_kwargs)
         return self.__forward(model_inputs, **generate_kwargs)

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,30 +1,10 @@
+import os
+import ast
 import time
-from enum import Enum
+from enums import PromptType  # also supports imports from this file from other files
 
 non_hf_types = ['gpt4all_llama', 'llama', 'gptj']
 
-
-class PromptType(Enum):
-    plain = 0
-    instruct = 1
-    quality = 2
-    human_bot = 3
-    dai_faq = 4
-    summarize = 5
-    simple_instruct = 6
-    instruct_vicuna = 7
-    instruct_with_end = 8
-    human_bot_orig = 9
-    prompt_answer = 10
-    open_assistant = 11
-    wizard_lm = 12
-    wizard_mega = 13
-    instruct_vicuna2 = 14
-    instruct_vicuna3 = 15
-    wizard2 = 16
-    wizard3 = 17
-
-
 prompt_type_to_model_name = {
     'plain': [
         'EleutherAI/gpt-j-6B',
@@ -45,17 +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-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'],
@@ -68,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',
     ],
@@ -79,7 +72,17 @@ prompt_type_to_model_name = {
     "open_assistant": ['OpenAssistant/oasst-sft-7-llama-30b-xor', 'oasst-sft-7-llama-30b'],
     "wizard_lm": ['ehartford/WizardLM-7B-Uncensored', 'ehartford/WizardLM-13B-Uncensored'],
     "wizard_mega": ['openaccess-ai-collective/wizard-mega-13b'],
+    "instruct_simple": ['JosephusCheung/Guanaco'],
+    "wizard_vicuna": ['ehartford/Wizard-Vicuna-13B-Uncensored'],
+    "wizard2": ['llama', '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}
@@ -93,20 +96,53 @@ for p in PromptType:
     prompt_types.extend([p.name, p.value, str(p.value)])
 
 
-def get_prompt(prompt_type, chat, context, reduced):
-    if prompt_type in [PromptType.plain.value, str(PromptType.plain.value),
-                       PromptType.plain.name]:
-        promptA = promptB = PreInstruct = PreInput = PreResponse = ''
-        terminate_response = []
+def get_prompt(prompt_type, prompt_dict, chat, context, reduced, making_context, return_dict=False):
+    prompt_dict_error = ''
+    generates_leading_space = False
+
+    if prompt_type == PromptType.custom.name and not isinstance(prompt_dict, dict):
+        try:
+            prompt_dict = ast.literal_eval(prompt_dict)
+        except BaseException as e:
+            prompt_dict_error = str(e)
+    if prompt_dict_error:
+        promptA = None
+        promptB = None
+        PreInstruct = None
+        PreInput = ''
+        PreResponse = ''
+        terminate_response = None
         chat_sep = ''
+        chat_turn_sep = ''
         humanstr = ''
         botstr = ''
+        generates_leading_space = False
+    elif prompt_type in [PromptType.custom.value, str(PromptType.custom.value),
+                         PromptType.custom.name]:
+        promptA = prompt_dict.get('promptA', '')
+        promptB = prompt_dict('promptB', '')
+        PreInstruct = prompt_dict.get('PreInstruct', '')
+        PreInput = prompt_dict.get('PreInput', '')
+        PreResponse = prompt_dict.get('PreResponse', '')
+        terminate_response = prompt_dict.get('terminate_response', None)
+        chat_sep = prompt_dict.get('chat_sep', '\n')
+        chat_turn_sep = prompt_dict.get('chat_turn_sep', '\n')
+        humanstr = prompt_dict.get('humanstr', '')
+        botstr = prompt_dict.get('botstr', '')
+    elif prompt_type in [PromptType.plain.value, str(PromptType.plain.value),
+                         PromptType.plain.name]:
+        promptA = promptB = PreInstruct = PreInput = PreResponse = None
+        terminate_response = []
+        chat_turn_sep = chat_sep = ''
+        # plain should have None for human/bot, so nothing truncated out, not '' that would truncate after first token
+        humanstr = None
+        botstr = None
     elif prompt_type == 'simple_instruct':
         promptA = promptB = PreInstruct = PreInput = PreResponse = None
         terminate_response = []
-        chat_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),
@@ -132,7 +168,7 @@ def get_prompt(prompt_type, chat, context, reduced):
             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),
@@ -154,7 +190,7 @@ def get_prompt(prompt_type, chat, context, reduced):
 ### 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),
@@ -176,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:
@@ -191,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 = ''
@@ -210,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),
@@ -219,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),
@@ -239,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),
@@ -247,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
@@ -285,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),
@@ -301,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:
@@ -320,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:
@@ -339,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 = """
@@ -356,30 +409,136 @@ 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]:
+        promptB = promptA = '' if not (chat and reduced) else ''
 
+        PreInstruct = """
+### Instruction:
+"""
+
+        PreInput = """
+### Input:
+"""
+
+        PreResponse = """
+### Response:
+"""
+        terminate_response = None
+        chat_turn_sep = chat_sep = '\n'
+        humanstr = PreInstruct
+        botstr = PreResponse
+    elif prompt_type in [PromptType.openai.value, str(PromptType.openai.value),
+                         PromptType.openai.name]:
+        preprompt = """The following is a conversation with an AI assistant. The assistant is helpful, creative, clever, and very friendly.""" if not (
+                chat and reduced) else ''
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        PreInstruct = "\nHuman: "
+        PreInput = None
+        PreResponse = "\nAI:"
+        terminate_response = [PreResponse] + [" Human:", " AI:"]
+        chat_turn_sep = chat_sep = '\n'
+        humanstr = PreInstruct
+        botstr = PreResponse
+    elif prompt_type in [PromptType.gptj.value, str(PromptType.gptj.value),
+                         PromptType.gptj.name]:
+        preprompt = "### Instruction:\n The prompt below is a question to answer, a task to complete, or a conversation to respond to; decide which and write an appropriate response." if not (
+                chat and reduced) else ''
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        PreInstruct = "\n### Prompt: "
+        PreInput = None
+        PreResponse = "\n### Response: "
+        terminate_response = [PreResponse] + ["Prompt:", "Response:"]
+        chat_turn_sep = chat_sep = '\n'
+        humanstr = PreInstruct
+        botstr = PreResponse
+    elif prompt_type in [PromptType.openai_chat.value, str(PromptType.openai_chat.value),
+                         PromptType.openai_chat.name]:
+        # prompting and termination all handled by endpoint
+        preprompt = """"""
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        PreInstruct = ""
+        PreInput = None
+        PreResponse = ""
+        terminate_response = []
+        chat_turn_sep = chat_sep = '\n'
+        humanstr = None
+        botstr = None
+    elif prompt_type in [PromptType.vicuna11.value, str(PromptType.vicuna11.value),
+                         PromptType.vicuna11.name]:
+        preprompt = """A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. """ if not (
+                chat and reduced) else ''
+        start = ''
+        promptB = promptA = '%s%s' % (preprompt, start)
+        eos = '</s>'
+        PreInstruct = """USER: """
+        PreInput = None
+        PreResponse = """ASSISTANT:"""
+        terminate_response = [PreResponse]
+        chat_sep = ' '
+        chat_turn_sep = eos
+        humanstr = PreInstruct
+        botstr = PreResponse
+
+        if making_context:
+            # when making context, want it to appear as-if LLM generated, which starts with space after :
+            PreResponse = PreResponse + ' '
+        else:
+            # normally LLM adds space after this, because was how trained.
+            # if add space here, non-unique tokenization will often make LLM produce wrong output
+            PreResponse = PreResponse
     else:
         raise RuntimeError("No such prompt_type=%s" % prompt_type)
 
-    return promptA, promptB, PreInstruct, PreInput, PreResponse, terminate_response, chat_sep, humanstr, botstr
+    if isinstance(terminate_response, (tuple, list)):
+        assert '' not in terminate_response, "Bad terminate_response"
+
+    ret_dict = dict(promptA=promptA, promptB=promptB, PreInstruct=PreInstruct, PreInput=PreInput,
+                    PreResponse=PreResponse, terminate_response=terminate_response, chat_sep=chat_sep,
+                    chat_turn_sep=chat_turn_sep,
+                    humanstr=humanstr, botstr=botstr,
+                    generates_leading_space=generates_leading_space)
 
+    if return_dict:
+        return ret_dict, prompt_dict_error
+    else:
+        return tuple(list(ret_dict.values()))
 
-def generate_prompt(data_point, prompt_type, 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 = ''
@@ -387,11 +546,15 @@ def generate_prompt(data_point, prompt_type, chat, reduced):
     input = data_point.get('input')
     output = data_point.get('output')
     prompt_type = data_point.get('prompt_type', prompt_type)
+    prompt_dict = data_point.get('prompt_dict', prompt_dict)
     assert prompt_type in prompt_types, "Bad prompt type: %s" % prompt_type
     promptA, promptB, PreInstruct, PreInput, PreResponse, \
-        terminate_response, chat_sep, humanstr, botstr = get_prompt(prompt_type, 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}"""
@@ -400,37 +563,37 @@ def generate_prompt(data_point, prompt_type, 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}"""
@@ -441,23 +604,21 @@ def generate_prompt(data_point, prompt_type, 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
 
 
 class Prompter(object):
-    def __init__(self, prompt_type, debug=False, chat=False, stream_output=False, repeat_penalty=True,
+    def __init__(self, prompt_type, prompt_dict, debug=False, chat=False, stream_output=False, repeat_penalty=True,
                  allowed_repeat_line_length=10):
         self.prompt_type = prompt_type
-        data_point = dict(instruction='', input='', output='')
-        _, self.pre_response, self.terminate_response, self.chat_sep = \
-            generate_prompt(data_point, prompt_type, chat, False)
+        self.prompt_dict = prompt_dict
         self.debug = debug
         self.chat = chat
         self.stream_output = stream_output
@@ -466,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(prompt_type, chat, context, reduced)
-
-    def generate_prompt(self, data_point):
-        reduced = False
-        prompt, _, _, _ = generate_prompt(data_point, self.prompt_type, 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
 
@@ -493,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):
@@ -515,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:
@@ -547,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:
@@ -557,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
@@ -572,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.31.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.5
+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
@@ -53,17 +53,66 @@ bioc==2.0
 
 # falcon
 einops==0.6.1
+instructorembedding==1.0.1
+
+# for gpt4all .env file, but avoid worrying about imports
+python-dotenv==1.0.0
+
+text-generation==0.6.0
+# for tokenization when don't have HF tokenizer
+tiktoken==0.4.0
+# optional: for OpenAI endpoint or embeddings (requires key)
+openai==0.27.8
+# optional for chat with PDF
+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
+
+# local vector db
+chromadb==0.3.25
+# server vector db
+#pymilvus==2.2.8
+
+# weak url support, if can't install opencv etc. If comment-in this one, then comment-out unstructured[local-inference]==0.6.6
+# unstructured==0.6.6
+
+# strong support for images
+# Requires on Ubuntu: sudo apt-get install libmagic-dev poppler-utils tesseract-ocr libreoffice
+unstructured[local-inference]==0.7.4
+#pdf2image==1.16.3
+#pytesseract==0.3.10
+pillow
+
+pdfminer.six==20221105
+urllib3
+requests_file
+
+#pdf2image==1.16.3
+#pytesseract==0.3.10
+tabulate==0.9.0
+# FYI pandoc already part of requirements.txt
+
+# JSONLoader, but makes some trouble for some users
+# jq==1.4.1
+
+# to check licenses
+# Run: pip-licenses|grep -v 'BSD\|Apache\|MIT'
+pip-licenses==4.3.0
+
+# weaviate vector db
+weaviate-client==3.20.0
 # optional for chat with PDF
-langchain==0.0.183
-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
@@ -75,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
@@ -97,4 +146,8 @@ 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
+# extract-msg==0.41.1  # GPL3

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,16 @@ 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, 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:
             # encounters = [prompt.count(human) + 1, prompt.count(bot) + 1]
@@ -66,7 +71,8 @@ def get_stopping(prompt_type, tokenizer, device, human='<human>:', bot="<bot>:")
         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

@@ -14,6 +14,7 @@ import time
 import traceback
 import zipfile
 from datetime import datetime
+
 import filelock
 import requests, uuid
 from typing import Tuple, Callable, Dict
@@ -68,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()
@@ -97,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
@@ -166,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"
             )
 
@@ -800,6 +826,7 @@ def get_kwargs(func, exclude_names=None, **kwargs):
 
 
 import pkg_resources
+
 have_faiss = False
 
 try:
@@ -827,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:
@@ -835,9 +862,67 @@ 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()
         md5 = None
     return md5.hexdigest()
+
+
+def start_faulthandler():
+    # If hit server or any subprocess with signal SIGUSR1, it'll print out all threads stack trace, but wont't quit or coredump
+    # If more than one fork tries to write at same time, then looks corrupted.
+    import faulthandler
+
+    # SIGUSR1 in h2oai/__init__.py as well
+    faulthandler.enable()
+    if hasattr(faulthandler, 'register'):
+        # windows/mac
+        import signal
+        faulthandler.register(signal.SIGUSR1)
+
+
+def get_hf_server(inference_server):
+    inf_split = inference_server.split("    ")
+    assert len(inf_split) == 1 or len(inf_split) == 3
+    inference_server = inf_split[0]
+    if len(inf_split) == 3:
+        headers = {"authorization": "%s %s" % (inf_split[1], inf_split[2])}
+    else:
+        headers = None
+    return inference_server, headers
+
+
+class FakeTokenizer:
+    """
+    1) For keeping track of model_max_length
+    2) For when model doesn't directly expose tokenizer but need to count tokens
+    """
+
+    def __init__(self, model_max_length=2048, encoding_name="cl100k_base"):
+        # dont' push limit, since if using fake tokenizer, only estimate, and seen underestimates by order 250
+        self.model_max_length = model_max_length - 250
+        self.encoding_name = encoding_name
+        # The first time this runs, it will require an internet connection to download. Later runs won't need an internet connection.
+        import tiktoken
+        self.encoding = tiktoken.get_encoding(self.encoding_name)
+
+    def encode(self, x, *args, return_tensors="pt", **kwargs):
+        input_ids = self.encoding.encode(x, disallowed_special=())
+        if return_tensors == 'pt' and isinstance(input_ids, list):
+            import torch
+            input_ids = torch.tensor(input_ids)
+        return dict(input_ids=input_ids)
+
+    def decode(self, x, *args, **kwargs):
+        # input is input_ids[0] form
+        return self.encoding.decode(x)
+
+    def num_tokens_from_string(self, prompt: str) -> int:
+        """Returns the number of tokens in a text string."""
+        num_tokens = len(self.encoding.encode(prompt))
+        return num_tokens
+
+    def __call__(self, x, *args, **kwargs):
+        return self.encode(x, *args, **kwargs)

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