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()