flexeval 0.12.2__tar.gz → 0.13.1__tar.gz

{flexeval-0.12.2 → flexeval-0.13.1}/PKG-INFO

@@ -1,12 +1,11 @@
 Metadata-Version: 2.3
 Name: flexeval
-Version: 0.12.2
+Version: 0.13.1
 Summary: 
 Author: ryokan-ri
 Author-email: ryokan.ri@sbintuitions.co.jp
-Requires-Python: >=3.9, !=2.7.*, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, !=3.6.*, !=3.7.*, !=3.8.*, !=3.13.*
+Requires-Python: >=3.10,<3.13
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -21,6 +20,7 @@ Requires-Dist: jiwer (>=3.0.4,<4.0.0)
 Requires-Dist: jsonargparse[jsonnet] (>=4.26.1,<5.0.0)
 Requires-Dist: litellm (>=1.52.9,<2.0.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
+Requires-Dist: math-verify[antlr4-13-2] (>=0.7.0,<0.8.0)
 Requires-Dist: openai (>=1.52.2,<2.0.0)
 Requires-Dist: peft (>=0.10.0,<0.11.0)
 Requires-Dist: pyarrow (==16.1.0)
@@ -33,7 +33,7 @@ Requires-Dist: smart-open (>=7.1.0,<8.0.0)
 Requires-Dist: sudachipy (>=0.6.10)
 Requires-Dist: tiktoken (>=0.9.0,<0.10.0)
 Requires-Dist: transformers[ja,sentencepiece,torch] (>=4.34.1,<5.0.0)
-Requires-Dist: vllm (>=0.8.4,<0.9.0) ; extra == "vllm"
+Requires-Dist: vllm (==0.9.2) ; extra == "vllm"
 Requires-Dist: wandb (>=0.17.2,<0.18.0) ; extra == "wandb"
 Description-Content-Type: text/markdown
 

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/chat_dataset/base.py

@@ -28,7 +28,34 @@ class ChatInstance:
         }
     ]
     ```
-    """
+
+    Tool-Calling message must follow the same format as the OpenAI ChatCompletion API.
+    https://platform.openai.com/docs/guides/function-calling?api-mode=chat#defining-functions
+    ```json
+    {
+        "role": "assistant",
+        "content": "content", # `None` is also allowed if `tool_calls` exists.
+        "tool_calls": [
+            {
+                "id": "dummy1",
+                "function": {
+                    "name": "search_web",
+                    "arguments": "{\"query\": \"flexeval developer\"}"  # Note that this is a json string, not a dictionary.
+                }
+            }
+        ]
+    }
+    ```
+
+    The results from tools should be represented as messages with the role "tool":
+    ```
+    {
+        "role": "tool",
+        "tool_call_id": "dummy1", # Optional, models on OpenAI APIs requires this field.
+        "name": "search_web", # Optional, Some HuggingFace models require this field.
+        "content": "[{\"title\": \"sbintuitions/flexeval: Flexible evaluation tool...\", \"description\": \"...\"}]",
+    }
+    """  # noqa: E501
     tools: list[dict[str, Any]] | None = None
     """
     A list of definitions of tools in the chat.

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/chat_dataset/openai_messages.py

@@ -16,6 +16,10 @@ class OpenAIMessagesDataset(ChatDataset):
     The difference lies in that this class has 'tool_definition' field, in which
     available tools are listed.
 
+    Tool-Calling (Function-Calling) is supported in this class.
+    It must follow the same format as the OpenAI ChatCompletion API.
+    https://platform.openai.com/docs/guides/function-calling?api-mode=chat#defining-functions
+
     Parameters:
         file_path (str | list[str] | None): Path or list of paths to `.jsonl` file(s).
         message_key (str): Key used to extract the list of messages from each JSON object.
@@ -24,25 +28,54 @@ class OpenAIMessagesDataset(ChatDataset):
         drop_if_last_from_assistant (bool): If true, when the last utterance is given by assistant, drop it.
 
     In Jsonl, each line must have a following structure:
