def prep_next_item_dataset(
    k,
    item_sequence,
    user_negative_items,
    candidate_sample_prob=None,
):
  """Converts raw user-item dataset to a dataset with a dictionary of features.

  The produced dataset uses the kth item from the end as the next item to be
  predicted and items before that as the input sequence. Typically, `k` will be
  1 or 2 depending on the data split_type (user/step). When using validation and
  split_type=step, `k` will 3 for training data, 2 for validation data and 1 for
  testing data.

  Args:
    k: Integer index denoting the kth item from the end.
    item_sequence: The complete user item history.
    user_negative_items: Negative item samples used for evaluation. See
      evaluation in the proposal doc
      (http://shortn/_PO6OdvUuAs#bookmark=id.4dm5yebd3f21) for details.
    candidate_sample_prob: The candidate sampling probability. This is used for
      correcting the sample bias of in-batch negatives.

  Returns:
    A dictionary containing the following features in the dataset:
    user_item_sequence: List of user item sequences.
    next_item: List of next items for each user.
    user_negative_items: List of negative items for each user. This is used to
      compute metrics for real datasets.
    candidate_sample_prob: The candidate sampling probability.
  """

  user_item_input = {
      'user_item_sequence': item_sequence[:-k],
      'next_item': item_sequence[-k],
  }

  if user_negative_items is not None:
    user_item_input['user_negative_items'] = user_negative_items
  if candidate_sample_prob is not None:
    user_item_input['candidate_sampling_probability'] = candidate_sample_prob

  return user_item_input