text: from argparse import ArgumentParser
  from pathlib import Path
  import os
  
  import torch
  from tokenizers import Tokenizer
  from transformers import GPT2LMHeadModel, GPT2TokenizerFast
  from tqdm import tqdm
  
  
  def generate(token_ids, model, tokenizer):
      input_ids = torch.LongTensor(token_ids).unsqueeze(0).cuda()
  
      batch_ids = model.generate(input_ids=input_ids,
                                 num_return_sequences=1,
                                 max_length=100,
                                 do_sample=True,
                                 top_k=20,
                                 top_p=0.9,
                                 temperature=2.0,
                                 repetition_penalty=10.,
                                 num_beams=10,
                                 no_repeat_ngram_size=4)
  
      ids = batch_ids[0].flatten().tolist()
      txt = f'[{tokenizer.decode(token_ids).strip()}]{tokenizer.decode(ids[len(token_ids):], skip_special_tokens=True)}'.replace(
          '\n\n', '\n')
      # txt = tokenizer.decode(ids, skip_special_tokens=True).strip()
  
      tqdm.write(f' > {txt}\n')
      return txt + '\n\n'
  
  
  def main():
      parser = ArgumentParser()
      parser.add_argument('lang', choices=['nld', 'ita'])
      parser.add_argument('models', nargs='+')
      parser.add_argument('--src', default='small', choices=['full', 'small'])
      parser.add_argument('--file', default='full')
      parser.add_argument('-n', default=5, type=int)
      parser.add_argument('-f', '--force', action='store_true')
      args = parser.parse_args()
  
      base_path = Path(
          'data') / args.lang / 'evaluation' / 'examples' / args.src / args.file
  
      src_path = base_path / 'gold.txt'
      if not src_path.exists():
          print(f' > gold path {src_path} does not exist')
          exit(1)
  
      print(' > loading tokenizer')
      os.environ['TOKENIZERS_PARALLELISM'] = 'false'
      if args.lang == 'ita':
          tokenizer = GPT2TokenizerFast.from_pretrained(
              'LorenzoDeMattei/GePpeTto')
      else:
          tokenizer_path = Path(
              'data') / args.lang / 'vocabularies' / 'tokenizer.json'
          tokenizer = Tokenizer.from_file(str(tokenizer_path))
          args.n += 1
  
      print(f' > loading examples from {src_path}')
      examples = []
      with open(src_path) as f:
          for line in f:
              token_ids = tokenizer.encode(line.strip())
              if type(token_ids) != list:
                  token_ids = [0] + token_ids.ids
              examples.append(token_ids[:args.n])
      print(f' > loaded {len(examples)} examples')
  
      for model_name in args.models:
          tgt_path = base_path / f'{model_name.replace("/", "_")}.txt'
          if not args.force and tgt_path.exists():
              print(f'{tgt_path} already exists. skipping')
              continue
  
          model_path = Path('data') / args.lang / 'models' / model_name
          if not model_path.exists():
              model_path = model_name
  
          print(f' > loading model {model_path}')
          model = GPT2LMHeadModel.from_pretrained(model_path).cuda()
          model.eval()
  
          print(' > generating endings for examples')
          generated = [
              generate(input_ids, model, tokenizer)
              for input_ids in tqdm(examples, ncols=80)
          ]
          with open(tgt_path, 'w') as f:
              f.writelines(generated)
  
          print(f'\nsaved to {tgt_path}')
  
  
  if __name__ == '__main__':
      main()
  