Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Testing streaming outputs #852

Draft
wants to merge 6 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions cookbooks/Gradio/hf_model_parsers.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
from aiconfig_extension_hugging_face import (
HuggingFaceAutomaticSpeechRecognitionTransformer, # Haven't tested yet

HuggingFaceImage2TextTransformer, # Tested, model doesn't support streaming
HuggingFaceTextSummarizationTransformer, # Tested
HuggingFaceText2ImageDiffusor, # Tested, model doesn't support streaming
HuggingFaceText2SpeechTransformer, # Tested, model doesn't support streaming
HuggingFaceTextGenerationTransformer, # Tested
HuggingFaceTextTranslationTransformer, # Tested
)
from aiconfig import (AIConfigRuntime, ModelParserRegistry)

def register_model_parsers() -> None:
"""Register model parsers for HuggingFace models.
"""
# Audio --> Text
# AIConfigRuntime.register_model_parser(HuggingFaceAutomaticSpeechRecognitionTransformer(), "AutomaticSpeechRecognition")

# # Image --> Text
# AIConfigRuntime.register_model_parser(HuggingFaceImage2TextTransformer(), "Image2Text")

# # Text --> Image
# AIConfigRuntime.register_model_parser(HuggingFaceText2ImageDiffusor(), "Text2Image")

# # Text --> Audio
# AIConfigRuntime.register_model_parser(HuggingFaceText2SpeechTransformer(), "Text2Speech")

# # Text --> Text
# AIConfigRuntime.register_model_parser(HuggingFaceTextGenerationTransformer(), "TextGeneration")
# AIConfigRuntime.register_model_parser(HuggingFaceTextSummarizationTransformer(), "TextSummarization")
# ModelParserRegistry.register_model_parser(HuggingFaceText2SpeechTransformer())
# ModelParserRegistry.register_model_parser(HuggingFaceTextGenerationTransformer())
ModelParserRegistry.register_model_parser(HuggingFaceImage2TextTransformer())

# ModelParserRegistry.register_model_parser(HuggingFaceAutomaticSpeechRecognitionTransformer())
# ModelParserRegistry.register_model_parser(HuggingFaceTextSummarizationTransformer())
# ModelParserRegistry.register_model_parser(HuggingFaceText2ImageDiffusor())
# ModelParserRegistry.register_model_parser(HuggingFaceTextTranslationTransformer())
# AIConfigRuntime.register_model_parser(HuggingFaceTextTranslationTransformer(), "TextTranslation")
75 changes: 75 additions & 0 deletions cookbooks/Gradio/huggingface.aiconfig.json

Large diffs are not rendered by default.

5 changes: 3 additions & 2 deletions extensions/HuggingFace/python/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,12 @@ huggingface_hub

#Hugging Face Libraries - Local Inference Tranformers & Diffusors
accelerate # Used to help speed up image generation
diffusers # Used for image + audio generation
diffusers # Used for image generation
scipy # array -> wav file, text-speech. torchaudio.save seems broken.
sentencepiece # Used for text translation
torch
torchvision
torchaudio
scipy # array -> wav file, text-speech. torchaudio.save seems broken.
transformers # Used for text generation

#Other
Expand Down
Original file line number Diff line number Diff line change
@@ -1,11 +1,24 @@
from typing import Any, Dict, Optional, List, TYPE_CHECKING
import base64
import json
from io import BytesIO
from PIL import Image
from typing import Any, Dict, Optional, List, TYPE_CHECKING, Union
from transformers import (
Pipeline,
pipeline,
)

from aiconfig import ParameterizedModelParser, InferenceOptions
from aiconfig.callback import CallbackEvent
import torch
from aiconfig.schema import Prompt, Output, ExecuteResult, Attachment

from transformers import pipeline, Pipeline

from aiconfig.schema import (
Attachment,
ExecuteResult,
Output,
OutputDataWithValue,
Prompt,
)

# Circular Dependency Type Hints
if TYPE_CHECKING:
from aiconfig import AIConfigRuntime

Expand Down Expand Up @@ -93,10 +106,11 @@ async def deserialize(
await aiconfig.callback_manager.run_callbacks(CallbackEvent("on_deserialize_start", __name__, {"prompt": prompt, "params": params}))

# Build Completion data
completion_params = self.get_model_settings(prompt, aiconfig)

inputs = validate_and_retrieve_image_from_attachments(prompt)
model_settings = self.get_model_settings(prompt, aiconfig)
completion_params = refine_completion_params(model_settings)

#Add image inputs
inputs = validate_and_retrieve_images_from_attachments(prompt)
completion_params["inputs"] = inputs

await aiconfig.callback_manager.run_callbacks(CallbackEvent("on_deserialize_complete", __name__, {"output": completion_params}))
Expand All @@ -110,24 +124,93 @@ async def run_inference(self, prompt: Prompt, aiconfig: "AIConfigRuntime", optio
{"prompt": prompt, "options": options, "parameters": parameters},
)
)
model_name = aiconfig.get_model_name(prompt)

