Update app.py

app.py

@@ -49,6 +49,10 @@ 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,19 +93,22 @@ Importantly, always respond to and rebut the previous speaker's points in Trump'
     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,
-        repetition_penalty=repetition_penalty,
-        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,
     )
 
     with torch.no_grad():