code: def main(args):
      assert (
          args.split.startswith("train") or args.num_shards == 1
      ), "--num-shards should be set to 1 for valid and test sets"
      assert (
          args.split.startswith("train")
          or args.split.startswith("valid")
          or args.split.startswith("test")
      ), "--split should be set to train[n]/valid[n]/test[n]"
  
      source_sents = read_text_file(args.input_source)
      target_sents = read_text_file(args.input_target)
  
      num_sents = len(source_sents)
      assert num_sents == len(
          target_sents
      ), f"{args.input_source} and {args.input_target} should have the same number of sentences."
  
      hypo_sents = read_text_file(args.input_hypo)
      assert (
          len(hypo_sents) % args.beam == 0
      ), f"Number of hypotheses ({len(hypo_sents)}) cannot be divided by beam size ({args.beam})."
  
      hypo_sents = [
          hypo_sents[i : i + args.beam]
          for i in range(0, len(hypo_sents), args.beam)
      ]
      assert num_sents == len(
          hypo_sents
      ), f"{args.input_hypo} should contain {num_sents * args.beam} hypotheses but only has {len(hypo_sents) * args.beam}. (--beam={args.beam})"
  
      output_dir = args.output_dir / args.metric
      for ns in range(args.num_shards):
          print(f"processing shard {ns+1}/{args.num_shards}")
          shard_output_dir = output_dir / f"split{ns+1}"
          source_output_dir = shard_output_dir / "input_src"
          hypo_output_dir = shard_output_dir / "input_tgt"
          metric_output_dir = shard_output_dir / args.metric
  
          source_output_dir.mkdir(parents=True, exist_ok=True)
          hypo_output_dir.mkdir(parents=True, exist_ok=True)
          metric_output_dir.mkdir(parents=True, exist_ok=True)
  
          if args.n_proc > 1:
              with Pool(
                  args.n_proc,
                  initializer=init,
                  initargs=(args.sentencepiece_model,),
              ) as p:
                  output = p.starmap(
                      process,
                      [
                          (
                              source_sents[i],
                              target_sents[i],
                              hypo_sents[i],
                              args.metric,
                          )
                          for i in range(ns, num_sents, args.num_shards)
                      ],
                  )
          else:
              init(args.sentencepiece_model)
              output = [
                  process(
                      source_sents[i],
                      target_sents[i],
                      hypo_sents[i],
                      args.metric,
                  )
                  for i in range(ns, num_sents, args.num_shards)
              ]
  
          with open(
              source_output_dir / f"{args.split}.bpe", "w"
          ) as s_o, open(
              hypo_output_dir / f"{args.split}.bpe", "w"
          ) as h_o, open(
              metric_output_dir / f"{args.split}.{args.metric}", "w"
          ) as m_o:
              for source_bpe, hypo_bpe, score_str in output:
                  assert len(hypo_bpe) == len(score_str)
                  for h, m in zip(hypo_bpe, score_str):
                      s_o.write(f"{source_bpe}\n")
                      h_o.write(f"{h}\n")
                      m_o.write(f"{m}\n")
  
file_path: fairseq/examples/discriminative_reranking_nmt/scripts/prep_data.py
imports: import argparse
  from multiprocessing import Pool
  from pathlib import Path
  import sacrebleu
  import sentencepiece
label_desc: Use image node_id 6 for calling a global function with example usage: main(args) without return types