self.pipelines[model_name] = pipeline(task="image-to-text", model=model_name)

captioner = self.pipelines[model_name]
completion_data = await self.deserialize(prompt, aiconfig, parameters)
inputs = completion_data.pop("inputs")
model = completion_data.pop("model")
response = captioner(inputs, **completion_data)

output = ExecuteResult(output_type="execute_result", data=response, metadata={})
model_name: str | None = aiconfig.get_model_name(prompt)
if isinstance(model_name, str) and model_name not in self.pipelines:
self.pipelines[model_name] = pipeline(task="image-to-text", model=model_name)
captioner = self.pipelines[model_name]

outputs: List[Output] = []
response: List[Any] = captioner(inputs, **completion_data)
for count, result in enumerate(response):
output: Output = construct_regular_output(result, count)
outputs.append(output)

prompt.outputs = [output]
await aiconfig.callback_manager.run_callbacks(CallbackEvent("on_run_complete", __name__, {"result": prompt.outputs}))
prompt.outputs = outputs
print(f"{prompt.outputs=}")
await aiconfig.callback_manager.run_callbacks(
CallbackEvent(
"on_run_complete",
__name__,
{"result": prompt.outputs},
)
)
return prompt.outputs

def get_output_text(self, response: dict[str, Any]) -> str:
raise NotImplementedError("get_output_text is not implemented for HuggingFaceImage2TextTransformer")
def get_output_text(
self,
prompt: Prompt,
aiconfig: "AIConfigRuntime",
output: Optional[Output] = None,
) -> str:
if output is None:
output = aiconfig.get_latest_output(prompt)

if output is None:
return ""

# TODO (rossdanlm): Handle multiple outputs in list
# https://github.com/lastmile-ai/aiconfig/issues/467
if output.output_type == "execute_result":
output_data = output.data
if isinstance(output_data, str):
return output_data
if isinstance(output_data, OutputDataWithValue):
if isinstance(output_data.value, str):
return output_data.value
# HuggingFace Text summarization does not support function
# calls so shouldn't get here, but just being safe
return json.dumps(output_data.value, indent=2)
return ""


def refine_completion_params(model_settings: Dict[str, Any]) -> Dict[str, Any]:
"""
Refines the completion params for the HF image to text api. Removes any unsupported params.
The supported keys were found by looking at the HF ImageToTextPipeline.__call__ method
"""
supported_keys = {
"max_new_tokens",
"timeout",
}

completion_data = {}
for key in model_settings:
if key.lower() in supported_keys:
completion_data[key.lower()] = model_settings[key]

return completion_data

# Helper methods
def construct_regular_output(result: Dict[str, str], execution_count: int) -> Output:
"""
Construct regular output per response result, without streaming enabled
"""
output = ExecuteResult(
**{
"output_type": "execute_result",
# For some reason result is always in list format we haven't found
# a way of being able to return multiple sequences from the image
# to text pipeline
"data": result[0]["generated_text"],
"execution_count": execution_count,
"metadata": {},
}
)
return output


def validate_attachment_type_is_image(attachment: Attachment):
Expand All @@ -138,7 +221,7 @@ def validate_attachment_type_is_image(attachment: Attachment):
raise ValueError(f"Invalid attachment mimetype {attachment.mime_type}. Expected image mimetype.")


def validate_and_retrieve_image_from_attachments(prompt: Prompt) -> list[str]:
def validate_and_retrieve_images_from_attachments(prompt: Prompt) -> list[Union[str, Image]]:
"""
Retrieves the image uri's from each attachment in the prompt input.

Expand All @@ -152,15 +235,23 @@ def validate_and_retrieve_image_from_attachments(prompt: Prompt) -> list[str]:
if not hasattr(prompt.input, "attachments") or len(prompt.input.attachments) == 0:
raise ValueError(f"No attachments found in input for prompt {prompt.name}. Please add an image attachment to the prompt input.")

image_uris: list[str] = []
images: list[Union[str, Image]] = []

for i, attachment in enumerate(prompt.input.attachments):
validate_attachment_type_is_image(attachment)

if not isinstance(attachment.data, str):
input_data = attachment.data
if not isinstance(input_data, str):
# See todo above, but for now only support uri's
raise ValueError(f"Attachment #{i} data is not a uri. Please specify a uri for the image attachment in prompt {prompt.name}.")

image_uris.append(attachment.data)
# Really basic heurestic to check if the data is a base64 encoded str
# vs. uri. This will be fixed once we have standardized inputs
# See https://github.com/lastmile-ai/aiconfig/issues/829
if len(input_data) > 10000:
pil_image : Image = Image.open(BytesIO(base64.b64decode(input_data)))
images.append(pil_image)
else:
images.append(input_data)

return image_uris
return images
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@

# Step 1: define Helpers
def refine_pipeline_creation_params(model_settings: Dict[str, Any]) -> List[Dict[str, Any]]:
# There are from the transformers Github repo:
# https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L2534
supported_keys = {
"torch_dtype",
"force_download",
Expand Down Expand Up @@ -61,9 +63,11 @@ def refine_pipeline_creation_params(model_settings: Dict[str, Any]) -> List[Dict


def refine_completion_params(unfiltered_completion_params: Dict[str, Any]) -> Dict[str, Any]:
supported_keys = {
# ???
}
# Note: There seems to be no public API docs on what completion
# params are supported for text to speech:
# https://huggingface.co/docs/transformers/tasks/text-to-speech#inference
# The only one mentioned is `forward_params` which can contain `speaker_embeddings`
supported_keys = {}

completion_params: Dict[str, Any] = {}
for key in unfiltered_completion_params:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ def __init__(self):
config.register_model_parser(parser)
"""
super().__init__()
self.generators : dict[str, Pipeline]= {}
self.generators: dict[str, Pipeline]= {}

