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(