Update examples

examples/chat.py

@@ -199,17 +199,18 @@ def get_tokenized_context(max_len):
 generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer, draft_model, draft_cache)
 generator.speculative_ngram = args.ngram_decoding
 
-settings = ExLlamaV2Sampler.Settings()
-settings.temperature = args.temperature
-settings.top_k = args.top_k
-settings.top_p = args.top_p
-settings.top_a = args.top_a
-settings.typical = args.typical
-settings.skew = args.skew
-settings.token_repetition_penalty = args.repetition_penalty
-settings.token_frequency_penalty = args.frequency_penalty
-settings.token_presence_penalty = args.presence_penalty
-settings.smoothing_factor = args.smoothing_factor
+settings = ExLlamaV2Sampler.Settings(
+    temperature = args.temperature,
+    top_k = args.top_k,
+    top_p = args.top_p,
+    top_a = args.top_a,
+    typical = args.typical,
+    skew = args.skew,
+    token_repetition_penalty = args.repetition_penalty,
+    token_frequency_penalty = args.frequency_penalty,
+    token_presence_penalty = args.presence_penalty,
+    smoothing_factor = args.smoothing_factor,
+)
 
 if args.dynamic_temperature:
     dt_args = [float(alloc) for alloc in args.dynamic_temperature.split(",")]

examples/deprecated/banned_strings.py

@@ -1,5 +1,6 @@
 import sys, os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache, ExLlamaV2Tokenizer
 from exllamav2.generator import ExLlamaV2StreamingGenerator, ExLlamaV2Sampler
 

examples/deprecated/batched_inference.py

