text: import argparse
  from pathlib import Path
  import torch
  from tqdm import tqdm
  from transformers import BartForConditionalGeneration, BartTokenizer
  
  DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
  
  
  def chunks(lst, n):
      for i in range(0, len(lst), n):
          yield lst[i: i + n]
  
  
  def generate_summaries(
      examples: list, out_file: str, model_name: str, batch_size: int = 8, device: str = DEFAULT_DEVICE
  ):
      fout = Path(out_file).open("w")
      model = BartForConditionalGeneration.from_pretrained(model_name).to(device)
      tokenizer = BartTokenizer.from_pretrained("bart-large")
  
      max_length = 140
      min_length = 55
  
      for batch in tqdm(list(chunks(examples, batch_size))):
          dct = tokenizer.batch_encode_plus(
              batch, max_length=1024, return_tensors="pt", pad_to_max_length=True)
          summaries = model.generate(
              input_ids=dct["input_ids"].to(device),
              attention_mask=dct["attention_mask"].to(device),
              num_beams=4,
              length_penalty=2.0,
              # +2 from original because we start at step=1 and stop before max_length
              max_length=max_length + 2,
              min_length=min_length + 1,  # +1 from original because we start at step=1
              no_repeat_ngram_size=3,
              early_stopping=True,
              decoder_start_token_id=model.config.eos_token_id,
          )
          dec = [tokenizer.decode(g, skip_special_tokens=True,
                                  clean_up_tokenization_spaces=False) for g in summaries]
          for hypothesis in dec:
              fout.write(hypothesis + "\n")
              fout.flush()
  
  
  def run_generate():
      parser = argparse.ArgumentParser()
      parser.add_argument(
          "source_path", type=str, help="like cnn_dm/test.source",
      )
      parser.add_argument(
          "output_path", type=str, help="where to save summaries",
      )
      parser.add_argument(
          "model_name", type=str, default="bart-large-cnn", help="like bart-large-cnn",
      )
      parser.add_argument(
          "--device", type=str, required=False, default=DEFAULT_DEVICE, help="cuda, cuda:1, cpu etc.",
      )
      parser.add_argument(
          "--bs", type=int, default=8, required=False, help="batch size: how many to summarize at a time",
      )
      args = parser.parse_args()
      examples = [" " + x.rstrip() for x in open(args.source_path).readlines()]
      generate_summaries(examples, args.output_path,
                         args.model_name, batch_size=args.bs, device=args.device)
  
  
  if __name__ == "__main__":
      run_generate()
  