kaiko-eva 0.3.3__py3-none-any.whl → 0.4.1__py3-none-any.whl

eva/language/models/wrappers/huggingface.py

@@ -1,22 +1,34 @@
 """LLM wrapper for HuggingFace `transformers` models."""
 
-from typing import Any, Dict, List, Literal
+from typing import Any, Callable, Dict, List, Literal
 
 from transformers.pipelines import pipeline
 from typing_extensions import override
 
-from eva.core.models.wrappers import base
+from eva.language.models.typings import ModelOutput, TextBatch
+from eva.language.models.wrappers import base
+from eva.language.utils.text import messages as message_utils
 
 
-class HuggingFaceTextModel(base.BaseModel[List[str], List[str]]):
+class HuggingFaceModel(base.LanguageModel):
     """Wrapper class for loading HuggingFace `transformers` models using pipelines."""
 
+    _default_generation_kwargs = {
+        "temperature": 0.0,
+        "max_new_tokens": 1024,
+        "do_sample": False,
+        "top_p": 1.0,
+    }
+    """Default HF model parameters for evaluation."""
+
     def __init__(
         self,
         model_name_or_path: str,
         task: Literal["text-generation"] = "text-generation",
         model_kwargs: Dict[str, Any] | None = None,
+        system_prompt: str | None = None,
         generation_kwargs: Dict[str, Any] | None = None,
+        chat_mode: bool = True,
     ) -> None:
         """Initializes the model.
 
@@ -26,21 +38,26 @@ class HuggingFaceTextModel(base.BaseModel[List[str], List[str]]):
                 model hub.
             task: The pipeline task. Defaults to "text-generation".
             model_kwargs: Additional arguments for configuring the pipeline.
+            system_prompt: System prompt to use.
             generation_kwargs: Additional generation parameters (temperature, max_length, etc.).
+            chat_mode: Whether the specified model expects chat style messages. If set to False
+                the model is assumed to be a standard text completion model and will expect
+                plain text string inputs.
         """
-        super().__init__()
+        super().__init__(system_prompt=system_prompt)
 
         self._model_name_or_path = model_name_or_path
         self._task = task
         self._model_kwargs = model_kwargs or {}
-        self._generation_kwargs = generation_kwargs or {}
+        self._generation_kwargs = self._default_generation_kwargs | (generation_kwargs or {})
+        self._chat_mode = chat_mode
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
-    def load_model(self) -> None:
+    def load_model(self) -> Callable:
         """Loads the model as a Hugging Face pipeline."""