def id(self) -> str:
"""
Expand Down Expand Up @@ -217,14 +217,14 @@ async def deserialize(
# Build Completion data
model_settings = self.get_model_settings(prompt, aiconfig)
completion_data = refine_chat_completion_params(model_settings)

#Add resolved prompt
resolved_prompt = resolve_prompt(prompt, params, aiconfig)
completion_data["prompt"] = resolved_prompt
return completion_data

async def run_inference(
self, prompt: Prompt, aiconfig : "AIConfigRuntime", options : InferenceOptions, parameters: Dict[str, Any]
self, prompt: Prompt, aiconfig: "AIConfigRuntime", options: InferenceOptions, parameters: Dict[str, Any]
) -> List[Output]:
"""
Invoked to run a prompt in the .aiconfig. This method should perform
Expand All @@ -239,8 +239,8 @@ async def run_inference(
"""
completion_data = await self.deserialize(prompt, aiconfig, options, parameters)
completion_data["text_inputs"] = completion_data.pop("prompt", None)
model_name : str = aiconfig.get_model_name(prompt)

model_name: str | None = aiconfig.get_model_name(prompt)
if isinstance(model_name, str) and model_name not in self.generators:
self.generators[model_name] = pipeline('text-generation', model=model_name)
generator = self.generators[model_name]
Expand All @@ -251,14 +251,14 @@ async def run_inference(
not "stream" in completion_data or completion_data.get("stream") != False
)
if should_stream:
tokenizer : AutoTokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(model_name)
streamer = TextIteratorStreamer(tokenizer)
completion_data["streamer"] = streamer

outputs : List[Output] = []
outputs: List[Output] = []
output = None
if not should_stream:
response : List[Any] = generator(**completion_data)
response: List[Any] = generator(**completion_data)
for count, result in enumerate(response):
output = construct_regular_output(result, count)
outputs.append(output)
Expand All @@ -267,7 +267,7 @@ async def run_inference(
raise ValueError("Sorry, TextIteratorStreamer does not support multiple return sequences, please set `num_return_sequences` to 1")
if not streamer:
raise ValueError("Stream option is selected but streamer is not initialized")

# For streaming, cannot call `generator` directly otherwise response will be blocking
thread = threading.Thread(target=generator, kwargs=completion_data)
thread.start()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -128,13 +128,18 @@ def construct_stream_output(
"metadata": {},
}
)

accumulated_message = ""
for new_text in streamer:
if isinstance(new_text, str):
# For some reason these symbols aren't filtered out by the streamer
new_text = new_text.replace("</s>", "")
new_text = new_text.replace("<s>", "")

accumulated_message += new_text
options.stream_callback(new_text, accumulated_message, 0)

output.data = accumulated_message

return output


Expand Down Expand Up @@ -245,18 +250,18 @@ async def run_inference(self, prompt: Prompt, aiconfig: "AIConfigRuntime", optio

# if stream enabled in runtime options and config, then stream. Otherwise don't stream.
streamer = None
should_stream = (options.stream if options else False) and (not "stream" in completion_data or completion_data.get("stream") != False)
should_stream = (options.stream if options else False) and (
not "stream" in completion_data or completion_data.get("stream") != False
)
if should_stream:
tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(model_name)
streamer = TextIteratorStreamer(tokenizer)
completion_data["streamer"] = streamer

outputs: List[Output] = []
output = None

def _summarize():
return summarizer(inputs, **completion_data)

outputs: List[Output] = []
if not should_stream:
response: List[Any] = _summarize()
for count, result in enumerate(response):
Expand Down
Loading