evalscope 0.7.2__py3-none-any.whl → 0.8.0__py3-none-any.whl

evalscope/backend/rag_eval/clip_benchmark/utils/webdataset_convert.py

@@ -1,11 +1,12 @@
 # Convert datasets to webdataset format
 import os
-from tqdm import tqdm
 import torch
 import torch.utils.data
 import webdataset
+from tqdm import tqdm
+
 from evalscope.backend.rag_eval.clip_benchmark.dataset_builder import DatasetWrapper
-from evalscope.backend.rag_eval.utils.tools import path_to_bytes, PIL_to_bytes
+from evalscope.backend.rag_eval.utils.tools import PIL_to_bytes, path_to_bytes
 from evalscope.utils.logger import get_logger
 
 logger = get_logger()
@@ -17,19 +18,21 @@ def convert_dataset(
     output_folder,
     *,
     transform=None,
-    image_format="webp",
+    image_format='webp',
     max_count=10_000,
     max_size=1_000_000_000,
     multilabel=False,
     verbose=True,
 ):
     """
-    Convert an iterable `dataset` of (image, label) pairs to webdataset (.tar) format, and store in `output_folder/split`.
+    Convert an iterable `dataset` of (image, label) pairs to webdataset (.tar) format, and store in
+       `output_folder/split`.
 
     Images may be passed in as either:
     * File paths: pass in `transform=path_to_bytes`;
     * PIL images: pass in `transform=PIL_to_bytes(image_format)` where `image_format` is e.g. "webp"; or
-    * Raw binary data: use a PyTorch `Dataset` that supports `transform=PIL_to_bytes(image_format)`, and pass in `transform=None` here.
+    * Raw binary data: use a PyTorch `Dataset` that supports `transform=PIL_to_bytes(image_format)`, and
+        pass in `transform=None` here.
         Be sure that the transform is not applied twice.
 
     Copying image files directly or writing raw binary data is fastest since it allows multiprocessing;
@@ -37,9 +40,10 @@ def convert_dataset(
 
     Labels must be zero-indexed integers (for multilabel datasets, labels must be arrays/tensors).
 
-    Classnames and zero-shot classification templates can be provided as attributes of the dataset (`.classes` and `.templates`)
-    or filled in manually afterward. `dataset.classes` should be a list of strings indexed by the labels,
-    and `dataset.templates` should be a list of strings containing `{c}` to specify where classnames are to be inserted.
+    Classnames and zero-shot classification templates can be provided as attributes of the dataset (`.classes`
+    and `.templates`) or filled in manually afterward. `dataset.classes` should be a list of strings indexed by
+    the labels, and `dataset.templates` should be a list of strings containing `{c}` to specify where classnames
+    are to be inserted.
     """
     # Create output directory
     os.makedirs(os.path.join(output_folder, split), exist_ok=True)
@@ -52,52 +56,44 @@ def convert_dataset(
     )
     if verbose:
         try:
-            logger.info(f"Dataset size: {len(dataset)}")
+            logger.info(f'Dataset size: {len(dataset)}')
         except TypeError:
-            logger.info("IterableDataset has no len()")
+            logger.info('IterableDataset has no len()')
     # Save classnames
-    if hasattr(dataset, "classes") and dataset.classes:
-        classnames_fname = os.path.join(output_folder, "classnames.txt")
-        with open(classnames_fname, "w") as classnames_file:
-            logger.info(*dataset.classes, sep="\n", end="\n", file=classnames_file)
+    if hasattr(dataset, 'classes') and dataset.classes:
+        classnames_fname = os.path.join(output_folder, 'classnames.txt')
+        with open(classnames_fname, 'w') as classnames_file:
+            logger.info(*dataset.classes, sep='\n', end='\n', file=classnames_file)
         if verbose:
             logger.info("Saved class names to '%s'" % classnames_fname)
     elif verbose:
-        logger.info("WARNING: No class names found")
+        logger.info('WARNING: No class names found')
     # Save zeroshot templates
