Update app.py

app.py

@@ -21,7 +21,7 @@ examples = ["I have been feeling more and more down for over a month. I have sta
 
 class StopOnTokens(StoppingCriteria):
  def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
- stop_ids = [29, 0]
+ stop_ids = [1, 2]
  for stop_id in stop_ids:
  if input_ids[0][-1] == stop_id:
  return True
@@ -36,9 +36,16 @@ def predict(message, history):
 
  messages = "".join(["".join([sys_msg + "\n###USER:"+item[0], "\n###ASSISTANT:"+item[1]]) #curr_system_message +
  for item in history_transformer_format])
+
+ # def format_prompt(q):
+ # return f"""{sys_msg}
+ # ###USER: {q}
+ # ###ASSISTANT:"""
+
+ # messages = format_prompt(message)
  
  model_inputs = tokenizer([messages], return_tensors="pt").to(device)
- streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+ streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=False)
  generate_kwargs = dict(
  model_inputs,
  streamer=streamer,
@@ -46,8 +53,10 @@ def predict(message, history):
  do_sample=True,
  top_p=0.95,
  top_k=1000,
- temperature=1.0,
+ temperature=0.2,
  num_beams=1,
+ eos_token_id=[tokenizer.eos_token_id],
+ pad_token_id=tokenizer.eos_token_id,
  stopping_criteria=StoppingCriteriaList([stop])
  )
  t = Thread(target=model.generate, kwargs=generate_kwargs)
@@ -55,13 +64,12 @@ def predict(message, history):
 
  partial_message = ""
  for new_token in streamer:
- if new_token != '#':
+ if new_token != '<':
+ # if "#" in new_token:
+ # break
+ # else:
  partial_message += new_token
  yield partial_message
- else:
- print("new token = #")
- partial_message += new_token
- yield partial_message
 
 
 gr.ChatInterface(