chat.py: Add (SVG) save option

Author: turboderp

examples/chat.py

@@ -274,6 +274,21 @@ def get_input_ids(_prefix):
         # Add response to context
         response = s.all_text.strip()
 
+        # Optionally save output
+        if args.save:
+            sr = response
+            if args.save_svg:
+                m = re.search(r"<svg\b[^>]*>.*?</svg>", sr, flags = re.IGNORECASE | re.DOTALL)
+                sr = m.group(0).strip() if m else None
+                if sr: print_info(f"Found SVG: {len(sr)} characters")
+                else: print_error(f"No SVG block found")
+            if sr:
+                print_info(f"Writing response to: {args.save}")
+                with open(args.save, "w") as f:
+                    f.write(sr)
+            else:
+                print_info(f"Nothing to write")
+
         context[-1] = (user_prompt, response)
 
 
@@ -303,5 +318,7 @@ def get_input_ids(_prefix):
     parser.add_argument("-topp", "--top_p", type = float, help = "Top-P truncation, 1 to disable (default: disabled)", default = 1.0)
     parser.add_argument("-tps", "--show_tps", action = "store_true", help = "Show tokens/second after every reply")
     parser.add_argument("-prompt", "--prompt", type = str, help = "Run single prompt, then exit")
+    parser.add_argument("-save", "--save", type = str, help = "Save output to file (use with --prompt)")
+    parser.add_argument("-save_svg", "--save_svg", action = "store_true", help = "Extract SVG from response (use with --save)")
     _args = parser.parse_args()
     main(_args)