[HF][streaming][4/n] Image2Text (no streaming, but lots of fixing)

extensions/HuggingFace/python/src/aiconfig_extension_hugging_face/local_inference/image_2_text.py

@@ -1,11 +1,21 @@
+import json
 from typing import Any, Dict, Optional, List, TYPE_CHECKING
+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
 
@@ -93,10 +103,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)
+        model_settings = self.get_model_settings(prompt, aiconfig)
+        completion_params = refine_completion_params(model_settings)
 
+        #Add image inputs
         inputs = validate_and_retrieve_image_from_attachments(prompt)
-
         completion_params["inputs"] = inputs
 
         await aiconfig.callback_manager.run_callbacks(CallbackEvent("on_deserialize_complete", __name__, {"output": completion_params}))
@@ -110,24 +121,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):

extensions/HuggingFace/python/src/aiconfig_extension_hugging_face/local_inference/text_generation.py

@@ -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:
         """
@@ -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
@@ -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]
@@ -255,10 +255,10 @@ async def run_inference(
             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)
@@ -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()

extensions/HuggingFace/python/src/aiconfig_extension_hugging_face/local_inference/text_summarization.py

@@ -258,12 +258,10 @@ async def run_inference(self, prompt: Prompt, aiconfig: "AIConfigRuntime", optio
             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):