codingutf8 2023 huggingface inc licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 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 abstract class to be reimplemented to define types that can be returned by agents these objects serve three purposes they behave as they were the type they re meant to be e g a string for text a pil image for images they can be stringified strobject in order to return a string defining the object they should be displayed correctly in ipython notebookscolabjupyter text type returned by the agent behaves as a string image type returned by the agent behaves as a pil image displays correctly this type in an ipython notebook ipython colab jupyter returns the raw version of that object in the case of an agentimage it is a pil image returns the stringified version of that object in the case of an agentimage it is a path to the serialized version of the image there is likely simpler than load into image into save audio type returned by the agent displays correctly this type in an ipython notebook ipython colab jupyter returns the raw version of that object it is a torch tensor object returns the stringified version of that object in the case of an agentaudio it is a path to the serialized version of the audio if the class has defined outputs we can map directly according to the class definition if the class does not have defined output then we map according to the type if the class has defined outputs we can map directly according to the class definition if the class does not have defined output then we map according to the type if the class has defined outputs we can map directly according to the class definition if the class does not have defined output then we map according to the type coding utf 8 2023 huggingface inc licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 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 abstract class to be reimplemented to define types that can be returned by agents these objects serve three purposes they behave as they were the type they re meant to be e g a string for text a pil image for images they can be stringified str object in order to return a string defining the object they should be displayed correctly in ipython notebooks colab jupyter text type returned by the agent behaves as a string image type returned by the agent behaves as a pil image displays correctly this type in an ipython notebook ipython colab jupyter returns the raw version of that object in the case of an agentimage it is a pil image returns the stringified version of that object in the case of an agentimage it is a path to the serialized version of the image there is likely simpler than load into image into save audio type returned by the agent displays correctly this type in an ipython notebook ipython colab jupyter returns the raw version of that object it is a torch tensor object returns the stringified version of that object in the case of an agentaudio it is a path to the serialized version of the audio if the class has defined outputs we can map directly according to the class definition if the class does not have defined output then we map according to the type if the class has defined outputs we can map directly according to the class definition if the class does not have defined output then we map according to the type if the class has defined outputs we can map directly according to the class definition if the class does not have defined output then we map according to the type
import os
import pathlib
import tempfile
import uuid

import numpy as np

from ..utils import is_soundfile_availble, is_torch_available, is_vision_available, logging


logger = logging.get_logger(__name__)

if is_vision_available():
    import PIL.Image
    from PIL import Image
    from PIL.Image import Image as ImageType
else:
    ImageType = object

if is_torch_available():
    import torch

if is_soundfile_availble():
    import soundfile as sf


class AgentType:
    

    def __init__(self, value):
        self._value = value

    def __str__(self):
        return self.to_string()

    def to_raw(self):
        logger.error(
            "This is a raw AgentType of unknown type. Display in notebooks and string conversion will be unreliable"
        )
        return self._value

    def to_string(self) -> str:
        logger.error(
            "This is a raw AgentType of unknown type. Display in notebooks and string conversion will be unreliable"
        )
        return str(self._value)


class AgentText(AgentType, str):
    

    def to_raw(self):
        return self._value

    def to_string(self):
        return self._value


class AgentImage(AgentType, ImageType):
    

    def __init__(self, value):
        super().__init__(value)

        if not is_vision_available():
            raise ImportError("PIL must be installed in order to handle images.")

        self._path = None
        self._raw = None
        self._tensor = None

        if isinstance(value, ImageType):
            self._raw = value
        elif isinstance(value, (str, pathlib.Path)):
            self._path = value
        elif isinstance(value, torch.Tensor):
            self._tensor = value
        else:
            raise ValueError(f"Unsupported type for {self.__class__.__name__}: {type(value)}")

    def _ipython_display_(self, include=None, exclude=None):
        
        from IPython.display import Image, display

        display(Image(self.to_string()))

    def to_raw(self):
        
        if self._raw is not None:
            return self._raw

        if self._path is not None:
            self._raw = Image.open(self._path)
            return self._raw

    def to_string(self):
        
        if self._path is not None:
            return self._path

        if self._raw is not None:
            directory = tempfile.mkdtemp()
            self._path = os.path.join(directory, str(uuid.uuid4()) + ".png")
            self._raw.save(self._path)

            return self._path

        if self._tensor is not None:
            array = self._tensor.cpu().detach().numpy()

            
            img = Image.fromarray((array * 255).astype(np.uint8))

            directory = tempfile.mkdtemp()
            self._path = os.path.join(directory, str(uuid.uuid4()) + ".png")

            img.save(self._path)

            return self._path


