Update with h2oGPT hash cf3886c550581e34d9f05d69d2e3438b2a46d7b2

generate.py

@@ -5,6 +5,8 @@ import traceback
 import typing
 from threading import Thread
 
+import filelock
+
 from utils import set_seed, clear_torch_cache, save_generate_output, NullContext, wrapped_partial
 
 SEED = 1236
@@ -809,46 +811,52 @@ def evaluate(
                                     )
 
     with torch.no_grad():
-        # decoded tokenized prompt can deviate from prompt due to special characters
-        inputs_decoded = decoder(input_ids[0])
-        inputs_decoded_raw = decoder_raw(input_ids[0])
-        if inputs_decoded == prompt:
-            # normal
-            pass
-        elif inputs_decoded.lstrip() == prompt.lstrip():
-            # sometimes extra space in front, make prompt same for prompt removal
-            prompt = inputs_decoded
-        elif inputs_decoded_raw == prompt:
-            # some models specify special tokens that are part of normal prompt, so can't skip them
-            inputs_decoded_raw = inputs_decoded
-            decoder = decoder_raw
-        else:
-            print("WARNING: Special characters in prompt", flush=True)
-        decoded_output = None
-        if stream_output:
-            skip_prompt = False
-            streamer = TextIteratorStreamer(tokenizer, skip_prompt=skip_prompt)
-            gen_kwargs.update(dict(streamer=streamer))
-            target_func = generate_with_exceptions
-            target = wrapped_partial(generate_with_exceptions, model.generate, prompt, inputs_decoded,
-                                     raise_generate_gpu_exceptions, **gen_kwargs)
-            thread = Thread(target=target)
-            thread.start()
-            outputs = ""
-            for new_text in streamer:
-                outputs += new_text
+        # protection for gradio not keeping track of closed users,
+        # else hit bitsandbytes lack of thread safety:
+        # https://github.com/h2oai/h2ogpt/issues/104
+        # but only makes sense if concurrency_count == 1
+        context_class = NullContext if concurrency_count > 1 else filelock.FileLock
+        with context_class("generate.lock"):
+            # decoded tokenized prompt can deviate from prompt due to special characters
+            inputs_decoded = decoder(input_ids[0])
+            inputs_decoded_raw = decoder_raw(input_ids[0])
+            if inputs_decoded == prompt:
+                # normal
+                pass
+            elif inputs_decoded.lstrip() == prompt.lstrip():
+                # sometimes extra space in front, make prompt same for prompt removal
+                prompt = inputs_decoded
+            elif inputs_decoded_raw == prompt:
+                # some models specify special tokens that are part of normal prompt, so can't skip them
+                inputs_decoded_raw = inputs_decoded
+                decoder = decoder_raw
+            else:
+                print("WARNING: Special characters in prompt", flush=True)
+            decoded_output = None
+            if stream_output:
+                skip_prompt = False
+                streamer = TextIteratorStreamer(tokenizer, skip_prompt=skip_prompt)
+                gen_kwargs.update(dict(streamer=streamer))
+                target_func = generate_with_exceptions
+                target = wrapped_partial(generate_with_exceptions, model.generate, prompt, inputs_decoded,
+                                         raise_generate_gpu_exceptions, **gen_kwargs)
+                thread = Thread(target=target)
+                thread.start()
+                outputs = ""
+                for new_text in streamer:
+                    outputs += new_text
+                    yield prompter.get_response(outputs, prompt=inputs_decoded,
+                                                sanitize_bot_response=sanitize_bot_response)
+                decoded_output = outputs
+            else:
+                outputs = model.generate(**gen_kwargs)
+                outputs = [decoder(s) for s in outputs.sequences]
                 yield prompter.get_response(outputs, prompt=inputs_decoded,
                                             sanitize_bot_response=sanitize_bot_response)
-            decoded_output = outputs
-        else:
-            outputs = model.generate(**gen_kwargs)
-            outputs = [decoder(s) for s in outputs.sequences]
-            yield prompter.get_response(outputs, prompt=inputs_decoded,
-                                        sanitize_bot_response=sanitize_bot_response)
-            if outputs and len(outputs) >= 1:
-                decoded_output = prompt + outputs[0]
-        if save_dir and decoded_output:
-            save_generate_output(output=decoded_output, base_model=base_model, save_dir=save_dir)
+                if outputs and len(outputs) >= 1:
+                    decoded_output = prompt + outputs[0]
+            if save_dir and decoded_output:
+                save_generate_output(output=decoded_output, base_model=base_model, save_dir=save_dir)
 
 
 def generate_with_exceptions(func, prompt, inputs_decoded, raise_generate_gpu_exceptions, **kwargs):