-
+    ```json
     {
       '<message_key>': [
         {
           'role': 'user',
-          'content': 'こんにちわ。元気になる言葉を教えて下さい。'
+          'content': 'こんにちは。元気が出る言葉を教えて下さい。'
         },
         {
           'role': 'assistant',
           'content': 'こんなのはどうでしょう。どんどんやってください！'
+        },
+      ],
+    }
+    ```
+
+    Example with tool-calling:
+    ```json
+    {
+      '<message_key>': [
+        {
+          'role': 'user',
+          'content': 'こんにちは。元気が出る偉人の言葉を教えて下さい。'
+        },
+        {
+          'role': 'assistant',
+          'content': '調べてみますね。',
+          'tool_calls': [
+            {
+              'id': 'dummy1',
+              'function': {
+                'name': 'web_search',
+                'arguments': '{"query": "元気が出る言葉 偉人"}',
+              }
+            }
+          ]
         }
       ],
       '<tool_definitions_key>': [
         {
-          'type': 'function',
-          'function': { ... }
+          "type": "function",
+          "function": {
+            "name": "web_search",
+            ...
+          }
         }
       ]
     }
+    ```
     """
 
     def __init__(

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/evaluate_chat_response.py

@@ -185,6 +185,8 @@ def evaluate_chat_response(  # noqa: C901,PLR0912, PLR0915
             if mes["content"] is None:
                 mes["content"] = ""
 
+    language_model.cleanup_resources()
+
     # Evaluate the generated responses
     metrics_summary_dict: dict[str, float] = {}
     instance_metrics_list: list[dict[str, Any]] = [{} for _ in range(len(all_messages_list))]
@@ -197,6 +199,7 @@ def evaluate_chat_response(  # noqa: C901,PLR0912, PLR0915
                 for messages, extra_info in zip(all_messages_list, extra_info_list)
             ],
         )
+        metric.cleanup_resources()
 
         metrics_summary_dict.update(metric_result.summary)
 
@@ -217,6 +220,10 @@ def evaluate_chat_response(  # noqa: C901,PLR0912, PLR0915
                 tool_call_validation_result_counter[mes["tool_call_validation_result"]] += 1
     for finish_reason, count in finish_reason_counter.items():
         metrics_summary_dict[f"finish_reason_ratio-{finish_reason}"] = count / sum(finish_reason_counter.values())
+    for validation_result, count in tool_call_validation_result_counter.items():
+        metrics_summary_dict[f"tool_call_validation_result_ratio-{validation_result}"] = count / sum(
+            tool_call_validation_result_counter.values()
+        )
 
     logger.info(metrics_summary_dict)
 
@@ -229,6 +236,11 @@ def evaluate_chat_response(  # noqa: C901,PLR0912, PLR0915
             **instance_metrics,
         }
         | ({"raw_lm_output": messages[-1]["raw_content"]} if "raw_content" in messages[-1] else {})
+        | (
+            {"tool_call_validation_result": messages[-1]["tool_call_validation_result"]}
+            if "tool_call_validation_result" in messages[-1]
+            else {}
+        )
         for messages, references, extra_info, instance_metrics in zip(
             all_messages_list,
             references_list,

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/evaluate_from_data.py

@@ -45,6 +45,7 @@ def evaluate_from_data(
             references_list=references_list,
             extra_info_list=extra_info_list,
         )
+        metric.cleanup_resources()
 
         metrics_summary_dict.update(metric_result.summary)
 

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/evaluate_generation.py

@@ -69,6 +69,8 @@ def evaluate_generation(  # noqa: C901
 
             pbar.update(len(batch))
 
+    language_model.cleanup_resources()
+
     # Evaluate the generated continuations
     metrics_summary_dict: dict[str, float] = {}
     instance_metrics_list: list[dict[str, Any]] = [{} for _ in range(len(eval_instances))]
@@ -78,6 +80,7 @@ def evaluate_generation(  # noqa: C901
             references_list=[i.references for i in eval_instances],
             extra_info_list=[i.inputs for i in eval_instances],
         )
+        metric.cleanup_resources()
 
         metrics_summary_dict.update(metric_result.summary)
 

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/language_model/__init__.py

@@ -4,3 +4,4 @@ from .litellm_api import LiteLLMChatAPI
 from .openai_api import OpenAIChatAPI, OpenAICompletionAPI
 from .openai_batch_api import OpenAIChatBatchAPI
 from .vllm_model import VLLM
+from .vllm_serve_lm import VLLMServeLM

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/language_model/base.py

@@ -240,6 +240,15 @@ class LanguageModel:
             return self._batch_compute_chat_log_probs([prompt], [response])[0]
         return self._batch_compute_chat_log_probs(prompt, response)
 
+    def cleanup_resources(self) -> None:
+        """
+        Clean up resources if necessary.
+        This method is called when the language model is no longer needed.
+        """
+
+    def __del__(self) -> None:
+        self.cleanup_resources()
+
 
 def normalize_stop_sequences(
     stop_sequences_list: list[str | list[str] | None],

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/language_model/hf_lm.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 import contextlib
-from typing import Any, Literal, TypeVar
+import copy
+import gc
+import json
+from typing import Any, Callable, Literal, TypeVar
 
 import torch
 import torch.nn.functional as F  # noqa: N812
@@ -138,6 +141,25 @@ def decode_for_lm_continuation(
     return entire_text[len(input_text) :]
 
 
+def deserialize_tool_calls_in_messages(messages: list[dict[str, Any]]) -> None:
+    """
+    We adopt the standard OpenAI format, where the 'arguments' field in tool_calls is expected to be a JSON string.
+    However, huggingface/transformers expects the 'arguments' field to be a dict.
+    https://huggingface.co/docs/transformers/v4.48.2/chat_templating#a-complete-tool-use-example
+
+    To resolve this mismatch, this function deserializes 'arguments' before passing messages to apply_chat_template.
+
+    Args:
+        messages: A list of messages to deserialize.
+    """
+    deserialized_messages = copy.deepcopy(messages)
+    for message in deserialized_messages:
+        if message["role"] == "assistant" and "tool_calls" in message:
+            for item in message["tool_calls"]:
+                item["function"]["arguments"] = json.loads(item["function"]["arguments"])
+    return deserialized_messages
+
+
 class HuggingFaceLM(LanguageModel):
     """
     LanguageModel implementation using Hugging Face Transformers.
