Update summarize.py

summarize.py

@@ -27,9 +27,9 @@ 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, **kwargs):
     """
-    summarize_and_score - given a batch of ids and a mask, return a summary and a score for the summary
+    summarize - given a batch of ids and a mask, returns a summary and the token length of the output summary
     Args:
         ids (): the batch of ids
         mask (): the attention mask for the batch
@@ -53,8 +53,6 @@ def summarize_and_score(ids, mask, model, tokenizer, **kwargs):
         input_ids,
         attention_mask=attention_mask,
         #global_attention_mask=global_attention_mask,
-        output_scores=True,
-        return_dict_in_generate=True,
         **kwargs,
     )
     summary = tokenizer.batch_decode(
@@ -62,9 +60,8 @@ def summarize_and_score(ids, mask, model, tokenizer, **kwargs):
         skip_special_tokens=True,
         remove_invalid_values=True,
     )
-    score = round(summary_pred_ids.sequences_scores.cpu().numpy()[0], 4)
     len_res = len(summary_pred_ids.sequences.cpu().numpy()[0])
-    return summary, score, len_res
+    return summary, len_res
 
 
 def summarize_via_tokenbatches(
@@ -111,23 +108,21 @@ def summarize_via_tokenbatches(
 
     for _id, _mask in zip(in_id_arr, att_arr):
 
-        result, score, l = summarize_and_score(
+        result, l = summarize(
             ids=_id,
             mask=_mask,
             model=model,
             tokenizer=tokenizer,
             **kwargs,
         )
-        score = round(float(score), 4)
         rate = round(float((len(_id)-l)/len(_id)),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]}\nCompression:\t{rate}")
         pbar.update()
 
     pbar.close()