# Path: prompt_tuning/train/prompts.py
# def prefix_prompt(prompt: Array, x_embed: Array, x: Array) -> Array:
# def prefix_prompt_after_bos(prompt: Array, x_embed: Array, x: Array) -> Array:
# def suffix_prompt(prompt: Array,
#                   x_embed: Array,
#                   x: Array,
#                   decoder_only: bool = False,
#                   before_eos: bool = True) -> Array:
#   def __call__(self, x, x_embed):
# class Prompt(nn.Module):


from typing import Callable
from prompt_tuning.train import prompts
from flaxformer.types import Array
import flax.linen as nn
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Training subclasses of our prompt that actually do the concatenation.

This also creates a unified API between single task prompts, where the prompt is
unbatched, and multi-task prompts, where the input needs to have the first token
removed.
"""



class MultiTaskPrompt(nn.Module):
  """Geneate a MultiTaskPrompt and concatenate it with the input.

  This is the training time version of prompting a model. Calling the injected
  `prompt` module will generate your prompt. This prompt should be batched. This
  module then slices off the task index from the input and concatenates the
  prompt. This can be used in conjunction with the `multitask=True` arguments
  for attention mask creation to do multi-task prompting without need multitask
  subclasses of various flaxformer modules.

  Attributes:
    prompt: The model that actually generates the batched prompt.
    combine: A function that combines the prompt and the embedded input.
  """
  prompt: nn.Module

# Path: prompt_tuning/train/prompts.py
# def prefix_prompt(prompt: Array, x_embed: Array, x: Array) -> Array:
# def prefix_prompt_after_bos(prompt: Array, x_embed: Array, x: Array) -> Array:
# def suffix_prompt(prompt: Array,
#                   x_embed: Array,
#                   x: Array,
#                   decoder_only: bool = False,
#                   before_eos: bool = True) -> Array:
#   def __call__(self, x, x_embed):
# class Prompt(nn.Module):


# Path: prompt_tuning/extended/train/multitask_prompts.py
from typing import Callable
from prompt_tuning.train import prompts
from flaxformer.types import Array
import flax.linen as nn

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Training subclasses of our prompt that actually do the concatenation.

This also creates a unified API between single task prompts, where the prompt is
unbatched, and multi-task prompts, where the input needs to have the first token
removed.
"""



class MultiTaskPrompt(nn.Module):
  """Geneate a MultiTaskPrompt and concatenate it with the input.

  This is the training time version of prompting a model. Calling the injected
  `prompt` module will generate your prompt. This prompt should be batched. This
  module then slices off the task index from the input and concatenates the
  prompt. This can be used in conjunction with the `multitask=True` arguments
  for attention mask creation to do multi-task prompting without need multitask
  subclasses of various flaxformer modules.

  Attributes:
    prompt: The model that actually generates the batched prompt.
    combine: A function that combines the prompt and the embedded input.
  """
  prompt: nn.Module

  combine: Callable[[Array, Array, Array], Array] = prompts.prefix_prompt