@@ -194,45 +216,56 @@ class HuggingFaceLM(LanguageModel):
         self.chat_template_kwargs = chat_template_kwargs or {}
         self.add_special_tokens = add_special_tokens
         self.default_gen_kwargs = default_gen_kwargs or {}
-
-        model_kwargs = get_default_model_kwargs(model_kwargs)
-        if not load_peft:
-            self.model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
-                model,
-                **model_kwargs,
-            )
-        else:
-            from peft import AutoPeftModelForCausalLM
-
-            self.model = AutoPeftModelForCausalLM.from_pretrained(
-                model,
-                **model_kwargs,
-            )
-
-        self.model.eval()
-
+        # `self.model` is initialized lazily to avoid unnecessary memory usage.
+        self.model: PreTrainedModel | None = None
+        self.model_kwargs = get_default_model_kwargs(model_kwargs)
+        self.load_peft = load_peft
         self.amp_dtype = amp_dtype
-        if model_limit_tokens == "default":
-            hf_config = self.model.config.to_dict()
-            if "n_positions" in hf_config:
-                model_limit_tokens = hf_config["n_positions"]
-            elif "max_position_embeddings" in hf_config:
-                model_limit_tokens = hf_config["max_position_embeddings"]
-            else:
-                msg = (
-                    "`model_limit_tokens` was set to “default”, but the default max_position_embedeings "
-                    "could not be found in the config. Set it to `None`."
-                )
-                logger.warning(msg)
         self.model_limit_tokens = model_limit_tokens
         self.tool_parser = tool_parser
-
-        transformers.set_seed(random_seed)
-
-        logger.info(f"model device: {self.model.device}")
-        logger.info(f"model dtype: {self.model.dtype}")
         logger.info(f"amp_dtype: {amp_dtype}")
         logger.info(f"random seed: {random_seed}")
