Update summarize.py

test coming back to orig settings

summarize.py

@@ -27,7 +27,7 @@ def load_model_and_tokenizer(model_name):
     return model, tokenizer
 
 
-def summarize(ids, mask, model, tokenizer, model_arch, **kwargs):
+def summarize_and_score(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
     Args:
@@ -35,7 +35,6 @@ def summarize(ids, mask, model, tokenizer, model_arch, **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
     """
@@ -45,32 +44,27 @@ def summarize(ids, mask, model, tokenizer, model_arch, **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
-    
-    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,
-        )
+
+    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,
+    )
     summary = tokenizer.batch_decode(
         summary_pred_ids.sequences,
         skip_special_tokens=True,
         remove_invalid_values=True,
     )
-    return summary
+    score = round(summary_pred_ids.sequences_scores.cpu().numpy()[0], 4)
+
+    return summary, score
 
 
 def summarize_via_tokenbatches(
@@ -116,28 +110,22 @@ def summarize_via_tokenbatches(
     pbar = tqdm(total=len(in_id_arr))
 
     for _id, _mask in zip(in_id_arr, att_arr):
-    
-        if model=='Blaise-g/led_pubmed_sumpubmed_1' or model=='Blaise-g/led_large_sumpbumed_scitldr':
-           model_arch = 'LED'
-        else:
-           model_arch = 'LongT5'
-        
-        result = summarize(
+
+        result, score = summarize_and_score(
             ids=_id,
             mask=_mask,
             model=model,
-            model_arch=model_arch,
             tokenizer=tokenizer,
             **kwargs,
         )
-        rate = round(float((len(input_text)-len(result))/len(input_text)), 3)
+        score = round(float(score), 4)
         _sum = {
             "input_tokens": _id,
             "summary": result,
-            "compression_rate": rate,
+            "summary_score": score,
         }
         gen_summaries.append(_sum)
-        print(f"\t{result[0]}\nRate:\t{rate}")
+        print(f"\t{result[0]}\nScore:\t{score}")
         pbar.update()
 
     pbar.close()