class_docstr: Handles sampling, eval, and training with shared queries.
class_docstr_tok: ["Handles","sampling",",","eval",",","and","training","with","shared","queries","."]
docstr: :return: A dict with structure:
      samples: [batch, num_responses, sample_len]
      logprobs: [batch, num_responses, sample_len]
      [head]: {
          response: [batch, num_responses, sample_len+1]
      }
code: def _sample(
      self,
      context_tokens,
      sample_len,
      partial_responses=None,
      responses_per_query=1,
      sample_H=None,
      **model_call_kwargs,
  ):
      """
      :return: A dict with structure:
          samples: [batch, num_responses, sample_len]
          logprobs: [batch, num_responses, sample_len]
          [head]: {
              response: [batch, num_responses, sample_len+1]
          }
      """
      context_tokens = context_tokens.to(self.device)
      self.model.eval()
      n_batch, query_length = context_tokens.size()
      assert self.logit_head, f"Cannot sample without logit_head"
      # NOTE: could do this more efficiently by sharing context work
      repeated_context_tokens = context_tokens.unsqueeze(1).repeat(
          1, responses_per_query, 1
      )
  
      # Combine query and response so far into new query to be passed to _sample()
      if partial_responses is not None:
          partial_responses = partial_responses.to(self.device)
          repeated_context_tokens = torch.cat(
              (repeated_context_tokens, partial_responses), 2
          )
  
      sample_fn = _get_sample_fn(sample_H)
  
      flat_context_tokens = repeated_context_tokens.flatten(0, 1)
      flat_n_batch, context_len = flat_context_tokens.shape
  
      assert (
          sample_len + context_len <= self.model_hparams["n_ctx"] + 1
      ), (
          f"Requested completion {sample_len} is too long for"
          f"context {context_len} and model context_len {self.model_hparams.n_ctx}"
      )
  
      results = sample(
          self.model,
          flat_context_tokens,
          sample_len=sample_len,
          sample_fn=sample_fn,
          model_output_keys=self.heads,
          **model_call_kwargs,
      )
  
      samples = results["tokens"]
      logprobs = results["logprobs"]
      assert samples.size(-2) == n_batch * responses_per_query
      assert logprobs.size(-2) == n_batch * responses_per_query
      assert (
          samples.size(-1) == sample_len
      ), f"{samples.size()} vs {sample_len}"
      assert (
          logprobs.size(-1) == sample_len
      ), f"{logprobs.size()} vs {sample_len}"
      samples = samples.view(n_batch, responses_per_query, sample_len)
      logprobs = logprobs.view(n_batch, responses_per_query, sample_len)
  
      output = dict(
          contexts=context_tokens, samples=samples, logprobs=logprobs
      )
  
      mask, _ = _zero_padding_tokens(output["samples"])
      # NOTE: sample doesn't return eval'ed values on final token
      mask = mask[:, :, :-1]
      for k in self.heads:
          reshaped = results[k].view(
              n_batch, responses_per_query, *results[k].shape[1:]
          )
          output[k] = _split_query_response_output_parts(
              reshaped, query_length, mask
          )
      return output
  
imports: import json
  import os
  import re
  import shutil
  import uuid
  from dataclasses import dataclass, field
  from typing import Callable, Optional, Set, List, Dict
  import blobfile
  import numpy
  import torch
  from mpi4py import MPI
  import summarize_from_feedback
  from summarize_from_feedback.model_layout import ModelLayout
  from summarize_from_feedback.models import sample_fns
  from summarize_from_feedback.models.loss_functions import softmax_xent_loss_fn
  from summarize_from_feedback.models.transformer import Hyperparams
  from summarize_from_feedback.models.transformer import build_with_random_weights
  from summarize_from_feedback.utils import exact_div, hyperparams
  from summarize_from_feedback.utils import blobs
  from summarize_from_feedback.utils.dist_utils import setup_cuda_device_and_dist, create_data_parallel_comm, create_within_replica_comm
  from summarize_from_feedback.utils.nested import map_nested
  from summarize_from_feedback.utils.torch_utils import nans, to_numpy
label_desc: Use image node_id 7 for calling the QueryResponseModel obj's underlying member method code with example usage: obj._sample(context_tokens, sample_len, partial_responses, responses_per_query, sample_H) and returns: output