code: def format_input_with_prompt(
      input_col_name: str,
      dataset_df: DataFrame,
      backend: "Backend",
      task_str: str,
      retrieval_model: Optional[RetrievalModel] = None,
      k: int = -1,
      template: Optional[str] = None,
  ) -> Series:
      """Returns a new Series with the input column data formatted with the prompt.
  
      A prompt can either be zero-shot or few-shot. A zero-shot prompt is comprised of some (unlabeled) input and a task
      to be completed given the input. A few-shot prompt additionally includes some dynamically retrieved context, which
      is retrieved using the `retrieval_model.search` function.
  
      A template can be provided to customize the prompt. The template must be a string with the following fields:
          - __sample__ or at least one column from the input dataset: The input sample.
          - __context__: The context retrieved by the `search_fn` function. Only required if `search_fn` is provided.
          - __task__: The task to be completed given the input. Only required if `task` is set in the prompt config.
  
      Zero-shot example:
  
      Before formatting:
  
          input_col = ["I am happy"]
          task_str = "sentiment analysis"
  
      After formatting:
  
          input_col = ["SAMPLE INPUT: I am happy\n\nUSER: Complete the following task: sentiment analysis\n\nASSISTANT:"]
  
      Args:
          input_col_name (str): The name of the input column.
          dataset_df (DataFrame): The input dataset.
          backend (Backend): The backend used for map operations.
          task_str (str): The task to be completed given the input.
          retrieval_model (Optional[RetrievalModel]): The retrieval model used to retrieve context. If provided, the
              prompt will be few-shot. If not provided, the prompt will be zero-shot.
          k (int): The number of samples to retrieve. Only required if `retrieval_model` is provided.
          template (Optional[str]): The template to use for the prompt. If not provided, the default will be used.
  
      Returns:
          Series: A new Series with the input column data formatted with the prompt.
      """
      # determine if this is a few-shot or zero-shot prompt
      # few-shot prompts require a search function that returns samples from some dataset
      is_few_shot = retrieval_model is not None
  
      # if no template is provided, use the default template
      if template is None:
          if is_few_shot:
              template = DEFAULT_FEW_SHOT_PROMPT_TEMPLATE
          else:
              template = DEFAULT_ZERO_SHOT_PROMPT_TEMPLATE
  
      # ensure that the prompt template has all required fields
      template_fields, field_to_dtype = _get_template_fields(template)
      try:
          _validate_prompt_template(
              template_fields,
              task_str,
              is_few_shot,
              dataset_df.columns,
              input_col_name,
          )
      except ValueError as e:
          raise ValueError(
              f"template invalid for {'few-shot' if is_few_shot else 'zero-shot'} prompt: {e}"
          )
  
      def generate_prompt(df: pd.DataFrame):
          if CONTEXT in template_fields:
              df[CONTEXT] = retrieval_model.search(
                  df, backend, k=k, return_data=True
              )
          if SAMPLE in template_fields:
              # During preprocessing, we're inserting quotes that change the token IDs completely if we
              # don't remove the " from the string. For parity with expected user output, we need to get rid of them.
              # TODO(Arnav): see if there's a way to only remove them if the entry does't have quotes. This currently
              # removes all " from the string (even those not added by json.dumps), which is not ideal.
              df[SAMPLE] = df[input_col_name].map(
                  lambda entry: json.dumps(entry, indent=2).strip('"')
              )
          if TASK in template_fields:
              df[TASK] = task_str
  
          def generate_prompt_for_row(row):
              kwargs = {
                  col: field_to_dtype[col](row[col])
                  for col in template_fields
              }
              return template.format(**kwargs)
  
          return df.apply(generate_prompt_for_row, axis=1)
  
      result = backend.df_engine.map_partitions(
          dataset_df, generate_prompt, meta=(input_col_name, "object")
      )
      result = backend.df_engine.persist(
          result
      )  # persist to prevent re-computation
      return result
  
file_path: ludwig/ludwig/data/prompt.py
imports: import json
  import logging
  import os
  import string
  from typing import Any, Dict, List, Optional, Set, Tuple, Type, TYPE_CHECKING
  import pandas
  from ludwig.models.retrieval import df_checksum, get_retrieval_model, RetrievalModel
  from ludwig.utils.fs_utils import get_default_cache_location, makedirs, path_exists
  from ludwig.utils.types import DataFrame, Series
label_desc: Use image node_id 2 for calling a global function with example usage: format_input_with_prompt(input_col_name, dataset_df, backend, task_str, retrieval_model, k, template) and returns: result, df, template