+        transformers.set_seed(random_seed)
+
+    @staticmethod
+    def load_model(method: Callable) -> Callable:
+        """Decorator to load the model lazily."""
+
+        def wrapper(self: HuggingFaceLM, *args: tuple, **kwargs: dict) -> Callable:
+            if self.model is None:
+                if not self.load_peft:
+                    self.model = AutoModelForCausalLM.from_pretrained(
+                        self._model_name_or_path,
+                        **self.model_kwargs,
+                    )
+                else:
+                    from peft import AutoPeftModelForCausalLM
+
+                    self.model = AutoPeftModelForCausalLM.from_pretrained(
+                        self._model_name_or_path,
+                        **self.model_kwargs,
+                    )
+
+                self.model.eval()
+
+                if self.model_limit_tokens == "default":
+                    hf_config = self.model.config.to_dict()
+                    if "n_positions" in hf_config:
+                        self.model_limit_tokens = hf_config["n_positions"]
+                    elif "max_position_embeddings" in hf_config:
+                        self.model_limit_tokens = hf_config["max_position_embeddings"]
+                    else:
+                        msg = (
+                            "`model_limit_tokens` was set to “default”, but the default max_position_embedeings "
+                            "could not be found in the config. Set it to `None`."
+                        )
+                        logger.warning(msg)
+
+                logger.info(f"model device: {self.model.device}")
+                logger.info(f"model dtype: {self.model.dtype}")
+            return method(self, *args, **kwargs)
+
+        return wrapper
 
     def _get_amp_context(self) -> contextlib.AbstractContextManager:
         if self.amp_dtype is None:
@@ -276,6 +309,7 @@ class HuggingFaceLM(LanguageModel):
         return stop_token_ids
 
     @torch.inference_mode()
