Update summarize.py

summarize.py

@@ -27,7 +27,7 @@ def load_model_and_tokenizer(model_name):
     return model, tokenizer
 
 
-def summarize_and_score(ids, mask, model, tokenizer, **kwargs):
+def summarize(ids, mask, model, tokenizer, model_arch, **kwargs):
     """
     summarize_and_score - given a batch of ids and a mask, return a summary and a score for the summary
     Args:
@@ -35,6 +35,7 @@ def summarize_and_score(ids, mask, model, tokenizer, **kwargs):
         mask (): the attention mask for the batch
         model   (): the model to use for summarization
         tokenizer (): the tokenizer to use for summarization
+        model
     Returns:
         str: the summary of the batch
     """
@@ -44,27 +45,32 @@ def summarize_and_score(ids, mask, model, tokenizer, **kwargs):
 
     input_ids = ids.to("cuda") if torch.cuda.is_available() else ids
     attention_mask = mask.to("cuda") if torch.cuda.is_available() else mask
-
-    #global_attention_mask = torch.zeros_like(attention_mask)
-    # put global attention on <s> token
-    #global_attention_mask[:, 0] = 1
-
-    summary_pred_ids = model.generate(
-        input_ids,
-        attention_mask=attention_mask,
-        #global_attention_mask=global_attention_mask,
-        output_scores=True,
-        return_dict_in_generate=True,
-        **kwargs,
-    )
+    
+    if model_arch == 'LED':
+       global_attention_mask = torch.zeros_like(attention_mask)
+       # put global attention on <s> token
+       global_attention_mask[:, 0] = 1
+       summary_pred_ids = model.generate(
+           input_ids,
+           attention_mask=attention_mask,
+           global_attention_mask=global_attention_mask,
+           return_dict_in_generate=True,
+           **kwargs,
+       )
+       
+    else:   
+       summary_pred_ids = model.generate(
+           input_ids,
+           attention_mask=attention_mask,
+           return_dict_in_generate=True,
+           **kwargs,
+       )
     summary = tokenizer.batch_decode(
         summary_pred_ids.sequences,
         skip_special_tokens=True,
         remove_invalid_values=True,
     )
-    score = round(summary_pred_ids.sequences_scores.cpu().numpy()[0], 4)
-
-    return summary, score
+    return summary
 
 
 def summarize_via_tokenbatches(
@@ -111,21 +117,21 @@ def summarize_via_tokenbatches(
 
     for _id, _mask in zip(in_id_arr, att_arr):
 
-        result, score = summarize_and_score(
+        result = summarize(
             ids=_id,
             mask=_mask,
             model=model,
             tokenizer=tokenizer,
             **kwargs,
         )
-        score = round(float(score), 4)
+        rate = round(float(len()), 3)
         _sum = {
             "input_tokens": _id,
             "summary": result,
-            "summary_score": score,
+            "compression_rate": rate,
         }
         gen_summaries.append(_sum)
-        print(f"\t{result[0]}\nScore:\t{score}")
+        print(f"\t{result[0]}\nRate:\t{rate}")
         pbar.update()
 
     pbar.close()