Update README.md

README.md

@@ -5,7 +5,6 @@ base_model:
 
 # Sample Inference Script
 ```py
-import re
 from argparse import ArgumentParser
 
 import torch
@@ -35,8 +34,8 @@ parser.add_argument("-a", "--audio", default="")
 parser.add_argument("-t", "--transcript", default="")
 parser.add_argument("-o", "--output", default="output.wav")
 parser.add_argument("-d", "--debug", action="store_true")
-parser.add_argument("--max_seq_len", type=int, default=2048)
 parser.add_argument("--sample_rate", type=int, default=16000)
+parser.add_argument("--max_seq_len", type=int, default=2048)
 parser.add_argument("--temperature", type=float, default=0.8)
 parser.add_argument("--top_p", type=float, default=1.0)
 args = parser.parse_args()
@@ -105,24 +104,18 @@ with Timer() as timer:
     input = template.render(messages=messages, eos_token="")
     input_ids = tokenizer.encode(input, add_bos=True, encode_special_tokens=True)
 
-    if args.debug:
-        print(input)
-
 print(f"Encoded input in {timer.interval:.2f} seconds.")
 
 with Timer() as timer:
-    max_new_tokens = config.max_seq_len - input_ids.shape[-1]
     gen_settings = ExLlamaV2Sampler.Settings()
     gen_settings.temperature = args.temperature
     gen_settings.top_p = args.top_p
-    stop_conditions = ["<|SPEECH_GENERATION_END|>"]
 
     job = ExLlamaV2DynamicJob(
         input_ids=input_ids,
-        max_new_tokens=max_new_tokens,
+        max_new_tokens=config.max_seq_len - input_ids.shape[-1],
         gen_settings=gen_settings,
-        stop_conditions=stop_conditions,
-        decode_special_tokens=True,
+        stop_conditions=["<|SPEECH_GENERATION_END|>"],
     )
 
     generator.enqueue(job)
@@ -131,11 +124,13 @@ with Timer() as timer:
     while generator.num_remaining_jobs():
         for result in generator.iterate():
             if result.get("stage") == "streaming":
-                text = result.get("text", "")
-                output.append(text)
+                text = result.get("text")
 
-                if args.debug:
-                    print(text, end="", flush=True)
+                if text:
+                    output.append(text)
+
+                    if args.debug:
+                        print(text, end="", flush=True)
 
             if result.get("eos"):
                 generator.clear_queue()
@@ -146,8 +141,7 @@ with Timer() as timer:
 print(f"Generated {len(output)} tokens in {timer.interval:.2f} seconds.")
 
 with Timer() as timer:
-    output = "".join(output)
-    output = [int(o) for o in re.findall(r"<\|s_(\d+)\|>", output)]
+    output = [int(o[4:-2]) for o in output]
     output = torch.tensor([[output]]).cuda()
     output = vocoder.decode_code(output)
     output = output[0, 0, :]