+    @load_model
     def _batch_complete_text(
         self,
         text_list: list[str],
@@ -350,6 +384,7 @@ class HuggingFaceLM(LanguageModel):
             lm_outputs.append(LMOutput(text=decoded_text, finish_reason=finish_reason))
         return lm_outputs
 
+    @load_model
     def _batch_generate_chat_response(
         self,
         chat_messages_list: list[list[dict[str, Any]]],
@@ -366,7 +401,7 @@ class HuggingFaceLM(LanguageModel):
                 chat_messages.insert(0, {"role": "system", "content": self.system_message})
         chat_messages_as_string = [
             self.tokenizer.apply_chat_template(
-                chat_messages,
+                deserialize_tool_calls_in_messages(chat_messages),
                 tools=tools,
                 tokenize=False,
                 add_generation_prompt=True,
@@ -390,6 +425,7 @@ class HuggingFaceLM(LanguageModel):
         return lm_outputs
 
     @torch.inference_mode()
+    @load_model
     def _batch_compute_log_probs(
         self,
         text_list: list[str],
@@ -496,6 +532,7 @@ class HuggingFaceLM(LanguageModel):
             total_log_probs = (log_prob_of_next * log_prob_mask).sum(dim=-1)
         return total_log_probs.tolist()
 
+    @load_model
     def _batch_compute_chat_log_probs(
         self, prompt_list: list[list[dict[str, Any]]], response_list: list[dict[str, Any]]
     ) -> list[float]:
@@ -512,6 +549,14 @@ class HuggingFaceLM(LanguageModel):
             response_as_string.append(response_as_string_i)
         return self._batch_compute_log_probs(response_as_string, prefix_list=prompt_as_string)
 
+    def cleanup_resources(self) -> None:
+        del self._model
+        self._model = None
+        gc.collect()
+        if torch.cuda.is_available():
+            logger.info("Cleaning up CUDA resources...")
+            torch.cuda.empty_cache()
+
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(model={self._model_name_or_path!r})"
 

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/language_model/openai_batch_api.py

@@ -279,7 +279,7 @@ class OpenAIChatBatchAPI(LanguageModel):
             for res in api_responses
         ]
 
-    def close(self) -> None:
+    def cleanup_resources(self) -> None:
         # in case that the program fails before the file is initialized in __init__
         if not hasattr(self, "temp_jsonl_file"):
             return
@@ -333,7 +333,7 @@ class OpenAIChatBatchAPI(LanguageModel):
         return log_probs
 
     def __del__(self) -> None:
-        self.close()
+        self.cleanup_resources()
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(model={self.model})"

{flexeval-0.12.2 → flexeval-0.13.1}/flexeval/core/language_model/vllm_model.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import Any, Literal
+import time
+from typing import TYPE_CHECKING, Any, Callable, Literal
 
 import torch
 from loguru import logger
@@ -10,7 +11,10 @@ from flexeval.core.string_processor import StringProcessor
 from flexeval.core.tool_parser.base import ToolParser
 
 from .base import LanguageModel, LMOutput, normalize_stop_sequences
-from .hf_lm import decode_for_lm_continuation, get_prefix_and_completion_from_chat
+from .hf_lm import decode_for_lm_continuation, deserialize_tool_calls_in_messages, get_prefix_and_completion_from_chat
+
+if TYPE_CHECKING:
+    from vllm import LLM
 
 
 def tokenize_text_for_lm_prefix(
@@ -122,9 +126,6 @@ class VLLM(LanguageModel):
         if "max_new_tokens" in self.default_gen_kwargs:
             self.default_gen_kwargs["max_tokens"] = self.default_gen_kwargs.pop("max_new_tokens")
 
-        # import from vllm here because it is an extra dependency
-        from vllm import LLM
-
         model_kwargs = model_kwargs or {}
         # automatically set tensor_parallel_size to the number of GPUs
         if "tensor_parallel_size" not in model_kwargs:
@@ -132,13 +133,28 @@ class VLLM(LanguageModel):
         if "enable_chunked_prefill" not in model_kwargs:
             model_kwargs["enable_chunked_prefill"] = True
             model_kwargs["disable_sliding_window"] = True
-        self.llm = LLM(model, **model_kwargs)
-
-        if model_limit_tokens == "default":
-            model_limit_tokens = self.llm.llm_engine.get_model_config().max_model_len
+        self.model_kwargs = model_kwargs
+        # `self.llm` is initialized lazily to avoid unnecessary memory usage.
+        self.llm: LLM | None = None
         self.model_limit_tokens = model_limit_tokens
         self.tool_parser = tool_parser
 
+    @staticmethod
+    def load_model(method: Callable) -> Callable:
+        """Decorator to load the model lazily."""
+
+        def wrapper(self: VLLM, *args: tuple, **kwargs: dict) -> Callable:
+            if self.llm is None:
+                from vllm import LLM
+
+                self.llm = LLM(self.model_name, **self.model_kwargs)
+                if self.model_limit_tokens == "default":
+                    self.model_limit_tokens = self.llm.llm_engine.get_model_config().max_model_len
+            return method(self, *args, **kwargs)
+
+        return wrapper
+
+    @load_model
     def _batch_complete_text(
         self,
         text_list: list[str],
@@ -215,6 +231,7 @@ class VLLM(LanguageModel):
             outputs.append(LMOutput(text=decoded_text, finish_reason=finish_reason))
         return outputs
 
+    @load_model
     def _batch_generate_chat_response(
         self,
         chat_messages_list: list[list[dict[str, Any]]],
@@ -228,7 +245,7 @@ class VLLM(LanguageModel):
                 chat_messages.insert(0, {"role": "system", "content": self.system_message})
         chat_messages_as_string = [
             self.tokenizer.apply_chat_template(
-                chat_messages,
+                deserialize_tool_calls_in_messages(chat_messages),
                 tools=tools,
                 tokenize=False,
                 add_generation_prompt=True,
@@ -250,6 +267,7 @@ class VLLM(LanguageModel):
 
         return lm_outputs
 
+    @load_model
     def _batch_compute_log_probs(
         self, text_list: list[str], prefix_list: list[str] | None = None, stride: int | None = None
     ) -> list[float]:
@@ -329,6 +347,7 @@ class VLLM(LanguageModel):
 
         return batch_logprobs
 
+    @load_model
     def _batch_compute_chat_log_probs(
         self, prompt_list: list[list[dict[str, Any]]], response_list: list[dict[str, Any]]
     ) -> list[float]:
@@ -345,5 +364,14 @@ class VLLM(LanguageModel):
             response_as_string.append(response_as_string_i)
         return self._batch_compute_log_probs(response_as_string, prefix_list=prompt_as_string)
 
+    def cleanup_resources(self) -> None:
+        from vllm.distributed import cleanup_dist_env_and_memory
+
+        del self.llm
+        logger.info("cleaning up vLLM resources...")
+        cleanup_dist_env_and_memory()
+        time.sleep(10)  # wait for the vLLM server to release resources
+        self.llm = None
+
     def __repr__(self) -> str:
         return f"VLLM(model_name={self.model_name})"