Update README.md

README.md

@@ -24,12 +24,15 @@ from transformers import (
 
 
 def summarize(text, tokenizer, model, num_beams=4, temperature=1, max_new_tokens=512):
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
     if len(text) < 20:
         raise ValueError('Text must be at least 20 characters long.')
     # This text template is important.
     inputs = tokenizer(f'{text}\n### סיכום:', return_tensors="pt")
     in_data = inputs.input_ids.to('cuda')
-    output_ids = model.generate(input_ids=in_data, num_beams=num_beams, max_new_tokens = max_new_tokens, do_sample=True, early_stopping=True, use_cache=True, temperature=temperature, eos_token_id=tokenizer.eos_token_id)
+    attention_mask = inputs.attention_mask.to('cuda')
+    output_ids = model.generate(input_ids=in_data, attention_mask=attention_mask, num_beams=num_beams, max_new_tokens=max_new_tokens, do_sample=True, early_stopping=True, use_cache=True, temperature=temperature, pad_token_id=tokenizer.eos_token_id, eos_token_id=tokenizer.eos_token_id)
     generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=False)
     
     return generated_text