class_docstr: Represents a RL policy + value function, containing a value and logit head.
  Only returns pre-token values (never evaluates on the final response token)
docstr: Samples from the policy given the context (query_tokens). partial_responses, if provided,
  should be a torch tensor of size (batch_size, responses_per_query, X) where X = response length so far.
  :return: A dict with structure:
      tokens: [batch, num_responses, sample_len]
      logprobs: [batch, num_responses, sample_len]
      value: [batch, num_responses, sample_len]
docstr_tok: ["Samples","from","the","policy","given","the","context","(","query_tokens",")",".","partial_responses",",","if","provided",",","should","be","a","torch","tensor","of","size","(","batch_size",",","responses_per_query",",","X",")","where","X","=","response","length","so","far",".",":","return",":","A","dict","with","structure",":","tokens",":","[","batch",",","num_responses",",","sample_len","]","logprobs",":","[","batch",",","num_responses",",","sample_len","]","value",":","[","batch",",","num_responses",",","sample_len","]"]
code: def sample(
      self,
      query_tokens,
      partial_responses=None,
      responses_per_query=1,
      sample_len=None,
      **kwargs,
  ):
      """
      Samples from the policy given the context (query_tokens). partial_responses, if provided,
      should be a torch tensor of size (batch_size, responses_per_query, X) where X = response length so far.
      :return: A dict with structure:
          tokens: [batch, num_responses, sample_len]
          logprobs: [batch, num_responses, sample_len]
          value: [batch, num_responses, sample_len]
      """
  
      response_len = 0  # length of responses so far
      if partial_responses is not None:
          response_len = partial_responses.size(2)
  
      assert self.logit_head
  
      # sample_len is length of sample to be completed
      if sample_len is None:
          assert self.task_hparams is not None
          sample_len = self.task_hparams.response.length - response_len
  
      if self.task_hparams is not None:
          assert (
              query_tokens.size(1) == self.task_hparams.query.length
          ), f"{query_tokens.size()}"
  
      ret = self._sample(
          query_tokens,
          partial_responses=partial_responses,
          sample_len=sample_len,
          responses_per_query=responses_per_query,
          **kwargs,
      )
      for key in self.heads:
          ret[key] = ret[key]["response"]
  
      return ret
  
file_path: summarize-from-feedback/summarize_from_feedback/policy.py
imports: import functools
  from summarize_from_feedback import tasks
  from summarize_from_feedback.query_response_model import QueryResponseModel
label_desc: Use image node_id 2 for calling the Policy obj's underlying member method code with example usage: obj.sample(query_tokens, partial_responses, responses_per_query, sample_len) and returns: ret