add interactive mode

app.py

@@ -2,6 +2,7 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 import gradio as gr
+import re
 import torch
 
 # load model and tokenizer
@@ -29,20 +30,31 @@ def chat(user_input, max_new_tokens=512):
     #create streamer
     streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
 
+    def get_start_idx(response, input):
+        match = re.search(re.escape(response), input)
+        if not match:
+            return -1
+        return match.end()
+
     def generate():
         model.generate(**inputs, max_new_tokens=max_new_tokens, streamer=streamer)
 
     thread = Thread(target=generate)
     thread.start()
 
-    prompt_len = len(prompt)
-    print(prompt)
+    start_idx = -1
     generated_text = ""
     for new_text in streamer:
         generated_text += new_text
-        print(generated_text)
+
+        if (start_idx == -1):
+            start_idx = get_start_idx(generated_text, user_input)
+            if (start_idx != -1):
+                start_idx += len("ASSISTANT")
+        #print(generated_text)
         #yield generated_text
-        yield generated_text[prompt_len:]
+        if (start_idx > 0):
+            yield generated_text[start_idx:]
 
     thread.join()