code: def prepare_ud_dataset(
      treebank,
      udbase_dir,
      tokenizer_dir,
      short_name,
      short_language,
      dataset,
      augment=True,
      input_conllu=None,
      output_conllu=None,
  ):
      if input_conllu is None:
          input_conllu = common.find_treebank_dataset_file(
              treebank, udbase_dir, dataset, "conllu", fail=True
          )
      if output_conllu is None:
          output_conllu = common.tokenizer_conllu_name(
              tokenizer_dir, short_name, dataset
          )
      print(
          "Reading from %s and writing to %s"
          % (input_conllu, output_conllu)
      )
  
      if short_name == "te_mtg" and dataset == "train" and augment:
          write_augmented_dataset(
              input_conllu, output_conllu, augment_telugu
          )
      elif short_name == "ar_padt" and dataset == "train" and augment:
          write_augmented_dataset(
              input_conllu, output_conllu, augment_arabic_padt
          )
      elif short_name.startswith("ko_") and short_name.endswith("_seg"):
          remove_spaces(input_conllu, output_conllu)
      elif dataset == "train" and augment:
          write_augmented_dataset(
              input_conllu, output_conllu, augment_punct
          )
      else:
          sents = read_sentences_from_conllu(input_conllu)
          write_sentences_to_conllu(output_conllu, sents)
  
imports: import argparse
  import glob
  import os
  import random
  import re
  import tempfile
  from collections import Counter
  from stanza.models.common.constant import treebank_to_short_name
  import stanza.utils.datasets.common
  from stanza.utils.datasets.common import read_sentences_from_conllu, write_sentences_to_conllu, INT_RE, MWT_RE, MWT_OR_COPY_RE
  import stanza.utils.datasets.tokenization.convert_ml_cochin
  import stanza.utils.datasets.tokenization.convert_my_alt
  import stanza.utils.datasets.tokenization.convert_vi_vlsp
  import stanza.utils.datasets.tokenization.convert_th_best
  import stanza.utils.datasets.tokenization.convert_th_lst20
  import stanza.utils.datasets.tokenization.convert_th_orchid
label_desc: Use image node_id 37 for calling a global function with example usage: prepare_ud_dataset(treebank, udbase_dir, tokenizer_dir, short_name, short_language, dataset, augment, input_conllu, output_conllu) without return types