Update app.py

app.py

@@ -1,11 +1,12 @@
 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
 import os
-# auth_key = os.environ["HF_ACCESS_TOKEN"]
+from threading import Thread
+
 print(f"Starting to load the model to memory")
 m = AutoModelForCausalLM.from_pretrained(
     "stabilityai/stablelm-tuned-alpha-7b", torch_dtype=torch.float16).cuda()
@@ -28,62 +29,40 @@ 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 generate(text, bad_text=None):
-    stop = StopOnTokens()
-    result = generator(text, max_new_tokens=1024, num_return_sequences=1, num_beams=1, do_sample=True,
-                       temperature=1.0, top_p=0.95, top_k=1000, stopping_criteria=StoppingCriteriaList([stop]))
-    return result[0]["generated_text"].replace(text, "")
-
-
 def user(user_message, history):
     history = history + [[user_message, ""]]
     return "", history, history
 
 
 def bot(history, curr_system_message):
+    stop = StopOnTokens()
     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]
+    model_inputs = tok([messages], return_tensors="pt").to("cuda")
+    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()
+
+    print(history)
+    for new_text in streamer:
+        print(new_text)
+        history[-1][1] += new_text
+        yield history, history
+        
     return history, history
 
 
@@ -107,5 +86,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.launch()
+demo.queue(concurrency_count=2)
+demo.launch()