Update app.py

app.py

@@ -71,18 +71,19 @@ def generate(
         )
     conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt")
-    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-    input_ids = input_ids.to(model.device)
-    attention_mask = input_ids["attention_mask"]
+    formatted = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(formatted, return_tensors="pt", padding=True)
+    #if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+    #    input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+    #    gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+    inputs = inputs.to(model.device)
+    attention_mask = inputs["attention_mask"]
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
-        {"input_ids": input_ids},
+        {"input_ids": inputs["input_ids"]},
         streamer=streamer,
         max_new_tokens=max_new_tokens,
-        eos_token_id=tokenizer.eos_token_id,
+        #eos_token_id=tokenizer.eos_token_id,
         pad_token_id=tokenizer.eos_token_id,
         attention_mask=attention_mask,
         do_sample=True,