Update app.py

app.py

@@ -49,6 +49,11 @@ quantization_config = BitsAndBytesConfig(
     bnb_4bit_quant_type="nf4")
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+        
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     torch_dtype=torch.bfloat16,
@@ -89,17 +94,20 @@ Crucially, always respond to and rebut the previous speaker's points in Harris's
     conversation.append({"role": "user", "content": message})
     
     input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
+    attention_mask = torch.ones_like(input_ids)
     
     streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
     
     generate_kwargs = dict(
         input_ids=input_ids,
+        attention_mask=attention_mask,
         max_new_tokens=max_new_tokens,
         do_sample=True,
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
-        eos_token_id=[128001,128008,128009],
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
         streamer=streamer,
         repetition_penalty=repetition_penalty,
     )