Update app.py

app.py

@@ -19,7 +19,6 @@ from transformers import (
 
 
 model_name = "WangZeJun/bloom-3b-moss-chat"
-max_new_tokens = 1024
 
 
 print(f"Starting to load the model {model_name} into memory")
@@ -43,14 +42,20 @@ class StopOnTokens(StoppingCriteria):
 
 
 def convert_history_to_text(history):
-
  user_input = history[-1][0]
- 
  input_pattern = "{}</s>"
  text = input_pattern.format(user_input)
  return text
 
-
+def convert_all_history_to_text(history):
+ text = ""
+ for instance in history:
+ text += instance[0]
+ text += "</s>"
+ if instance[1]:
+ text += instance[1]
+ text += "</s>"
+ return text
 
 def log_conversation(conversation_id, history, messages, generate_kwargs):
  logging_url = os.getenv("LOGGING_URL", None)
@@ -78,7 +83,7 @@ def user(message, history):
  return "", history + [[message, ""]]
 
 
-def bot(history, temperature, top_p, top_k, repetition_penalty, conversation_id):
+def bot(history, temperature, top_p, top_k, repetition_penalty, max_new_tokens, conversation_id):
  print(f"history: {history}")
  # Initialize a StopOnTokens object
  stop = StopOnTokens()
@@ -136,6 +141,64 @@ def bot(history, temperature, top_p, top_k, repetition_penalty, conversation_id)
  history[-1][1] = partial_text
  yield history
 
+def multi_bot(history, temperature, top_p, top_k, repetition_penalty, max_new_tokens, conversation_id):
+ print(f"history: {history}")
+ # Initialize a StopOnTokens object
+ stop = StopOnTokens()
+
+ # Construct the input message string for the model by concatenating the current system message and conversation history
+ messages = convert_all_history_to_text(history)
+
+ # Tokenize the messages string
+ input_ids = tok(messages, return_tensors="pt").input_ids
+ input_ids = input_ids.to(m.device)
+ streamer = TextIteratorStreamer(
+ tok, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+ generate_kwargs = dict(
+ input_ids=input_ids,
+ max_new_tokens=max_new_tokens,
+ temperature=temperature,
+ do_sample=temperature > 0.0,
+ top_p=top_p,
+ top_k=top_k,
+ repetition_penalty=repetition_penalty,
+ streamer=streamer,
+ stopping_criteria=StoppingCriteriaList([stop]),
+ )
+
+ stream_complete = Event()
+
+ def generate_and_signal_complete():
+ m.generate(**generate_kwargs)
+ stream_complete.set()
+
+ def log_after_stream_complete():
+ stream_complete.wait()
+ log_conversation(
+ conversation_id,
+ history,
+ messages,
+ {
+ "top_k": top_k,
+ "top_p": top_p,
+ "temperature": temperature,
+ "repetition_penalty": repetition_penalty,
+ },
+ )
+
+ t1 = Thread(target=generate_and_signal_complete)
+ t1.start()
+
+ t2 = Thread(target=log_after_stream_complete)
+ t2.start()
+
+ # Initialize an empty string to store the generated text
+ partial_text = ""
+ for new_text in streamer:
+ partial_text += new_text
+ history[-1][1] = partial_text
+ yield history
+
 
 def get_uuid():
  return str(uuid4())
@@ -162,7 +225,8 @@ with gr.Blocks(
  ).style(container=False)
  with gr.Column():
  with gr.Row():
- submit = gr.Button("Submit")
+ single_submit = gr.Button("单轮对话")
+ multi_submit = gr.Button("多轮对话")
  stop = gr.Button("Stop")
  clear = gr.Button("Clear")
  with gr.Row():
@@ -172,18 +236,30 @@ with gr.Blocks(
  with gr.Row():
  temperature = gr.Slider(
  label="Temperature",
- value=0.1,
+ value=0.3,
  minimum=0.0,
  maximum=1.0,
- step=0.1,
+ step=0.05,
  interactive=True,
  info="Higher values produce more diverse outputs",
  )
+ with gr.Column():
+ with gr.Row():
+ repetition_penalty = gr.Slider(
+ label="Repetition Penalty",
+ value=1.2,
+ minimum=1.0,
+ maximum=2.0,
+ step=0.05,
+ interactive=True,
+ info="Penalize repetition — 1.0 to disable.",
+ )
+ with gr.Row():
  with gr.Column():
  with gr.Row():
  top_p = gr.Slider(
  label="Top-p (nucleus sampling)",
- value=1.0,
+ value=0.85,
  minimum=0.0,
  maximum=1,
  step=0.01,
@@ -204,17 +280,16 @@ with gr.Blocks(
  interactive=True,
  info="Sample from a shortlist of top-k tokens — 0 to disable and sample from all tokens.",
  )
- with gr.Column():
- with gr.Row():
- repetition_penalty = gr.Slider(
- label="Repetition Penalty",
- value=1.2,
- minimum=1.0,
- maximum=2.0,
- step=0.1,
- interactive=True,
- info="Penalize repetition — 1.0 to disable.",
- )
+ with gr.Row():
+ max_new_tokens = gr.Slider(
+ label="Maximum new tokens",
+ value=1024,
+ minimum=0,
+ maximum=2048,
+ step=1,
+ interactive=True,
+ )
+ 
  # with gr.Row():
  # gr.Markdown(
  # "demo 2",
@@ -234,12 +309,13 @@ with gr.Blocks(
  top_p,
  top_k,
  repetition_penalty,
+ max_new_tokens,
  conversation_id,
  ],
  outputs=chatbot,
  queue=True,
  )
- submit_click_event = submit.click(
+ submit_click_event = single_submit.click(
  fn=user,
  inputs=[msg, chatbot],
  outputs=[msg, chatbot],
@@ -252,6 +328,26 @@ with gr.Blocks(
  top_p,
  top_k,
  repetition_penalty,
+ max_new_tokens,
+ conversation_id,
+ ],
+ outputs=chatbot,
+ queue=True,
+ )
+ multi_click_event = multi_submit.click(
+ fn=user,
+ inputs=[msg, chatbot],
+ outputs=[msg, chatbot],
+ queue=False,
+ ).then(
+ fn=multi_bot,
+ inputs=[
+ chatbot,
+ temperature,
+ top_p,
+ top_k,
+ repetition_penalty,
+ max_new_tokens,
  conversation_id,
  ],
  outputs=chatbot,