extract inference

llama_lora/lib/inference.py

@@ -0,0 +1,77 @@
+import torch
+import transformers
+
+from .streaming_generation_utils import Iteratorize, Stream
+
+
+def generate(
+    # model
+    model,
+    tokenizer,
+    # input
+    prompt,
+    generation_config,
+    max_new_tokens,
+    stopping_criteria=[],
+    # output options
+    stream_output=False
+):
+    device = get_device()
+
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"].to(device)
+    generate_params = {
+        "input_ids": input_ids,
+        "generation_config": generation_config,
+        "return_dict_in_generate": True,
+        "output_scores": True,
+        "max_new_tokens": max_new_tokens,
+        "stopping_criteria": transformers.StoppingCriteriaList() + stopping_criteria
+    }
+
+    if stream_output:
+        # Stream the reply 1 token at a time.
+        # This is based on the trick of using 'stopping_criteria' to create an iterator,
+        # from https://github.com/oobabooga/text-generation-webui/blob/ad37f396fc8bcbab90e11ecf17c56c97bfbd4a9c/modules/text_generation.py#L216-L243.
+
+        def generate_with_callback(callback=None, **kwargs):
+            kwargs["stopping_criteria"].insert(
+                0,
+                Stream(callback_func=callback)
+            )
+            with torch.no_grad():
+                model.generate(**kwargs)
+
+        def generate_with_streaming(**kwargs):
+            return Iteratorize(
+                generate_with_callback, kwargs, callback=None
+            )
+
+        with generate_with_streaming(**generate_params) as generator:
+            for output in generator:
+                decoded_output = tokenizer.decode(output, skip_special_tokens=True)
+                yield decoded_output, output
+                if output[-1] in [tokenizer.eos_token_id]:
+                    break
+        return  # early return for stream_output
+
+    # Without streaming
+    with torch.no_grad():
+        generation_output = model.generate(**generate_params)
+    output = generation_output.sequences[0]
+    decoded_output = tokenizer.decode(output, skip_special_tokens=True)
+    yield decoded_output, output
+    return
+
+
+def get_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    else:
+        return "cpu"
+
+    try:
+        if torch.backends.mps.is_available():
+            return "mps"
+    except:  # noqa: E722
+        pass

llama_lora/models.py

@@ -60,9 +60,10 @@ def get_new_base_model(base_model_name):
             base_model_name, device_map={"": device}, low_cpu_mem_usage=True
         )
 
-    model.config.pad_token_id = get_tokenizer(base_model_name).pad_token_id = 0
-    model.config.bos_token_id = 1
-    model.config.eos_token_id = 2
+    tokenizer = get_tokenizer(base_model_name)
+    model.config.pad_token_id = tokenizer.pad_token_id = 0
+    model.config.bos_token_id = tokenizer.bos_token_id = 1
+    model.config.eos_token_id = tokenizer.eos_token_id = 2
 
     return model
 

llama_lora/ui/inference_ui.py

@@ -8,12 +8,12 @@ from transformers import GenerationConfig
 
 from ..globals import Global
 from ..models import get_model, get_tokenizer, get_device
+from ..lib.inference import generate
 from ..utils.data import (
     get_available_template_names,
     get_available_lora_model_names,
     get_info_of_available_lora_model)
 from ..utils.prompter import Prompter
-from ..utils.callbacks import Iteratorize, Stream
 
 device = get_device()
 
@@ -103,8 +103,6 @@ def do_inference(
         tokenizer = get_tokenizer(base_model_name)
         model = get_model(base_model_name, lora_model_name)
 
-        inputs = tokenizer(prompt, return_tensors="pt")
-        input_ids = inputs["input_ids"].to(device)
         generation_config = GenerationConfig(
             temperature=temperature,
             top_p=top_p,
@@ -113,26 +111,56 @@ def do_inference(
             num_beams=num_beams,
         )
 
-        generate_params = {
-            "input_ids": input_ids,
-            "generation_config": generation_config,
-            "return_dict_in_generate": True,
-            "output_scores": True,
-            "max_new_tokens": max_new_tokens,
-        }
-
         def ui_generation_stopping_criteria(input_ids, score, **kwargs):
             if Global.should_stop_generating:
                 return True
             return False
 
         Global.should_stop_generating = False
-        generate_params.setdefault(
-            "stopping_criteria", transformers.StoppingCriteriaList()
-        )
-        generate_params["stopping_criteria"].append(
-            ui_generation_stopping_criteria
-        )
+
+        generation_args = {
+            'model': model,
+            'tokenizer': tokenizer,
+            'prompt': prompt,
+            'generation_config': generation_config,
+            'max_new_tokens': max_new_tokens,
+            'stopping_criteria': [ui_generation_stopping_criteria],
+            'stream_output': stream_output
+        }
+
+        for (decoded_output, output) in generate(**generation_args):
+            raw_output_str = None
+            if show_raw:
+                raw_output_str = str(output)
+            response = prompter.get_response(decoded_output)
+
+            if Global.should_stop_generating:
+                return
+
+            yield (
+                gr.Textbox.update(
+                    value=response, lines=inference_output_lines),
+                raw_output_str)
+
+            if Global.should_stop_generating:
+                # If the user stops the generation, and then clicks the
+                # generation button again, they may mysteriously landed
+                # here, in the previous, should-be-stopped generation
+                # function call, with the new generation function not be
+                # called at all. To workaround this, we yield a message
+                # and setting lines=1, and if the front-end JS detects
+                # that lines has been set to 1 (rows="1" in HTML),
+                # it will automatically click the generate button again
+                # (gr.Textbox.update() does not support updating
+                # elem_classes or elem_id).
+                # [WORKAROUND-UI01]
+                yield (
+                    gr.Textbox.update(
+                        value="Please retry", lines=1),
+                    None)
+
+        return
+
 
         if stream_output:
             # Stream the reply 1 token at a time.