Update summarize.py

summarize.py

@@ -47,24 +47,24 @@ def summarize(ids, mask, model, tokenizer, model_arch, **kwargs):
     attention_mask = mask.to("cuda") if torch.cuda.is_available() else mask
     
     if model_arch == 'LED':
-       global_attention_mask = torch.zeros_like(attention_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,
-           return_dict_in_generate=True,
-           **kwargs,
-       )
+        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_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,