Update onnxgenairun.py

onnxgenairun.py

@@ -3,16 +3,23 @@ import argparse
 import time
 import re
 
-
 def main(args):
     if args.verbose: print("Loading model...")
     if args.timings:
         started_timestamp = 0
         first_token_timestamp = 0
 
-    model = og.Model(f'{args.model}')
-    ##########model = og.Model(".\\")
+    config = og.Config(args.model_path)
+    config.clear_providers()
+    # if args.execution_provider != "cpu":
+    #    if args.verbose: print(f"Setting model to {args.execution_provider}")
+    #    config.append_provider(args.execution_provider)
+    
+    config.append_provider("dml")
+    model = og.Model(config)
+
     if args.verbose: print("Model loaded")
+    
     tokenizer = og.Tokenizer(model)
     tokenizer_stream = tokenizer.create_stream()
     if args.verbose: print("Tokenizer created")
@@ -26,6 +33,10 @@ def main(args):
 
     chat_template = '<|user|>\n{input} <|end|>\n<|assistant|>'
 
+    params = og.GeneratorParams(model)
+    params.set_search_options(**search_options)
+    # generator = og.Generator(model, params)
+
     # Keep asking for input prompts in a loop
     while True:
         text = input("Input: ")
@@ -40,10 +51,8 @@ def main(args):
 
         input_tokens = tokenizer.encode(prompt)
 
-        params = og.GeneratorParams(model)
-        params.set_search_options(**search_options)
-        # params.input_ids = input_tokens
         generator = og.Generator(model, params)
+        generator.append_tokens(input_tokens)
         if args.verbose: print("Generator created")
 
         if args.verbose: print("Running generation loop ...")
@@ -52,14 +61,13 @@ def main(args):
             new_tokens = []
 
         print()
-        print("Output:\n", end='', flush=True)
+        print("Output: ", end='', flush=True)
+
+        vPreviousDecoded = ""
+        vNewDecoded = ""
 
         try:
-            vPreviousDecoded = ""
-            vNewDecoded = ""
-            generator.append_tokens(input_tokens)
             while not generator.is_done():
-                # generator.compute_logits()
                 generator.generate_next_token()
                 if args.timings:
                     if first:
@@ -67,26 +75,25 @@ def main(args):
                         first = False
 
                 new_token = generator.get_next_tokens()[0]
-
-                ###print(tokenizer_stream.decode(new_token), end='', flush=True)
-
+                #print(tokenizer_stream.decode(new_token), end='', flush=True)
                 
                 vNewDecoded = tokenizer_stream.decode(new_token)
-                if  re.findall("^[\x2E\x3A\x3B]$", vPreviousDecoded) and vNewDecoded.startswith(" ") and (not vNewDecoded.startswith(" *")) :
-                    vNewDecoded = "\n" + vNewDecoded.replace(" ", "", 1)
+                #if  re.findall("^[\x2E\x3A\x3B]$", vPreviousDecoded) and vNewDecoded.startswith(" ") and (not vNewDecoded.startswith(" *")) :
+                if  re.fullmatch("^[\x2E\x3A\x3B]$", vPreviousDecoded) and vNewDecoded.startswith(" ") and (not vNewDecoded.startswith(" *")) :
 
-                print(vNewDecoded, end='', flush=True)
+                    # vNewDecoded = "\n" + vNewDecoded.replace(" ", "", 1)
+                    print("\n" + vNewDecoded.replace(" ", "", 1), end='', flush=True) 
+                else :
+                    print(vNewDecoded, end='', flush=True)
+                
                 vPreviousDecoded = vNewDecoded
-
+                
                 if args.timings: new_tokens.append(new_token)
         except KeyboardInterrupt:
             print("  --control+c pressed, aborting generation--")
         print()
         print()
 
-        # Delete the generator to free the captured graph for the next generator, if graph capture is enabled
-        del generator
-
         if args.timings:
             prompt_time = first_token_timestamp - started_timestamp
             run_time = time.time() - first_token_timestamp
@@ -95,7 +102,8 @@ def main(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS, description="End-to-end AI Question/Answer example for gen-ai")
-    parser.add_argument('-m', '--model', type=str, required=True, help='Onnx model folder path (must contain config.json and model.onnx)')
+    parser.add_argument('-m', '--model_path', type=str, required=True, help='Onnx model folder path (must contain genai_config.json and model.onnx)')
+    # parser.add_argument('-e', '--execution_provider', type=str, required=True, choices=["cpu", "cuda", "dml"], help="Execution provider to run ONNX model with")
     parser.add_argument('-i', '--min_length', type=int, help='Min number of tokens to generate including the prompt')
     parser.add_argument('-l', '--max_length', type=int, help='Max number of tokens to generate including the prompt')
     parser.add_argument('-ds', '--do_sample', action='store_true', default=False, help='Do random sampling. When false, greedy or beam search are used to generate the output. Defaults to false')
@@ -106,4 +114,4 @@ if __name__ == "__main__":
     parser.add_argument('-v', '--verbose', action='store_true', default=False, help='Print verbose output and timing information. Defaults to false')
     parser.add_argument('-g', '--timings', action='store_true', default=False, help='Print timing information for each generation step. Defaults to false')
     args = parser.parse_args()
-    main(args)
+    main(args)