sparknlp.base.prompt_assembler#

Classes#

class PromptAssembler[source]#

Assembles a sequence of messages into a single string using a template. These strings can then

be used as prompts for large language models.

This annotator expects an array of two-tuples as the type of the input column (one array of tuples per row). The first element of the tuples should be the role and the second element is the text of the message. Possible roles are “system”, “user” and “assistant”.

An assistant header can be added to the end of the generated string by using setAddAssistant(True).

At the moment, this annotator uses llama.cpp as a backend to parse and apply the templates. llama.cpp uses basic pattern matching to determine the type of the template, then applies a basic version of the template to the messages. This means that more advanced templates are not supported.

For an extended example see the example notebook.

Input Annotation types	Output Annotation type
NONE	DOCUMENT

Parameters:

inputCol

Input column name

outputCol

Output column name

chatTemplate

Template used for the chat

addAssistant

Whether to add an assistant header to the end of the generated string

You are a helpful assistant.<|eot_id|><|start_header_id|>assistant<|end_header_id|>

Hello there, how can I help you?<|eot_id|><|start_header_id|>user<|end_header_id|>

I need help with organizing my room.<|eot_id|><|start_header_id|>assistant<|end_header_id|>

]|

setInputCol(value)[source]#

Sets input column name.

Parameters:

valuestr

Name of the input column

setOutputCol(value)[source]#

Sets output column name.

Parameters:

valuestr

Name of the Output Column

setChatTemplate(value)[source]#

Sets the chat template.

Parameters:

valuestr

Template used for the chat

setAddAssistant(value)[source]#

Sets whether to add an assistant header to the end of the generated string.

Parameters:

valuebool

Whether to add an assistant header to the end of the generated string