Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, StoppingCriteria, StoppingCriteriaList
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 import time
 import numpy as np
 from torch.nn import functional as F
@@ -28,43 +28,11 @@ class StopOnTokens(StoppingCriteria):
                 return True
         return False
 
-
-def contrastive_generate(text, bad_text):
-    with torch.no_grad():
-        tokens = tok(text, return_tensors="pt")[
-            'input_ids'].cuda()[:, :4096-1024]
-        bad_tokens = tok(bad_text, return_tensors="pt")[
-            'input_ids'].cuda()[:, :4096-1024]
-        history = None
-        bad_history = None
-        curr_output = list()
-        for i in range(1024):
-            out = m(tokens, past_key_values=history, use_cache=True)
-            logits = out.logits
-            history = out.past_key_values
-            bad_out = m(bad_tokens, past_key_values=bad_history,
-                        use_cache=True)
-            bad_logits = bad_out.logits
-            bad_history = bad_out.past_key_values
-            probs = F.softmax(logits.float(), dim=-1)[0][-1].cpu()
-            bad_probs = F.softmax(bad_logits.float(), dim=-1)[0][-1].cpu()
-            logits = torch.log(probs)
-            bad_logits = torch.log(bad_probs)
-            logits[probs > 0.1] = logits[probs > 0.1] - bad_logits[probs > 0.1]
-            probs = F.softmax(logits)
-            out = int(torch.multinomial(probs, 1))
-            if out in [50278, 50279, 50277, 1, 0]:
-                break
-            else:
-                curr_output.append(out)
-            out = np.array([out])
-            tokens = torch.from_numpy(np.array([out])).to(
-                tokens.device)
-            bad_tokens = torch.from_numpy(np.array([out])).to(
-                tokens.device)
-        return tok.decode(curr_output)
-
-
+def streaming_generate(text, bad_text=None):
+    
+    return model_output
+    
+    
 def generate(text, bad_text=None):
     stop = StopOnTokens()
     result = generator(text, max_new_tokens=1024, num_return_sequences=1, num_beams=1, do_sample=True,
@@ -81,9 +49,29 @@ def bot(history, curr_system_message):
     messages = curr_system_message + \
         "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]])
                 for item in history])
-    output = generate(messages)
-    history[-1][1] = output
-    time.sleep(1)
+    
+    model_inputs = tok(messages, return_tensors="pt")['input_ids'].cuda()[:, :4096-1024]
+
+    streamer = TextIteratorStreamer(tok, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        top_p=0.95,
+        top_k=1000,
+        temperature=1.0,
+        num_beams=1,
+        stopping_criteria=StoppingCriteriaList([stop])
+    )
+    t = Thread(target=m.generate, kwargs=generate_kwargs)
+    t.start()
+
+    model_output = ""
+    for new_text in streamer:
+        history[-1][1] += new_text
+        yield history
+        
     return history, history
 
 
@@ -107,5 +95,5 @@ with gr.Blocks() as demo:
     submit.click(fn=user, inputs=[msg, history], outputs=[msg, chatbot, history], queue=False).then(
         fn=bot, inputs=[chatbot, system_msg], outputs=[chatbot, history], queue=True)
     clear.click(lambda: [None, []], None, [chatbot, history], queue=False)
-demo.queue(concurrency_count=5)
+demo.queue(concurrency_count=1)
 demo.launch()