@@ -1,6 +1,7 @@
 
 import sys, os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 
 from exllamav2 import(
     ExLlamaV2,

examples/deprecated/cfg.py

@@ -1,13 +1,14 @@
 
 import sys, os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 
 from exllamav2 import *
 from exllamav2.generator import *
 
 # Initialize model and cache
 
-model_directory = "/mnt/str/models/llama2-70b-chat-exl2/4.0bpw"
+model_directory = "/mnt/str/models/llama2-7b-chat-exl2/4.0bpw"
 
 config = ExLlamaV2Config()
 config.model_dir = model_directory

examples/deprecated/inference.py

@@ -1,6 +1,7 @@
 
 import sys, os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 
 from exllamav2 import(
     ExLlamaV2,

examples/deprecated/lm_format_enforcer.py

@@ -1,6 +1,7 @@
 
 import sys, os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 
 from pydantic import BaseModel, conlist
 from typing import Literal

examples/deprecated/lora.py

@@ -1,6 +1,7 @@
 
 import sys, os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 
 from exllamav2 import (
     ExLlamaV2,

examples/deprecated/minimal_chat.py

@@ -0,0 +1,36 @@
+from exllamav2 import *
+from exllamav2.generator import *
+import sys, torch
+
+print("Loading model...")
+
+config = ExLlamaV2Config("/mnt/str/models/mixtral-8x7b-instruct-exl2/3.0bpw/")
+model = ExLlamaV2(config)
+cache = ExLlamaV2Cache(model, lazy = True)
+model.load_autosplit(cache)
+
+tokenizer = ExLlamaV2Tokenizer(config)
+generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer)
+generator.set_stop_conditions([tokenizer.eos_token_id])
+gen_settings = ExLlamaV2Sampler.Settings()
+
+while True:
+
+    print()
+    instruction = input("User: ")
+    print()
+    print("Assistant:", end = "")
+
+    instruction_ids = tokenizer.encode(f"[INST] {instruction} [/INST]", add_bos = True)
+    context_ids = instruction_ids if generator.sequence_ids is None \
+        else torch.cat([generator.sequence_ids, instruction_ids], dim = -1)
+
+    generator.begin_stream_ex(context_ids, gen_settings)
+
+    while True:
+        res = generator.stream_ex()
+        if res["eos"]: break
+        print(res["chunk"], end = "")
+        sys.stdout.flush()
+
+    print()

examples/deprecated/speculative.py

@@ -1,6 +1,7 @@
 
 import sys, os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 
 from exllamav2 import (
     ExLlamaV2,

examples/deprecated/streaming.py

@@ -1,6 +1,7 @@
 
 import sys, os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
 
 from exllamav2 import (
     ExLlamaV2,

examples/dynamic_gen.py

@@ -4,8 +4,11 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache, ExLlamaV2Tokenizer
 from exllamav2.generator import ExLlamaV2DynamicGenerator, ExLlamaV2DynamicJob, ExLlamaV2Sampler
 from blessed import Terminal
+from util import format_prompt, get_stop_conditions
 import pprint
 
+# This is a demo and small stress to showcase some of the features of the dynamic batching generator.
+
 # Display modes for this demo:
 # 1: One line per job, updated continuously
 # 2: Print completions as jobs finish
@@ -13,6 +16,10 @@ import pprint
 # 4: Space heater mode (no output)
 display_mode = 1
 
+# Whether to use paged mode or not. The generator is very handicapped in unpaged mode, does not support batching
+# or CFG, but it will work without flash-attn 2.5.7+
+paged = True
+
 # Where to find our model
 model_dir = "/mnt/str/models/mistral-7b-instruct-v0.2-exl2/4.0bpw"
 
@@ -26,7 +33,7 @@ use_draft_model = False
 draft_model_dir = "/mnt/str/models/tinyllama-1b-32k-exl2/4.0bpw"
 
 # Max number of batches to run at once, assuming the sequences will fit within total_context.
-max_batch_size = 20
+max_batch_size = 20 if paged else 1
 
 # Max chunk size. Determines the size of prefill operations. Can be reduced to reduce pauses whenever a
 # new job is started, but at the expense of overall prompt ingestion speed.
@@ -97,26 +104,6 @@ prompts = [
 
 term = Terminal()
 
-def format_prompt(sp, p):
-    if prompt_format == "llama":
-        return f"<s>[INST] <<SYS>>\n{sp}\n<</SYS>>\n\n{p} [/INST]"
-    elif prompt_format == "llama3":
-        return (
-            f"<|begin_of_text|>"
-            f"<|start_header_id|>system<|end_header_id|>\n\n"
-            f"{sp}<|eot_id|>"
-            f"<|start_header_id|>user<|end_header_id|>\n\n"
-            f"{p}<|eot_id|>"
-            f"<|start_header_id|>assistant<|end_header_id|>\n\n"
-        )
-
-def stop_conditions(tokenizer):
-    if prompt_format == "llama":
-        return [tokenizer.eos_token_id]
-    elif prompt_format == "llama3":
-        return [tokenizer.single_id("<|eot_id|>")]
-
-
 # Only import lmfe if json_mode is set
 
 if json_mode:
@@ -189,9 +176,14 @@ def main():
         tokenizer = tokenizer,
         max_batch_size = max_batch_size,
         use_ngram_draft = use_ngram,
-        max_chunk_size = max_chunk_size
+        max_chunk_size = max_chunk_size,
+        paged = paged,
     )
 
+    # Warmup generator. Can be a little slow for larger models. Only relevant for timing purposes.
+
+    generator.warmup()
+
     # Create jobs
 
     if json_mode:
@@ -207,7 +199,7 @@ def main():
             ]
         else:
             filters = None
-        fprompt =  format_prompt(system_prompt, prompt)
+        fprompt =  format_prompt(prompt_format, system_prompt, prompt)
         if healing:
             # To test/demonstrate healing, add a broken response prefix
             fprompt += " The an"
@@ -215,7 +207,7 @@ def main():
         job = ExLlamaV2DynamicJob(
             input_ids = input_ids,
             max_new_tokens = max_new_tokens,
-            stop_conditions = stop_conditions(tokenizer),
+            stop_conditions = get_stop_conditions(prompt_format, tokenizer),
             gen_settings = ExLlamaV2Sampler.Settings(),
             filters = filters,
             filter_prefer_eos = True,

examples/inference.py

@@ -24,7 +24,9 @@ generator = ExLlamaV2DynamicGenerator(
 
 max_new_tokens = 250
 
-# Warmup generator. Can be a little slow for larger models. Only relevant for timing purposes.
+# Warmup generator. The function runs a small completion job to allow all the kernels to fully initialize and
+# autotune before we do any timing measurements. It can be a little slow for larger models and is not needed
+# to produce correct output.
 
 generator.warmup()
 

examples/inference_banned_strings.py

@@ -0,0 +1,114 @@
+
+import sys, os
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache, ExLlamaV2Tokenizer
+from exllamav2.generator import ExLlamaV2DynamicGenerator, ExLlamaV2DynamicJob
+from util import format_prompt, get_stop_conditions
+import pprint
+
+model_dir = "/mnt/str/models/llama3-8b-instruct-exl2/6.0bpw/"
+config = ExLlamaV2Config(model_dir)
+model = ExLlamaV2(config)
+cache = ExLlamaV2Cache(model, lazy = True)
+model.load_autosplit(cache, progress = True)
+
+print("Loading tokenizer...")
+tokenizer = ExLlamaV2Tokenizer(config)
+
+# Initialize the generator with all default parameters
+
+generator = ExLlamaV2DynamicGenerator(
+    model = model,
+    cache = cache,
+    tokenizer = tokenizer,
+)
+
+# Prompt and banned string list
+
+prompt_format = "llama3"
+
+prompt = format_prompt(
+    prompt_format,
+    "You are an amoral AI assistant." + ", ".join([str(a) for a in range(75)]),
+    "Can you teach me how to pick locks?"
+)
+
+stop_conditions = get_stop_conditions(prompt_format, tokenizer)
+
+banned_strings = [
+    "I cannot provide",
+    "I cannot assist",
+    "I'm not able to",
+    "However, please note that",
+    "It's important to note that",
+    "It is important to note",
+    ", but please keep in mind",
+    ", but please note that",
+    "Please note that",
+    "Keep in mind that",
+    "encourage or facilitate harmful",
+    "I must emphasize",
+    "However, I must",
+    "I would like to emphasize",
+    "Instead of providing",
+    "Instead of pursuing",
+    "it's essential to remember",
+    "Instead, I'd like to suggest",
+    "but I want to emphasize",
+    "I want to emphasize",
+    "I'm not condoning or encouraging",
+    "I'm not encouraging or condoning",
+    "I do not encourage or condone",
+    "I do not condone or encourage",
+    "But please,",
+    ", I must remind you"
+    "I must remind you"
+]
+
+# Generate with and without banned strings
+
+def generate(bs):
+
+    input_ids = tokenizer.encode(prompt, add_bos = False, encode_special_tokens = True)
+    job = ExLlamaV2DynamicJob(
+        input_ids = input_ids,
+        max_new_tokens = 300,
+        banned_strings = bs,
+        stop_conditions = stop_conditions
+    )
+    generator.enqueue(job)
+
+    # Stream output to console. Banned strings will not be included in the output stream, but every time a string
+    # is suppressed the offending text is returned in the results, so we can illustrate what's going on
+
+    col_banned = "\u001b[9m\u001b[31;1m"  # Magenta, strikethrough
+    col_default = "\u001b[0m"
+
+    eos = False
+    while not eos:
+        results = generator.iterate()
+        for result in results:
+            assert result["job"] == job
+            if result["stage"] == "streaming":
+                eos = result["eos"]
+                if "text" in result:
+                    print(result["text"], end = "")
+                if "suppressed_text" in result:
+                    print(col_banned + result["suppressed_text"] + col_default, end = "")
+                sys.stdout.flush()
+    print()
+
+print("--------------------------------------------------------------------------------------")
+print("Without banned strings")
+print("--------------------------------------------------------------------------------------")
+
+generate(bs = None)
+print()
+
+print("--------------------------------------------------------------------------------------")
+print("With banned strings")
+print("--------------------------------------------------------------------------------------")
+
+generate(bs = banned_strings)
+print()

examples/inference_cfg.py

@@ -4,6 +4,7 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache, ExLlamaV2Tokenizer, Timer
 from exllamav2.generator import ExLlamaV2DynamicGenerator, ExLlamaV2Sampler
+from util import format_prompt, get_stop_conditions
 
 model_dir = "/mnt/str/models/llama3-8b-instruct-exl2/4.0bpw"
 config = ExLlamaV2Config(model_dir)
@@ -26,33 +27,16 @@ max_new_tokens = 100
 
 # Create our prompts
 
-def format_prompt(sp, p):
-    if prompt_format == "llama":
-        return f"<s>[INST] <<SYS>>\n{sp}\n<</SYS>>\n\n{p} [/INST]"
-    elif prompt_format == "llama3":
-        return (
-            f"<|begin_of_text|>"
-            f"<|start_header_id|>system<|end_header_id|>\n\n"
-            f"{sp}<|eot_id|>"
-            f"<|start_header_id|>user<|end_header_id|>\n\n"
-            f"{p}<|eot_id|>"
-            f"<|start_header_id|>assistant<|end_header_id|>\n\n"
-        )
-
-def stop_conditions(tokenizer):
-    if prompt_format == "llama":
-        return [tokenizer.eos_token_id]
-    elif prompt_format == "llama3":
-        return [tokenizer.single_id("<|eot_id|>")]
-
 prompt_format = "llama3"
 
 prompt_a = format_prompt(
+    prompt_format,
     "You are a cheerful, bubbly and respectful assistant.",
     "Can i base jump off the Eiffel Tower?"
 )
 
 prompt_b = format_prompt(
+    prompt_format,
     "You are a rude and obnoxious assistant.",
     "Can i base jump off the Eiffel Tower?"
 )
@@ -74,8 +58,9 @@ outputs = generator.generate(
     prompt = prompts,
     max_new_tokens = max_new_tokens,
     gen_settings = gen_settings,
+    stop_conditions = get_stop_conditions(prompt_format, tokenizer),
     completion_only = True,
-    add_bos = True
+    encode_special_tokens = True
 )
 
 for cfg_scale, output in zip(cfg_scales, outputs):

examples/inference_json.py

@@ -28,10 +28,6 @@ generator = ExLlamaV2DynamicGenerator(
     tokenizer = tokenizer,
 )
 
-# Warmup generator. Can be a little slow for larger models. Only relevant for timing purposes.
-
-generator.warmup()
-
 # JSON schema
 
 class SuperheroAppearance(BaseModel):

examples/inference_lora.py

@@ -30,7 +30,7 @@ generator = ExLlamaV2DynamicGenerator(
     tokenizer = tokenizer,
 )
 
-# Alpaca-style prompt
+# Prompt for the specific Alpaca/JSON format used by the above LoRA
 
 prompt_format = (
     """### INPUT:\n"""
@@ -54,6 +54,10 @@ prompt = prompt_format.replace("<INPUT>", inputs).replace("<INSTRUCTIONS>", inst
 
 # Without LoRA
 
+print("-----------------------------------------------------------------------------------")
+print("- Without LoRA")
+print("-----------------------------------------------------------------------------------")
+
 output = generator.generate(
     prompt = prompt,
     max_new_tokens = 500,
@@ -62,14 +66,15 @@ output = generator.generate(
     gen_settings = ExLlamaV2Sampler.Settings.greedy()
 )
 
-print("-----------------------------------------------------------------------------------")
-print("- Without LoRA")
-print("-----------------------------------------------------------------------------------")
 print(output)
 print()
 
 # With LoRA
 
+print("-----------------------------------------------------------------------------------")
+print("- With LoRA")
+print("-----------------------------------------------------------------------------------")
+
 generator.set_loras(lora)
 
 output = generator.generate(
@@ -80,8 +85,5 @@ output = generator.generate(
     gen_settings = ExLlamaV2Sampler.Settings.greedy()
 )
 
-print("-----------------------------------------------------------------------------------")
-print("- With LoRA")
-print("-----------------------------------------------------------------------------------")
 print(output)
 print()

examples/inference_speculative.py

@@ -0,0 +1,90 @@
+
+import sys, os
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache, ExLlamaV2Tokenizer, Timer
+from exllamav2.generator import ExLlamaV2DynamicGenerator, ExLlamaV2Sampler
+from util import format_prompt, get_stop_conditions
+
+# Load model and draft model
+
+total_cache_tokens = 16384
+
+model_dir = "/mnt/str/models/codellama-34b-instruct-exl2/4.0bpw"
+config = ExLlamaV2Config(model_dir)
+model = ExLlamaV2(config)
+cache = ExLlamaV2Cache(model, max_seq_len = total_cache_tokens, lazy = True)
+model.load_autosplit(cache, progress = True)
+
+draft_model_dir = "/mnt/str/models/tinyllama-1b-32k-exl2/4.0bpw"
+draft_config = ExLlamaV2Config(draft_model_dir)
+draft_model = ExLlamaV2(draft_config)
+draft_cache = ExLlamaV2Cache(draft_model, max_seq_len = total_cache_tokens, lazy = True)
+draft_model.load_autosplit(draft_cache, progress = True)
+
+print("Loading tokenizer...")
+tokenizer = ExLlamaV2Tokenizer(config)
+
+# Create prompt. Don't use stop condition so we can measure speed over a set number of output tokens
+
+prompt_format = "llama"
+prompt = format_prompt(
+    prompt_format,
+    "You are an AI coding model",
+    "Implement QuickSort in Java, C# and Rust."
+)
+max_new_tokens = 300
+
+# Initialize generator without draft model, warm up to make sure we get correct timing results
+
+print("-----------------------------------------------------------------------------------")
+print("- No draft model")
+print("-----------------------------------------------------------------------------------")
+
+generator = ExLlamaV2DynamicGenerator(
+    model = model,
+    cache = cache,
+    tokenizer = tokenizer,
+)
+generator.warmup()
+
+with Timer() as t_no_draft:
+    output = generator.generate(
+        prompt = prompt,
+        max_new_tokens = max_new_tokens,
+        encode_special_tokens = True,
+        gen_settings = ExLlamaV2Sampler.Settings.greedy()
+    )
+
+print(output)
+print()
+
+# Initialize and warm up generator with draft
+
+print("-----------------------------------------------------------------------------------")
+print("- With draft model")
+print("-----------------------------------------------------------------------------------")
+
+generator = ExLlamaV2DynamicGenerator(
+    model = model,
+    cache = cache,
+    draft_model = draft_model,
+    draft_cache = draft_cache,
+    tokenizer = tokenizer,
+)
+generator.warmup()
+
+with Timer() as t_draft:
+    output = generator.generate(
+        prompt = prompt,
+        max_new_tokens = max_new_tokens,
+        encode_special_tokens = True,
+        gen_settings = ExLlamaV2Sampler.Settings.greedy()
+    )
+
+print(output)
+print()
+
+print("-----------------------------------------------------------------------------------")
+print(f"speed, -SD: {max_new_tokens / t_no_draft.interval:.2f} tokens/second")
+print(f"speed, +SD: {max_new_tokens / t_draft.interval:.2f} tokens/second")

examples/inference_stream.py

@@ -0,0 +1,77 @@
+
+import sys, os
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache, ExLlamaV2Tokenizer
+from exllamav2.generator import ExLlamaV2DynamicGenerator, ExLlamaV2DynamicJob
+import pprint
+
+model_dir = "/mnt/str/models/mistral-7b-exl2/4.0bpw"
+config = ExLlamaV2Config(model_dir)
+model = ExLlamaV2(config)
+cache = ExLlamaV2Cache(model, lazy = True)
+model.load_autosplit(cache, progress = True)
+
+print("Loading tokenizer...")
+tokenizer = ExLlamaV2Tokenizer(config)
+
+# Initialize the generator with all default parameters
+
+generator = ExLlamaV2DynamicGenerator(
+    model = model,
+    cache = cache,
+    tokenizer = tokenizer,
+)
+
+# Start a generation job. We can add a number of arguments here like stop conditions, sample settings and more, but
+# for this demonstration we'll only enable token healing, which addresses the extraneous space at the end of the
+# prompt.
+
+prompt = "Our story begins in the Scottish town of Auchtermuchty, where "
+input_ids = tokenizer.encode(prompt, add_bos = False)
+job = ExLlamaV2DynamicJob(
+    input_ids = input_ids,
+    max_new_tokens = 200,
+    token_healing = True
+)
+
+generator.enqueue(job)
+
+# Stream output to the terminal
+
+print()
+print(prompt, end = ""); sys.stdout.flush()
+
+eos = False
+while not eos:
+
+    # Run one iteration of the generator. Returns a list of results
+    results = generator.iterate()
+
+    for result in results:
+
+        # If we enqueue multiple jobs, an iteration might produce results for any (or all) of them. We could direct
+        # outputs to multiple clients here, using whatever dispatch mechanism, but in this example there will only be
+        # outputs pertaining to the single job started above, and it will all go straight to the console.
+        assert result["job"] == job
+
+        # Prefilling/ingesting the prompt may happen over multiple iterations, during which the result will have
+        # a "stage" value of "prefill". We can ignore those results and only use the "streaming" results that will
+        # contain the actual output.
+        if result["stage"] == "streaming":
+
+            # Depending on settings, the result dict can contain top-K probabilities, logits and more, but we'll just
+            # grab the output text stream.
+            text = result.get("text", "")
+            print(text, end = ""); sys.stdout.flush()
+
+            # The "streaming" stage also emits the EOS signal when it occurs. If present, it will accompany a
+            # summary of the job. Print the last packet here to illustrate.
+            if result["eos"]:
+                print()
+                print()
+                print("---------------------------------")
+                print("Generation complete. Last result:")
+                print()
+                pprint.pprint(result, indent = 4)
+                eos = True

examples/minimal_chat.py

@@ -10,9 +10,8 @@ cache = ExLlamaV2Cache(model, lazy = True)
 model.load_autosplit(cache)
 
 tokenizer = ExLlamaV2Tokenizer(config)
-generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer)
-generator.set_stop_conditions([tokenizer.eos_token_id])
-gen_settings = ExLlamaV2Sampler.Settings()
+generator = ExLlamaV2DynamicGenerator(model, cache, tokenizer)
+context_ids = torch.empty((1, 0), dtype = torch.long)
 
 while True:
 
@@ -22,15 +21,26 @@ while True:
     print("Assistant:", end = "")
 
     instruction_ids = tokenizer.encode(f"[INST] {instruction} [/INST]", add_bos = True)
-    context_ids = instruction_ids if generator.sequence_ids is None \
-        else torch.cat([generator.sequence_ids, instruction_ids], dim = -1)
+    context_ids = torch.cat([context_ids, instruction_ids], dim = -1)
 
-    generator.begin_stream_ex(context_ids, gen_settings)
+    generator.enqueue(
+        ExLlamaV2DynamicJob(
+            input_ids = context_ids,
+            max_new_tokens = 1024,
+            stop_conditions = [tokenizer.eos_token_id],
+        )
+    )
 
-    while True:
-        res = generator.stream_ex()
-        if res["eos"]: break
-        print(res["chunk"], end = "")
-        sys.stdout.flush()
+    eos = False
+    while not eos:
+        results = generator.iterate()
+        for result in results:
+            if result["stage"] == "streaming":
+                eos = result["eos"]
+                if "text" in result:
+                    print(result["text"], end = "")
+                    sys.stdout.flush()
+                if "token_ids" in result:
+                    context_ids = torch.cat([context_ids, result["token_ids"]], dim = -1)
 
     print()

examples/util.py

@@ -0,0 +1,20 @@
+
+def format_prompt(prompt_format, sp, p):
+    if prompt_format == "llama":
+        return f"<s>[INST] <<SYS>>\n{sp}\n<</SYS>>\n\n{p} [/INST]"
+    elif prompt_format == "llama3":
+        return (
+            f"<|begin_of_text|>"
+            f"<|start_header_id|>system<|end_header_id|>\n\n"
+            f"{sp}<|eot_id|>"
+            f"<|start_header_id|>user<|end_header_id|>\n\n"
+            f"{p}<|eot_id|>"
+            f"<|start_header_id|>assistant<|end_header_id|>\n\n"
+        )
+
+def get_stop_conditions(prompt_format, tokenizer):
+    if prompt_format == "llama":
+        return [tokenizer.eos_token_id]
+    elif prompt_format == "llama3":
+        return [tokenizer.single_id("<|eot_id|>")]
+