Update README.md

README.md

@@ -38,6 +38,38 @@ def summarize(text, tokenizer, model, num_beams=4, temperature=1, max_new_tokens
     return generated_text
 
 
+def summarize_batch(texts, tokenizer, model, num_beams=4, temperature=1, max_new_tokens=512):
+    for text in texts:
+        if len(text) < 20:
+            raise ValueError('Each text must be at least 20 characters long.')
+    
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    inputs = tokenizer([f'{text}\n### סיכום:' for text in texts], return_tensors="pt", padding=True)
+    
+    in_data = inputs.input_ids.to('cuda')
+    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
+    )
+
+    # Decode each generated summary
+    generated_texts = [tokenizer.decode(output, skip_special_tokens=False) for output in output_ids]
+    
+    return generated_texts
+
+
 model_path = 'maayanorner/hebrew-summarization-llm' # or maayanorner/hebrew-summarization-llm-4bit
 
 model = AutoModelForCausalLM.from_pretrained(