-    if hasattr(dataset, "templates") and dataset.templates:
-        templates_fname = os.path.join(
-            output_folder, "zeroshot_classification_templates.txt"
-        )
-        with open(templates_fname, "w") as templates_file:
-            logger.info(*dataset.templates, sep="\n", end="\n", file=templates_file)
+    if hasattr(dataset, 'templates') and dataset.templates:
+        templates_fname = os.path.join(output_folder, 'zeroshot_classification_templates.txt')
+        with open(templates_fname, 'w') as templates_file:
+            logger.info(*dataset.templates, sep='\n', end='\n', file=templates_file)
         if verbose:
             logger.info("Saved class names to '%s'" % templates_fname)
     elif verbose:
-        logger.info("WARNING: No zeroshot classification templates found")
+        logger.info('WARNING: No zeroshot classification templates found')
     # Save dataset type
     if multilabel:
-        type_fname = os.path.join(output_folder, "dataset_type.txt")
-        with open(type_fname, "w") as type_file:
-            logger.info("multilabel", end="\n", file=type_file)
+        type_fname = os.path.join(output_folder, 'dataset_type.txt')
+        with open(type_fname, 'w') as type_file:
+            logger.info('multilabel', end='\n', file=type_file)
             if verbose:
                 logger.info("Saved dataset type to '%s'" % type_fname)
     # Write to TAR files
-    data_fname = os.path.join(output_folder, split, r"%d.tar")
+    data_fname = os.path.join(output_folder, split, r'%d.tar')
     sink = webdataset.ShardWriter(data_fname, maxcount=max_count, maxsize=max_size)
     nsamples = 0
-    label_type = "npy" if multilabel else "cls"
-    for index, (input, output) in enumerate(tqdm(dataloader, desc="Converting")):
+    label_type = 'npy' if multilabel else 'cls'
+    for index, (input, output) in enumerate(tqdm(dataloader, desc='Converting')):
         nsamples += 1
         if isinstance(input, str) and transform is path_to_bytes:
             # If copying file, determine image format from extension
-            extension = (
-                os.path.splitext(input)[1]
-                .replace(".", "")
-                .lower()
-                .replace("jpeg", "jpg")
-                or image_format
-            )
+            extension = (os.path.splitext(input)[1].replace('.', '').lower().replace('jpeg', 'jpg') or image_format)
         else:
             extension = image_format
         # Convert label if necessary