class AgentAudio(AgentType):
    

    def __init__(self, value, samplerate=16_000):
        super().__init__(value)

        if not is_soundfile_availble():
            raise ImportError("soundfile must be installed in order to handle audio.")

        self._path = None
        self._tensor = None

        self.samplerate = samplerate

        if isinstance(value, (str, pathlib.Path)):
            self._path = value
        elif isinstance(value, torch.Tensor):
            self._tensor = value
        else:
            raise ValueError(f"Unsupported audio type: {type(value)}")

    def _ipython_display_(self, include=None, exclude=None):
        
        from IPython.display import Audio, display

        display(Audio(self.to_string(), rate=self.samplerate))

    def to_raw(self):
        
        if self._tensor is not None:
            return self._tensor

        if self._path is not None:
            tensor, self.samplerate = sf.read(self._path)
            self._tensor = torch.tensor(tensor)
            return self._tensor

    def to_string(self):
        
        if self._path is not None:
            return self._path

        if self._tensor is not None:
            directory = tempfile.mkdtemp()
            self._path = os.path.join(directory, str(uuid.uuid4()) + ".wav")
            sf.write(self._path, self._tensor, samplerate=self.samplerate)
            return self._path


AGENT_TYPE_MAPPING = {"text": AgentText, "image": AgentImage, "audio": AgentAudio}
INSTANCE_TYPE_MAPPING = {str: AgentText}

if is_vision_available():
    INSTANCE_TYPE_MAPPING[PIL.Image] = AgentImage


def handle_agent_inputs(*args, **kwargs):
    args = [(arg.to_raw() if isinstance(arg, AgentType) else arg) for arg in args]
    kwargs = {k: (v.to_raw() if isinstance(v, AgentType) else v) for k, v in kwargs.items()}
    return args, kwargs


def handle_agent_outputs(outputs, output_types=None):
    if isinstance(outputs, dict):
        decoded_outputs = {}
        for i, (k, v) in enumerate(outputs.items()):
            if output_types is not None:
                
                if output_types[i] in AGENT_TYPE_MAPPING:
                    decoded_outputs[k] = AGENT_TYPE_MAPPING[output_types[i]](v)
                else:
                    decoded_outputs[k] = AgentType(v)

            else:
                
                for _k, _v in INSTANCE_TYPE_MAPPING.items():
                    if isinstance(v, _k):
                        decoded_outputs[k] = _v(v)
                if k not in decoded_outputs:
                    decoded_outputs[k] = AgentType[v]

    elif isinstance(outputs, (list, tuple)):
        decoded_outputs = type(outputs)()
        for i, v in enumerate(outputs):
            if output_types is not None:
                
                if output_types[i] in AGENT_TYPE_MAPPING:
                    decoded_outputs.append(AGENT_TYPE_MAPPING[output_types[i]](v))
                else:
                    decoded_outputs.append(AgentType(v))
            else:
                
                found = False
                for _k, _v in INSTANCE_TYPE_MAPPING.items():
                    if isinstance(v, _k):
                        decoded_outputs.append(_v(v))
                        found = True

                if not found:
                    decoded_outputs.append(AgentType(v))

    else:
        if output_types[0] in AGENT_TYPE_MAPPING:
            
            decoded_outputs = AGENT_TYPE_MAPPING[output_types[0]](outputs)

        else:
            
            for _k, _v in INSTANCE_TYPE_MAPPING.items():
                if isinstance(outputs, _k):
                    return _v(outputs)
            return AgentType(outputs)

    return decoded_outputs