-        self._pipeline = pipeline(
+        return pipeline(
             task=self._task,
             model=self._model_name_or_path,
             trust_remote_code=True,
@@ -48,7 +65,34 @@ class HuggingFaceTextModel(base.BaseModel[List[str], List[str]]):
         )
 
     @override
-    def model_forward(self, prompts: List[str]) -> List[str]:
+    def format_inputs(self, batch: TextBatch) -> List[List[Dict[str, Any]]] | List[str]:
+        """Formats inputs for HuggingFace models.
+
+        Note: If multiple system messages are present, they will be combined
+        into a single message, given that many models only support a single
+        system prompt.
+
+        Args:
+            batch: A batch of text and image inputs.
+
+        Returns:
+            When in chat mode, returns a batch of message series following
+            OpenAI's API format {"role": "user", "content": "..."}, for non-chat
+            models returns a list of plain text strings.
+        """
+        message_batch, _, _ = TextBatch(*batch)
+        message_batch = message_utils.batch_insert_system_message(
+            message_batch, self.system_message
+        )
+        message_batch = list(map(message_utils.combine_system_messages, message_batch))
+
+        if self._chat_mode:
+            return list(map(message_utils.format_chat_message, message_batch))
+        else:
+            return list(map(message_utils.merge_message_contents, message_batch))
+
+    @override
+    def model_forward(self, prompts: List[str]) -> ModelOutput:
         """Generates text using the pipeline.
 
         Args:
@@ -57,13 +101,15 @@ class HuggingFaceTextModel(base.BaseModel[List[str], List[str]]):
         Returns:
             The generated text as a string.
         """
-        outputs = self._pipeline(prompts, return_full_text=False, **self._generation_kwargs)
+        outputs = self.model(prompts, return_full_text=False, **self._generation_kwargs)
         if outputs is None:
             raise ValueError("Outputs from the model are None.")
+
         results = []
         for output in outputs:
             if isinstance(output, list):
                 results.append(output[0]["generated_text"])  # type: ignore
             else:
                 results.append(output["generated_text"])  # type: ignore
-        return results
+
+        return ModelOutput(generated_text=results)

eva/language/models/wrappers/litellm.py

@@ -1,77 +1,122 @@
-"""LLM wrapper for litellm models."""
+"""LiteLLM language model wrapper."""
 
+import logging
 from typing import Any, Dict, List
 
-from litellm import batch_completion  # type: ignore
+import backoff
+import litellm
+from litellm import batch_completion
+from litellm.exceptions import (
+    APIConnectionError,
+    InternalServerError,
+    RateLimitError,
+    ServiceUnavailableError,
+    Timeout,
+)
 from loguru import logger
 from typing_extensions import override
 
-from eva.core.models.wrappers import base
+from eva.language.models.typings import ModelOutput, TextBatch
+from eva.language.models.wrappers import base
+from eva.language.utils.text import messages as message_utils
 
+RETRYABLE_ERRORS = (
+    RateLimitError,
+    Timeout,
+    InternalServerError,
+    APIConnectionError,
+    ServiceUnavailableError,
+)
 
-class LiteLLMTextModel(base.BaseModel[List[str], List[str]]):
-    """Wrapper class for using litellm for chat-based text generation.
 
-    This wrapper uses litellm's `completion` function which accepts a list of
-    message dicts. The `forward` method converts a string prompt into a chat
-    message with a default "user" role, optionally prepends a system message,
-    and includes an API key if provided.
-    """
+class LiteLLMModel(base.LanguageModel):
+    """Wrapper class for LiteLLM language models."""
+
+    _default_model_kwargs = {
+        "temperature": 0.0,
+        "max_completion_tokens": 1024,
+        "top_p": 1.0,
+        "seed": 42,
+    }
+    """Default API model parameters for evaluation."""
 
     def __init__(
         self,
-        model_name_or_path: str,
+        model_name: str,
         model_kwargs: Dict[str, Any] | None = None,
-    ) -> None:
-        """Initializes the litellm chat model wrapper.
+        system_prompt: str | None = None,
+        log_level: int | None = logging.INFO,
+    ):
+        """Initialize the LiteLLM Wrapper.
 
         Args:
-            model_name_or_path: The model identifier (or name) for litellm
-                (e.g.,"openai/gpt-4o" or "anthropic/claude-3-sonnet-20240229").
+            model_name: The name of the model to use.
             model_kwargs: Additional keyword arguments to pass during
                 generation (e.g., `temperature`, `max_tokens`).
+            system_prompt: The system prompt to use (optional).
+            log_level: Optional logging level for LiteLLM. Defaults to WARNING.
         """
-        super().__init__()
-        self._model_name_or_path = model_name_or_path
-        self._model_kwargs = model_kwargs or {}
-        self.load_model()
+        super().__init__(system_prompt=system_prompt)
 
-    @override
-    def load_model(self) -> None:
-        """Prepares the litellm model.
+        self.model_name = model_name
+        self.model_kwargs = self._default_model_kwargs | (model_kwargs or {})
 
-        Note:
-            litellm doesn't require an explicit loading step; models are called
-            directly during generation. This method exists for API consistency.
-        """
-        pass
+        litellm.suppress_debug_info = True
+        litellm.drop_params = True
+
+        if log_level is not None:
+            logging.getLogger("LiteLLM").setLevel(log_level)
 
     @override
-    def model_forward(self, prompts: List[str]) -> List[str]:
-        """Generates text using litellm.
+    def format_inputs(self, batch: TextBatch) -> List[List[Dict[str, Any]]]:
+        """Formats inputs for LiteLLM.
 
         Args:
-            prompts: A list of prompts to be converted into a "user" message.
+            batch: A batch of text inputs.
 
         Returns:
-            A list of generated text responses. Failed generations will contain
-            error messages instead of generated text.
+            A list of messages in the following format:
+            [
+                {
+                    "role": ...
+                    "content": ...
+                },
+                ...
+            ]
         """
-        messages = [[{"role": "user", "content": prompt}] for prompt in prompts]
+        message_batch, _, _ = TextBatch(*batch)
 
-        responses = batch_completion(
-            model=self._model_name_or_path,
-            messages=messages,
-            **self._model_kwargs,
+        message_batch = message_utils.batch_insert_system_message(
+            message_batch, self.system_message
         )
+        message_batch = list(map(message_utils.combine_system_messages, message_batch))
 
-        results = []
-        for i, response in enumerate(responses):
-            if isinstance(response, Exception):
-                error_msg = f"Error generating text for prompt {i}: {response}"
-                logger.error(error_msg)
-                raise RuntimeError(error_msg)
-            else:
-                results.append(response["choices"][0]["message"]["content"])
+        return list(map(message_utils.format_chat_message, message_batch))
 
-        return results
+    @override
+    @backoff.on_exception(
+        backoff.expo,
+        RETRYABLE_ERRORS,
+        max_tries=20,
+        jitter=backoff.full_jitter,
+        on_backoff=lambda details: logger.warning(
+            f"Retrying due to {details.get('exception') or 'Unknown error'}"
+        ),
+    )
+    def model_forward(self, batch: List[List[Dict[str, Any]]]) -> ModelOutput:
+        """Generates output text through API calls via LiteLLM's batch completion functionality."""
+        outputs = batch_completion(model=self.model_name, messages=batch, **self.model_kwargs)
+        self._raise_exceptions(outputs)
+
+        generated_text = [
+            output["choices"][0]["message"]["content"]
+            for output in outputs
+            if output["choices"][0]["message"]["role"] == "assistant"
+        ]
+        return ModelOutput(generated_text=generated_text)
+
+    def _raise_exceptions(self, outputs: list):
+        for output in outputs:
+            if isinstance(output, Exception):
+                logger.error(f"Model {self.model_name} encountered an error: {output}")
+                raise output

eva/language/models/wrappers/vllm.py

@@ -1,6 +1,6 @@
 """LLM wrapper for vLLM models."""
 
-from typing import Any, Dict, List, Sequence
+from typing import Any, Dict, List
 
 from loguru import logger
 from typing_extensions import override
@@ -11,17 +11,20 @@ try:
     from vllm.transformers_utils.tokenizer import AnyTokenizer  # type: ignore
 except ImportError as e:
     raise ImportError(
-        "vLLM is required for VLLMTextModel but not installed. "
+        "vLLM is required for VllmModel but not installed. "
         "vLLM must be installed manually as it requires CUDA and is not included in dependencies. "
         "Install with: pip install vllm "
         "Note: vLLM requires Linux with CUDA support for optimal performance. "
-        "For alternatives, consider using HuggingFaceTextModel or LiteLLMTextModel."
+        "For alternatives, consider using HuggingFaceModel or LiteLLMModel."
     ) from e
 
-from eva.core.models.wrappers import base
+from eva.language.data.messages import MessageSeries
+from eva.language.models.typings import TextBatch
+from eva.language.models.wrappers import base
+from eva.language.utils.text import messages as message_utils
 
 
-class VLLMTextModel(base.BaseModel):
+class VllmModel(base.LanguageModel):
     """Wrapper class for using vLLM for text generation.
 
     This wrapper loads a vLLM model, sets up the tokenizer and sampling
@@ -34,6 +37,7 @@ class VLLMTextModel(base.BaseModel):
         self,
         model_name_or_path: str,
         model_kwargs: Dict[str, Any] | None = None,
+        system_prompt: str | None = None,
         generation_kwargs: Dict[str, Any] | None = None,
     ) -> None:
         """Initializes the vLLM model wrapper.
@@ -44,12 +48,13 @@ class VLLMTextModel(base.BaseModel):
             model_kwargs: Arguments required to initialize the vLLM model,
                 see [link](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/llm.py)
                 for more information.
+            system_prompt: System prompt to use.
             generation_kwargs: Arguments required to generate the output,
                 need to align with the arguments of
                 [vllm.SamplingParams](https://github.com/vllm-project/vllm/blob/main/vllm/sampling_params.py).
 
         """
-        super().__init__()
+        super().__init__(system_prompt=system_prompt)
         self._model_name_or_path = model_name_or_path
         self._model_kwargs = model_kwargs or {}
         self._generation_kwargs = generation_kwargs or {}
@@ -71,11 +76,11 @@ class VLLMTextModel(base.BaseModel):
             raise RuntimeError("Model not initialized")
         self._llm_tokenizer = self._llm_model.get_tokenizer()
 
-    def _apply_chat_template(self, prompts: Sequence[str]) -> list[TokensPrompt]:
+    def _tokenize_messages(self, messages: List[MessageSeries]) -> List[TokensPrompt]:
         """Apply chat template to the messages.
 
         Args:
-            prompts: List of raw user strings.
+            messages: List of raw user strings.
 
         Returns:
             List of encoded messages.
@@ -90,7 +95,8 @@ class VLLMTextModel(base.BaseModel):
         if not hasattr(self._llm_tokenizer, "chat_template"):
             raise ValueError("Tokenizer does not have a chat template.")
 
-        chat_messages = [[{"role": "user", "content": p}] for p in prompts]
+        chat_messages = list(map(message_utils.format_chat_message, messages))
+
         encoded_messages = self._llm_tokenizer.apply_chat_template(
             chat_messages,  # type: ignore
             tokenize=True,
@@ -131,11 +137,30 @@ class VLLMTextModel(base.BaseModel):
 
         return result
 
-    def generate(self, prompts: List[str]) -> List[str]:
+    @override
+    def format_inputs(self, batch: TextBatch) -> List[TokensPrompt]:
+        """Formats inputs for vLLM models.
+
+        Args:
+            batch: A batch of text and image inputs.
+
+        Returns:
+            List of formatted prompts.
+        """
+        message_batch, _, _ = TextBatch(*batch)
+        message_batch = message_utils.batch_insert_system_message(
+            message_batch, self.system_message
+        )
+        message_batch = list(map(message_utils.combine_system_messages, message_batch))
+
+        return self._tokenize_messages(message_batch)
+
+    @override
+    def model_forward(self, batch: List[TokensPrompt]) -> List[str]:
         """Generates text for the given prompt using the vLLM model.
 
         Args:
-            prompts: A list of string prompts for generation.
+            batch: A list encoded / tokenized messages (TokensPrompt objects).
 
         Returns:
             The generated text response.
@@ -144,6 +169,5 @@ class VLLMTextModel(base.BaseModel):
         if self._llm_model is None:
             raise RuntimeError("Model not initialized")
 
-        prompt_tokens = self._apply_chat_template(prompts)
-        outputs = self._llm_model.generate(prompt_tokens, SamplingParams(**self._generation_kwargs))
+        outputs = self._llm_model.generate(batch, SamplingParams(**self._generation_kwargs))
         return [output.outputs[0].text for output in outputs]

eva/language/utils/__init__.py

@@ -1,5 +1,6 @@
 """Language utilities and helper functions."""
 
 from eva.language.utils.str_to_int_tensor import CastStrToIntTensor
+from eva.language.utils.text.messages import format_chat_message
 
-__all__ = ["CastStrToIntTensor"]
+__all__ = ["CastStrToIntTensor", "format_chat_message"]

eva/language/utils/str_to_int_tensor.py

@@ -16,11 +16,11 @@ class CastStrToIntTensor:
     Supports single values, lists of strings, or lists of integers.
 
     Example:
-        >>> # Default mapping for yes/no/maybe classification
-        >>> transform = CastStrToIntTensor()
-        >>> transform(['yes', 'no', 'maybe'])
+        >>> # Default mapping for A/B/C classification
+        >>> transform = CastStrToIntTensor(mapping={"A": 0, "B": 1, "C": 2})
+        >>> transform(['B', 'A', 'C'])
         tensor([1, 0, 2])
-        >>> transform('yes')
+        >>> transform('B')
         tensor([1])
 
         >>> # Custom mapping
@@ -29,20 +29,25 @@ class CastStrToIntTensor:
         tensor([1, 0])
     """
 
-    def __init__(self, mapping: Dict[str, int] | None = None):
-        """Initialize the transform with a regex-to-integer mapping.
+    def __init__(
+        self, mapping: Dict[str, int], standalone_words: bool = True, case_sensitive: bool = True
+    ) -> None:
+        r"""Initialize the transform with a regex-to-integer mapping.
 
         Args:
             mapping: Dictionary mapping regex patterns to integers. If None, uses default
                     yes/no/maybe mapping: {'no': 0, 'yes': 1, 'maybe': 2}
+            standalone_words: If True, patterns are treated as standalone words (e.g., '\bno\b').
+            case_sensitive: If True, regex patterns are case-sensitive.
         """
-        if mapping is None:
-            self.mapping = {r"\bno\b": 0, r"\byes\b": 1, r"\bmaybe\b": 2}
-        else:
-            self.mapping = mapping
+        self.mapping = mapping
+
+        if standalone_words:
+            self.mapping = {rf"\b{k}\b": v for k, v in mapping.items()}
 
         self.compiled_patterns = [
-            (re.compile(pattern, re.IGNORECASE), value) for pattern, value in self.mapping.items()
+            (re.compile(pattern, 0 if case_sensitive else re.IGNORECASE), value)
+            for pattern, value in self.mapping.items()
         ]
 
     def __call__(self, values: Union[str, List[str], List[int]]) -> torch.Tensor:
@@ -58,7 +63,10 @@ class CastStrToIntTensor:
             ValueError: If any value cannot be mapped to an integer.
         """
         return torch.tensor(
-            [self._cast_single(v) for v in (values if isinstance(values, list) else [values])],
+            [
+                self._cast_single(v)
+                for v in (values if isinstance(values, list | tuple) else [values])
+            ],
             dtype=torch.int,
         )
 

eva/language/utils/text/__init__.py

@@ -0,0 +1,5 @@
+"""Text utilities for language models."""
+
+from eva.language.utils.text.messages import format_chat_message
+
+__all__ = ["format_chat_message"]

eva/language/utils/text/messages.py

@@ -0,0 +1,113 @@
+"""Message formatting utilities for language models."""
+
+import functools
+import json
+from typing import Any, Dict, List
+
+from eva.language.data.messages import (
+    AssistantMessage,
+    MessageSeries,
+    Role,
+    SystemMessage,
+    UserMessage,
+)
+
+
+def format_chat_message(message: MessageSeries) -> List[Dict[str, Any]]:
+    """Formats a message series into a format following OpenAI's API specification."""
+    return [{"role": item.role, "content": item.content} for item in message]
+
+
+def combine_system_messages(message: MessageSeries, join_char: str = "\n") -> MessageSeries:
+    """Combine system messages into a single message.
+
+    This is useful when the MessageSeries contains multiple system messages such
+    as `ModelSystemMessage` and `TaskSystemMessage`. But given that most models / apis
+    expect a single system message, this function can be used to combines them into one.
+
+    Args:
+        message: The message series containing one or multiple messages.
+        join_char: The character to use to join the system messages. Default is newline.
+
+    Returns:
+        A new message series with system messages combined into one and the
+        remaining messages unchanged.
+    """
+    system_messages = list(filter(lambda item: item.role == Role.SYSTEM, message))
+    if len(system_messages) == 0:
+        return message
+
+    non_system_messages = list(filter(lambda item: item.role != Role.SYSTEM, message))
+    return [
+        SystemMessage(content=merge_message_contents(system_messages, join_char=join_char))
+    ] + non_system_messages
+
+
+def merge_message_contents(message: MessageSeries, join_char: str = "\n") -> str:
+    """Merges the all contents within a message series into a string.
+
+    Args:
+        message: The message series to combine.
+        join_char: The character to use to join the message contents. Default is newline.
+
+    Returns:
+        A string containing the combined message contents.
+    """
+    return join_char.join(item.content for item in message)
+
+
+def insert_system_message(
+    message: MessageSeries, system_message: SystemMessage | None
+) -> MessageSeries:
+    """Insert a system message at the beginning of the message series."""
+    if system_message is None:
+        return message
+    return [system_message] + message
+
+
+def batch_insert_system_message(
+    messages: List[MessageSeries], system_message: SystemMessage | None
+) -> List[MessageSeries]:
+    """Insert a system message at the beginning of each message series in a batch."""
+    return list(
+        map(functools.partial(insert_system_message, system_message=system_message), messages)
+    )
+
+
+def serialize(messages: MessageSeries) -> str:
+    """Serialize a MessageSeries object into a JSON string.
+
+    Args:
+        messages: A list of message objects (MessagesSeries).
+
+    Returns:
+        A JSON string representing the message series, with the following format:
+            [{"role": "user", "content": "Hello"}, ...]
+    """
+    serialized_messages = format_chat_message(messages)
+    return json.dumps(serialized_messages)
+
+
+def deserialize(messages: str) -> MessageSeries:
+    """Convert a json string to a MessageSeries object.
+
+    Format: [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}]
+    """
+    message_dicts = json.loads(messages)
+
+    message_series = []
+    for message_dict in message_dicts:
+        if "role" not in message_dict or "content" not in message_dict:
+            raise ValueError("`role` or `content` keys are missing.")
+
+        match message_dict["role"]:
+            case Role.USER:
+                message_series.append(UserMessage(**message_dict))
+            case Role.ASSISTANT:
+                message_series.append(AssistantMessage(**message_dict))
+            case Role.SYSTEM:
+                message_series.append(SystemMessage(**message_dict))
+            case _:
+                raise ValueError(f"Unknown role: {message_dict['role']}")
+
+    return message_series

eva/multimodal/__init__.py

@@ -0,0 +1,6 @@
+"""Multimodal API."""
+
+from eva.multimodal import models
+from eva.multimodal.data import datasets
+
+__all__ = ["models", "datasets"]

eva/multimodal/callbacks/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal callbacks API."""
+
+from eva.multimodal.callbacks.writers import TextPredictionWriter
+
+__all__ = ["TextPredictionWriter"]

eva/multimodal/callbacks/writers/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal writers callbacks API."""
+
+from eva.multimodal.callbacks.writers.prediction import TextPredictionWriter
+
+__all__ = ["TextPredictionWriter"]

eva/multimodal/callbacks/writers/prediction.py

@@ -0,0 +1,39 @@
+"""Text prediction writer callbacks."""
+
+from typing import Dict, List, Literal, Tuple
+
+from torch import nn
+from typing_extensions import override
+
+from eva.language.callbacks import writers
+from eva.multimodal.models.typings import TextImageBatch
+
+
+class TextPredictionWriter(writers.TextPredictionWriter):
+    """Callback for writing generated text predictions to disk."""
+
+    def __init__(
+        self,
+        output_dir: str,
+        model: nn.Module,
+        dataloader_idx_map: Dict[int, str] | None = None,
+        metadata_keys: List[str] | None = None,
+        include_input: bool = True,
+        overwrite: bool = False,
+        save_format: Literal["jsonl", "parquet", "csv"] = "jsonl",
+    ) -> None:
+        """See docstring of base class."""
+        super().__init__(
+            output_dir=output_dir,
+            model=model,
+            dataloader_idx_map=dataloader_idx_map,
+            metadata_keys=metadata_keys,
+            include_input=include_input,
+            overwrite=overwrite,
+            save_format=save_format,
+        )
+
+    @override
+    def _unpack_batch(self, batch: TextImageBatch) -> Tuple[list, list | None, dict | None]:  # type: ignore
+        text_batch, _, target_batch, metadata_batch = TextImageBatch(*batch)
+        return text_batch, target_batch, metadata_batch

eva/multimodal/data/__init__.py

@@ -0,0 +1,5 @@
+"""Data components for multimodal learning."""
+
+from eva.multimodal.data import datasets
+
+__all__ = ["datasets"]