@@ -107,27 +103,22 @@ def convert_dataset(
             else:
                 output = output.item()
         # Write example
-        sink.write(
-            {
-                "__key__": "s%07d" % index,
-                extension: transform(input) if transform else input,
-                label_type: output,
-            }
-        )
+        sink.write({
+            '__key__': 's%07d' % index,
+            extension: transform(input) if transform else input,
+            label_type: output,
+        })
     num_shards = sink.shard
     sink.close()
     if verbose:
-        logger.info(
-            "Saved dataset to '%s'"
-            % data_fname.replace(r"%d", "{0..%d}" % (num_shards - 1))
-        )
+        logger.info("Saved dataset to '%s'" % data_fname.replace(r'%d', '{0..%d}' % (num_shards - 1)))
     # Save number of shards
-    nshards_fname = os.path.join(output_folder, split, "nshards.txt")
-    with open(nshards_fname, "w") as nshards_file:
-        logger.info(num_shards, end="\n", file=nshards_file)
+    nshards_fname = os.path.join(output_folder, split, 'nshards.txt')
+    with open(nshards_fname, 'w') as nshards_file:
+        logger.info(num_shards, end='\n', file=nshards_file)
     if verbose:
         logger.info("Saved number of shards = %d to '%s'" % (num_shards, nshards_fname))
-    logger.info("Final dataset size:", nsamples)
+    logger.info('Final dataset size:', nsamples)
 
 
 def convert_retrieval_dataset(
@@ -136,13 +127,14 @@ def convert_retrieval_dataset(
     output_folder,
     *,
     transform=None,
-    image_format="webp",
+    image_format='webp',
     max_count=10_000,
     max_size=1_000_000_000,
     verbose=True,
 ):
     """
-    Convert an iterable `dataset` of (image, [caption1, caption2, ...]) pairs to webdataset (.tar) format, and store in `output_folder/split`.
+    Convert an iterable `dataset` of (image, [caption1, caption2, ...]) pairs to webdataset (.tar) format,
+    and store in `output_folder/split`.
 
     Labels must be lists of strings, with no newlines.
 
@@ -159,72 +151,59 @@ def convert_retrieval_dataset(
     )
     if verbose:
         try:
-            logger.info(f"Dataset size: {len(dataset)}")
+            logger.info(f'Dataset size: {len(dataset)}')
         except TypeError:
-            logger.info("IterableDataset has no len()")
+            logger.info('IterableDataset has no len()')
     # No classnames
     # No zeroshot templates
     # Save dataset type
-    type_fname = os.path.join(output_folder, "dataset_type.txt")
-    with open(type_fname, "w") as type_file:
-        logger.info("retrieval", end="\n", file=type_file)
+    type_fname = os.path.join(output_folder, 'dataset_type.txt')
+    with open(type_fname, 'w') as type_file:
+        logger.info('retrieval', end='\n', file=type_file)
     if verbose:
         logger.info("Saved dataset type to '%s'" % type_fname)
     # Write to TAR files
-    data_fname = os.path.join(output_folder, split, r"%d.tar")
+    data_fname = os.path.join(output_folder, split, r'%d.tar')
     sink = webdataset.ShardWriter(data_fname, maxcount=max_count, maxsize=max_size)
     nsamples = 0
-    for index, (input, output) in enumerate(tqdm(dataloader, desc="Converting")):
+    for index, (input, output) in enumerate(tqdm(dataloader, desc='Converting')):
         nsamples += 1
         if isinstance(input, str) and transform is path_to_bytes:
             # If copying file, determine image format from extension
-            extension = (
-                os.path.splitext(input)[1]
-                .replace(".", "")
-                .lower()
-                .replace("jpeg", "jpg")
-                or image_format
-            )
+            extension = (os.path.splitext(input)[1].replace('.', '').lower().replace('jpeg', 'jpg') or image_format)
         else:
             extension = image_format
-        sink.write(
-            {
-                "__key__": "s%07d" % index,
-                extension: transform(input) if transform else input,
-                "txt": "\n".join(caption.replace("\n", r"\n") for caption in output),
-            }
-        )
+        sink.write({
+            '__key__': 's%07d' % index,
+            extension: transform(input) if transform else input,
+            'txt': '\n'.join(caption.replace('\n', r'\n') for caption in output),
+        })
     num_shards = sink.shard
     sink.close()
     if verbose:
-        logger.info(
-            "Saved dataset to '%s'"
-            % data_fname.replace(r"%d", "{0..%d}" % (num_shards - 1))
-        )
+        logger.info("Saved dataset to '%s'" % data_fname.replace(r'%d', '{0..%d}' % (num_shards - 1)))
     # Save number of shards
-    nshards_fname = os.path.join(output_folder, split, "nshards.txt")
-    with open(nshards_fname, "w") as nshards_file:
-        logger.info(num_shards, end="\n", file=nshards_file)
+    nshards_fname = os.path.join(output_folder, split, 'nshards.txt')
+    with open(nshards_fname, 'w') as nshards_file:
+        logger.info(num_shards, end='\n', file=nshards_file)
     if verbose:
         logger.info("Saved number of shards = %d to '%s'" % (num_shards, nshards_fname))
-    logger.info("Final dataset size:", nsamples)
+    logger.info('Final dataset size:', nsamples)
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     from modelscope.msdatasets import MsDataset
 
-    splits = ["train", "validation"]
+    splits = ['train', 'validation']
     for split in splits:
-        ds = MsDataset.load("modelscope/muge", split=split)
+        ds = MsDataset.load('modelscope/muge', split=split)
         hf_dataset = ds.to_hf_dataset()
-        pytorch_dataset = DatasetWrapper(
-            hf_dataset, image_key="image", text_key="query"
-        )
+        pytorch_dataset = DatasetWrapper(hf_dataset, image_key='image', text_key='query')
         convert_retrieval_dataset(
             pytorch_dataset,
             split,
-            "data/muge",
-            transform=PIL_to_bytes("jpg"),
-            image_format="jpg",
+            'data/muge',
+            transform=PIL_to_bytes('jpg'),
+            image_format='jpg',
             max_count=50_000,
         )

evalscope/backend/rag_eval/cmteb/__init__.py

@@ -1,4 +1,4 @@
-from evalscope.backend.rag_eval.cmteb.tasks import *
+from evalscope.backend.rag_eval.cmteb.arguments import EvalArguments, ModelArguments
 from evalscope.backend.rag_eval.cmteb.base import *
-from evalscope.backend.rag_eval.cmteb.arguments import ModelArguments, EvalArguments
-from evalscope.backend.rag_eval.cmteb.task_template import one_stage_eval, two_stage_eval
+from evalscope.backend.rag_eval.cmteb.task_template import one_stage_eval, two_stage_eval
+from evalscope.backend.rag_eval.cmteb.tasks import *

evalscope/backend/rag_eval/cmteb/arguments.py

@@ -1,38 +1,36 @@
 from dataclasses import dataclass, field
-from typing import List, Optional, Union, Dict, Any
+from typing import Any, Dict, List, Optional, Union
 
 
 @dataclass
 class ModelArguments:
     # Arguments for embeding model: sentence transformer or cross encoder
-    model_name_or_path: str = ""  # model name or path
+    model_name_or_path: str = ''  # model name or path
     is_cross_encoder: bool = False  # whether the model is a cross encoder
     # pooling mode: Either “cls”, “lasttoken”, “max”, “mean”, “mean_sqrt_len_tokens”, or “weightedmean”.
     pooling_mode: Optional[str] = None
     max_seq_length: int = 512  # max sequence length
     # prompt for llm based model
-    prompt: str = ""
+    prompt: str = ''
     # model kwargs
     model_kwargs: dict = field(default_factory=dict)
     # config kwargs
     config_kwargs: Dict[str, Any] = field(default_factory=dict)
     # encode kwargs
-    encode_kwargs: dict = field(
-        default_factory=lambda: {"show_progress_bar": True, "batch_size": 32}
-    )
-    hub: str = "modelscope"  # modelscope or huggingface
+    encode_kwargs: dict = field(default_factory=lambda: {'show_progress_bar': True, 'batch_size': 32})
+    hub: str = 'modelscope'  # modelscope or huggingface
 
     def to_dict(self) -> Dict[str, Any]:
         return {
-            "model_name_or_path": self.model_name_or_path,
-            "is_cross_encoder": self.is_cross_encoder,
-            "pooling_mode": self.pooling_mode,
-            "max_seq_length": self.max_seq_length,
-            "prompt": self.prompt,
-            "model_kwargs": self.model_kwargs,
-            "config_kwargs": self.config_kwargs,
-            "encode_kwargs": self.encode_kwargs,
-            "hub": self.hub,
+            'model_name_or_path': self.model_name_or_path,
+            'is_cross_encoder': self.is_cross_encoder,
+            'pooling_mode': self.pooling_mode,
+            'max_seq_length': self.max_seq_length,
+            'prompt': self.prompt,
+            'model_kwargs': self.model_kwargs,
+            'config_kwargs': self.config_kwargs,
+            'encode_kwargs': self.encode_kwargs,
+            'hub': self.hub,
         }
 
 
@@ -42,20 +40,20 @@ class EvalArguments:
     tasks: List[str] = field(default_factory=list)  # task names
     dataset_path: Optional[str] = None  # custom dataset path
     verbosity: int = 2  # verbosity level 0-3
-    output_folder: str = "outputs"  # output folder
+    output_folder: str = 'outputs'  # output folder
     overwrite_results: bool = True  # overwrite results
     limits: Optional[int] = None  # limit number of samples
-    hub: str = "modelscope"  # modelscope or huggingface
-    top_k: int = 5 # top k for reranking
+    hub: str = 'modelscope'  # modelscope or huggingface
+    top_k: int = 5  # top k for reranking
 
     def to_dict(self) -> Dict[str, Any]:
         return {
-            "tasks": self.tasks,
-            "dataset_path": self.dataset_path,
-            "verbosity": self.verbosity,
-            "output_folder": self.output_folder,
-            "overwrite_results": self.overwrite_results,
-            "limits": self.limits,
-            "hub": self.hub,
-            "top_k": self.top_k,
+            'tasks': self.tasks,
+            'dataset_path': self.dataset_path,
+            'verbosity': self.verbosity,
+            'output_folder': self.output_folder,
+            'overwrite_results': self.overwrite_results,
+            'limits': self.limits,
+            'hub': self.hub,
+            'top_k': self.top_k,
         }

evalscope/backend/rag_eval/cmteb/base.py

@@ -1,12 +1,14 @@
+import datasets
 from collections import defaultdict
-from typing import List
-from mteb import AbsTask
 from datasets import DatasetDict
 from modelscope import MsDataset
-import datasets
-from evalscope.backend.rag_eval.cmteb.tasks import CLS_DICT, CLS_RETRIEVAL, CLS_CUSTOM
+from mteb import AbsTask
+from typing import List
+
+from evalscope.backend.rag_eval.cmteb.tasks import CLS_CUSTOM, CLS_DICT, CLS_RETRIEVAL
+from evalscope.constants import HubType
 
-__all__ = ["TaskBase"]
+__all__ = ['TaskBase']
 
 
 class TaskBase:
@@ -28,7 +30,7 @@ class TaskBase:
             from mteb.overview import TASKS_REGISTRY
 
             task_cls = TASKS_REGISTRY[task_name]
-            if task_cls.metadata.type != "Retrieval":
+            if task_cls.metadata.type != 'Retrieval':
                 task_cls.load_data = load_data
 
         # init task instance
@@ -41,33 +43,30 @@ def load_data(self, **kwargs):
     if self.data_loaded:
         return
 
-    limits = kwargs.get("limits", None)
-    hub = kwargs.get("hub", "modelscope")
-    name = self.metadata_dict.get("name")
-    path = self.metadata_dict["dataset"].get("path")
+    limits = kwargs.get('limits', None)
+    hub = kwargs.get('hub', HubType.MODELSCOPE)
+    name = self.metadata_dict.get('name')
+    path = self.metadata_dict['dataset'].get('path')
 
-    assert path is not None, "Path must be specified in dataset"
+    assert path is not None, 'Path must be specified in dataset'
 
     # Loading the dataset based on the source hub
-    if hub == "modelscope":
+    if hub == HubType.MODELSCOPE:
         import re
 
-        path = re.sub(r"^mteb/", "MTEB/", path)
+        path = re.sub(r'^mteb/', 'MTEB/', path)
         dataset = MsDataset.load(path)
     else:
-        dataset = datasets.load_dataset(**self.metadata_dict["dataset"])  # type: ignore
+        dataset = datasets.load_dataset(**self.metadata_dict['dataset'])  # type: ignore
 
     if limits is not None:
-        dataset = {
-            split: dataset[split].select(range(min(limits, len(dataset[split]))))
-            for split in dataset.keys()
-        }
+        dataset = {split: dataset[split].select(range(min(limits, len(dataset[split])))) for split in dataset.keys()}
 
     if name in CLS_RETRIEVAL:
         self.corpus, self.queries, self.relevant_docs = load_retrieval_data(
             dataset,
             path,
-            self.metadata_dict["eval_splits"],
+            self.metadata_dict['eval_splits'],
         )
 
     self.dataset = dataset
@@ -77,13 +76,13 @@ def load_data(self, **kwargs):
 
 def load_retrieval_data(dataset, dataset_name: str, eval_splits: list) -> tuple:
     eval_split = eval_splits[0]
-    qrels = MsDataset.load(dataset_name + "-qrels")[eval_split]
+    qrels = MsDataset.load(dataset_name + '-qrels')[eval_split]
 
-    corpus = {e["id"]: {"text": e["text"]} for e in dataset["corpus"]}
-    queries = {e["id"]: e["text"] for e in dataset["queries"]}
+    corpus = {e['id']: {'text': e['text']} for e in dataset['corpus']}
+    queries = {e['id']: e['text'] for e in dataset['queries']}
     relevant_docs = defaultdict(dict)
     for e in qrels:
-        relevant_docs[e["qid"]][e["pid"]] = e["score"]
+        relevant_docs[e['qid']][e['pid']] = e['score']
 
     corpus = DatasetDict({eval_split: corpus})
     queries = DatasetDict({eval_split: queries})

evalscope/backend/rag_eval/cmteb/task_template.py

@@ -1,8 +1,8 @@
-import os
 import mteb
-from evalscope.backend.rag_eval import EmbeddingModel
-from evalscope.backend.rag_eval import cmteb
+import os
 from mteb.task_selection import results_to_dataframe
+
+from evalscope.backend.rag_eval import EmbeddingModel, cmteb
 from evalscope.utils.logger import get_logger
 
 logger = get_logger()
@@ -19,8 +19,8 @@ def show_results(output_folder, model, results):
         model_name,
         revision,
     )
-    logger.info(f"Evaluation results:\n{results_df.to_markdown()}")
-    logger.info(f"Evaluation results saved in {os.path.abspath(save_path)}")
+    logger.info(f'Evaluation results:\n{results_df.to_markdown()}')
+    logger.info(f'Evaluation results saved in {os.path.abspath(save_path)}')
 
 
 def one_stage_eval(
@@ -29,18 +29,16 @@ def one_stage_eval(
 ) -> None:
     # load model
     model = EmbeddingModel.load(**model_args)
-    custom_dataset_path = eval_args.pop("dataset_path", None)
+    custom_dataset_path = eval_args.pop('dataset_path', None)
     # load task first to update instructions
-    tasks = cmteb.TaskBase.get_tasks(
-        task_names=eval_args["tasks"], dataset_path=custom_dataset_path
-    )
+    tasks = cmteb.TaskBase.get_tasks(task_names=eval_args['tasks'], dataset_path=custom_dataset_path)
     evaluation = mteb.MTEB(tasks=tasks)
 
     # run evaluation
     results = evaluation.run(model, **eval_args)
 
     # save and log results
-    show_results(eval_args["output_folder"], model, results)
+    show_results(eval_args['output_folder'], model, results)
 
 
 def two_stage_eval(
@@ -56,7 +54,7 @@ def two_stage_eval(
     first_stage_path = f"{eval_args['output_folder']}/stage1"
     second_stage_path = f"{eval_args['output_folder']}/stage2"
 
-    tasks = cmteb.TaskBase.get_tasks(task_names=eval_args["tasks"])
+    tasks = cmteb.TaskBase.get_tasks(task_names=eval_args['tasks'])
     for task in tasks:
         evaluation = mteb.MTEB(tasks=[task])
 
@@ -66,19 +64,19 @@ def two_stage_eval(
             save_predictions=True,
             output_folder=first_stage_path,
             overwrite_results=True,
-            hub=eval_args["hub"],
-            limits=eval_args["limits"],
+            hub=eval_args['hub'],
+            limits=eval_args['limits'],
         )
         # stage 2: run cross encoder
         results = evaluation.run(
             cross_encoder,
-            top_k=eval_args["top_k"],
+            top_k=eval_args['top_k'],
             save_predictions=True,
             output_folder=second_stage_path,
-            previous_results=f"{first_stage_path}/{task.metadata.name}_default_predictions.json",
+            previous_results=f'{first_stage_path}/{task.metadata.name}_default_predictions.json',
             overwrite_results=True,
-            hub=eval_args["hub"],
-            limits=eval_args["limits"],
+            hub=eval_args['hub'],
+            limits=eval_args['limits'],
         )
 
         # save and log results