fix queue (#6)

- fix queue (9431ee6c6359396371b5551bbcb2e678cfa4b060)


Co-authored-by: Radamés Ajna <[email protected]>

app.py

@@ -4,11 +4,12 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, Stopping
 import time
 import numpy as np
 from torch.nn import functional as F
-import os 
-auth_key = os.environ["HF_ACCESS_TOKEN"]
+import os
+# auth_key = os.environ["HF_ACCESS_TOKEN"]
 print(f"Starting to load the model to memory")
-m = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=auth_key, torch_dtype=torch.float16).cuda()
-tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b",use_auth_token=auth_key)
+m = AutoModelForCausalLM.from_pretrained(
+ "stabilityai/stablelm-tuned-alpha-7b", torch_dtype=torch.float16).cuda()
+tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b")
 generator = pipeline('text-generation', model=m, tokenizer=tok, device=0)
 print(f"Sucessfully loaded the model to the memory")
 
@@ -30,8 +31,10 @@ class StopOnTokens(StoppingCriteria):
 
 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]
+ 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()
@@ -39,7 +42,8 @@ def contrastive_generate(text, bad_text):
  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_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()
@@ -60,39 +64,48 @@ def contrastive_generate(text, bad_text):
  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]))
+ 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):
- return "", history + [[user_message, ""]]
+ history = history + [[user_message, ""]]
+ return "", history, history
 
 
 def bot(history, curr_system_message):
- messages = curr_system_message + "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]]) for item in history])
+ 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)
- return history
-
-
+ return history, history
 
 
 with gr.Blocks() as demo:
- num = gr.State(value=0)
+ history = gr.State([])
  gr.Markdown("## StableLM-Tuned-Alpha-7b Chat")
  gr.HTML('''<center><a href="https://huggingface.co/spaces/stabilityai/stablelm-tuned-alpha-chat?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate the Space to skip the queue and run in a private space</center>''')
- chatbot = gr.Chatbot([])
- clear = gr.Button("Clear Chat History")
- system_msg = gr.Textbox(start_message, label="System Message", interactive=False,visible=False)
- #system_msg = start_message
- msg = gr.Textbox(label="Chat Message Box")
+ chatbot = gr.Chatbot().style(height=500)
+ with gr.Row():
+ with gr.Column(scale=0.70):
+ msg = gr.Textbox(label="", placeholder="Chat Message Box")
+ with gr.Column(scale=0.30, min_width=0):
+ with gr.Row():
+ submit = gr.Button("Submit")
+ clear = gr.Button("Clear") 
+ system_msg = gr.Textbox(
+ start_message, label="System Message", interactive=False, visible=False)
 
- msg.submit(user, [msg, chatbot], [msg, chatbot], queue=True).then(
- bot, [chatbot, system_msg], chatbot
- )
- clear.click(lambda: None, None, chatbot, queue=True)
+ msg.submit(fn=user, inputs=[msg, history], outputs=[msg, chatbot, history], queue=False).then(
+ fn=bot, inputs=[chatbot, system_msg], outputs=[chatbot, history], queue=True)
+ 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.launch()