Skip to content

Huggingface pipeline prompt driver

HuggingFacePipelinePromptDriver

Bases: BasePromptDriver

Attributes:

Name Type Description
params dict

Custom model run parameters.
model str

Hugging Face Hub model name.
Source code in griptape/drivers/prompt/huggingface_pipeline_prompt_driver.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@define
class HuggingFacePipelinePromptDriver(BasePromptDriver):
    """
    Attributes:
        params: Custom model run parameters.
        model: Hugging Face Hub model name.

    """

    max_tokens: int = field(default=250, kw_only=True, metadata={"serializable": True})
    model: str = field(kw_only=True, metadata={"serializable": True})
    params: dict = field(factory=dict, kw_only=True, metadata={"serializable": True})
    tokenizer: HuggingFaceTokenizer = field(
        default=Factory(
            lambda self: HuggingFaceTokenizer(model=self.model, max_output_tokens=self.max_tokens), takes_self=True
        ),
        kw_only=True,
    )
    pipe: TextGenerationPipeline = field(
        default=Factory(
            lambda self: import_optional_dependency("transformers").pipeline(
                "text-generation", model=self.model, max_new_tokens=self.max_tokens, tokenizer=self.tokenizer.tokenizer
            ),
            takes_self=True,
        )
    )

    def try_run(self, prompt_stack: PromptStack) -> TextArtifact:
        messages = [self._prompt_stack_input_to_message(input) for input in prompt_stack.inputs]

        result = self.pipe(
            messages,
            max_new_tokens=self.max_tokens,
            tokenizer=self.tokenizer.tokenizer,
            stop_strings=self.tokenizer.stop_sequences,
            temperature=self.temperature,
            do_sample=True,
        )

        if isinstance(result, list):
            if len(result) == 1:
                generated_text = result[0]["generated_text"][-1]["content"]

                return TextArtifact(value=generated_text)
            else:
                raise Exception("completion with more than one choice is not supported yet")
        else:
            raise Exception("invalid output format")

    def try_stream(self, prompt_stack: PromptStack) -> Iterator[TextArtifact]:
        raise NotImplementedError("streaming is not supported")

    def prompt_stack_to_string(self, prompt_stack: PromptStack) -> str:
        return self.tokenizer.tokenizer.decode(self.__prompt_stack_to_tokens(prompt_stack))

    def _prompt_stack_input_to_message(self, prompt_input: PromptStack.Input) -> dict:
        return {"role": prompt_input.role, "content": prompt_input.content}

    def __prompt_stack_to_tokens(self, prompt_stack: PromptStack) -> list[int]:
        tokens = self.tokenizer.tokenizer.apply_chat_template(
            [self._prompt_stack_input_to_message(i) for i in prompt_stack.inputs],
            add_generation_prompt=True,
            tokenize=True,
        )

        if isinstance(tokens, list):
            return tokens
        else:
            raise ValueError("Invalid output type.")

max_tokens: int = field(default=250, kw_only=True, metadata={'serializable': True}) class-attribute instance-attribute

model: str = field(kw_only=True, metadata={'serializable': True}) class-attribute instance-attribute

params: dict = field(factory=dict, kw_only=True, metadata={'serializable': True}) class-attribute instance-attribute

pipe: TextGenerationPipeline = field(default=Factory(lambda self: import_optional_dependency('transformers').pipeline('text-generation', model=self.model, max_new_tokens=self.max_tokens, tokenizer=self.tokenizer.tokenizer), takes_self=True)) class-attribute instance-attribute

tokenizer: HuggingFaceTokenizer = field(default=Factory(lambda self: HuggingFaceTokenizer(model=self.model, max_output_tokens=self.max_tokens), takes_self=True), kw_only=True) class-attribute instance-attribute

__prompt_stack_to_tokens(prompt_stack)

Source code in griptape/drivers/prompt/huggingface_pipeline_prompt_driver.py 
74
75
76
77
78
79
80
81
82
83
84
def __prompt_stack_to_tokens(self, prompt_stack: PromptStack) -> list[int]:
    tokens = self.tokenizer.tokenizer.apply_chat_template(
        [self._prompt_stack_input_to_message(i) for i in prompt_stack.inputs],
        add_generation_prompt=True,
        tokenize=True,
    )

    if isinstance(tokens, list):
        return tokens
    else:
        raise ValueError("Invalid output type.")

prompt_stack_to_string(prompt_stack)

Source code in griptape/drivers/prompt/huggingface_pipeline_prompt_driver.py 
68
69
def prompt_stack_to_string(self, prompt_stack: PromptStack) -> str:
    return self.tokenizer.tokenizer.decode(self.__prompt_stack_to_tokens(prompt_stack))

try_run(prompt_stack)

Source code in griptape/drivers/prompt/huggingface_pipeline_prompt_driver.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def try_run(self, prompt_stack: PromptStack) -> TextArtifact:
    messages = [self._prompt_stack_input_to_message(input) for input in prompt_stack.inputs]

    result = self.pipe(
        messages,
        max_new_tokens=self.max_tokens,
        tokenizer=self.tokenizer.tokenizer,
        stop_strings=self.tokenizer.stop_sequences,
        temperature=self.temperature,
        do_sample=True,
    )

    if isinstance(result, list):
        if len(result) == 1:
            generated_text = result[0]["generated_text"][-1]["content"]

            return TextArtifact(value=generated_text)
        else:
            raise Exception("completion with more than one choice is not supported yet")
    else:
        raise Exception("invalid output format")

try_stream(prompt_stack)

Source code in griptape/drivers/prompt/huggingface_pipeline_prompt_driver.py 
65
66
def try_stream(self, prompt_stack: PromptStack) -> Iterator[TextArtifact]:
    raise NotImplementedError("streaming is not supported")