docstr: Evaluate the model on test or valid datasets
docstr_tok: ["Evaluate","the","model","on","test","or","valid","datasets"]
code: def test_step(
      model, test_triples, args, entity_dict, random_sampling=False
  ):
      """
      Evaluate the model on test or valid datasets
      """
  
      model.eval()
  
      # Prepare dataloader for evaluation
      test_dataloader_head = DataLoader(
          TestDataset(
              test_triples,
              args,
              "head-batch",
              random_sampling,
              entity_dict,
          ),
          batch_size=args.test_batch_size,
          num_workers=max(1, args.cpu_num // 2),
          collate_fn=TestDataset.collate_fn,
      )
  
      test_dataloader_tail = DataLoader(
          TestDataset(
              test_triples,
              args,
              "tail-batch",
              random_sampling,
              entity_dict,
          ),
          batch_size=args.test_batch_size,
          num_workers=max(1, args.cpu_num // 2),
          collate_fn=TestDataset.collate_fn,
      )
  
      test_dataset_list = [test_dataloader_head, test_dataloader_tail]
  
      test_logs = defaultdict(list)
  
      step = 0
      total_steps = sum([len(dataset) for dataset in test_dataset_list])
  
      with torch.no_grad():
          for test_dataset in test_dataset_list:
              for (
                  positive_sample,
                  negative_sample,
                  mode,
              ) in test_dataset:
                  if args.cuda:
                      positive_sample = positive_sample.cuda()
                      negative_sample = negative_sample.cuda()
  
                  batch_size = positive_sample.size(0)
                  score = model(
                      (positive_sample, negative_sample), mode
                  )
  
                  batch_results = model.evaluator.eval(
                      {
                          "y_pred_pos": score[:, 0],
                          "y_pred_neg": score[:, 1:],
                      }
                  )
                  for metric in batch_results:
                      test_logs[metric].append(batch_results[metric])
  
                  if step % args.test_log_steps == 0:
                      logging.info(
                          "Evaluating the model... (%d/%d)"
                          % (step, total_steps)
                      )
  
                  step += 1
  
          metrics = {}
          for metric in test_logs:
              metrics[metric] = (
                  torch.cat(test_logs[metric]).mean().item()
              )
  
      return metrics
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import logging
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.utils.data import DataLoader
  from dataloader import TestDataset
  from collections import defaultdict
  from ogb.linkproppred import Evaluator
label_desc: Use image node_id 8 for calling the KGEModel obj's underlying member method code with example usage: obj.test_step(model, test_triples, args, entity_dict, random_sampling) and returns: metrics