code: def top_k_filter(logits, k, probs=False):
      """
      Masks everything but the k top entries as -infinity (1e10).
      Used to mask logits such that e^-infinity -> 0 won't contribute to the
      sum of the denominator.
      """
      if k == 0:
          return logits
      else:
          values = torch.topk(logits, k)[0]
          batch_mins = values[:, -1].view(-1, 1).expand_as(logits)
          if probs:
              return torch.where(
                  logits < batch_mins,
                  torch.ones_like(logits) * 0.0,
                  logits,
              )
          return torch.where(
              logits < batch_mins,
              torch.ones_like(logits) * -BIG_CONST,
              logits,
          )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/examples/research_projects/pplm/run_pplm.py
imports: import argparse
  import json
  from operator import add
  from typing import List, Optional, Tuple, Union
  import numpy
  import torch
  from torch import nn
  from tqdm import trange
  from pplm_classification_head import ClassificationHead
  from transformers import GPT2LMHeadModel, GPT2Tokenizer
  from transformers.file_utils import cached_path
label_desc: Use image node_id 1 for calling a global function with example usage: top_k_filter(logits, k, probs) and returns: logits, torch, torch