Refactorings

Author: lukasgarbas

.gitignore

@@ -1,12 +1,7 @@
-# Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
 *$py.class
-
-# C extensions
 *.so
-
-# Distribution / packaging
 .Python
 build/
 develop-eggs/
@@ -24,19 +19,12 @@ share/python-wheels/
 *.egg-info/
 .installed.cfg
 *.egg
+space-invaders.py
 MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
 *.manifest
 *.spec
-
-# Installer logs
 pip-log.txt
 pip-delete-this-directory.txt
-
-# Unit test / coverage reports
 htmlcov/
 .tox/
 .nox/
@@ -50,115 +38,31 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
 .pybuilder/
 target/
-
-# Jupyter Notebook
 .ipynb_checkpoints
-
-# IPython
 profile_default/
 ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.python-version
+poetry.lock
 .pdm.toml
 .pdm-python
 .pdm-build/
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
 __pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
 .env
 .venv
 env/
 venv/
 ENV/
 env.bak/
 venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
 /site
-
-# mypy
 .mypy_cache/
 .dmypy.json
 dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
 .pytype/
-
-# Cython debug symbols
 cython_debug/
-
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 .vscode/
-.DS_Store
+.DS_Store
+.ruff_cache

tests/test_ranker.py

@@ -47,3 +47,5 @@ def test_ranker_bestlayer(small_language_models, trec):
     ranker = TransformerRanker(dataset=trec, dataset_downsample=0.05)
     result = ranker.run(small_language_models, layer_aggregator='bestlayer')
     assert len(str(result).split("\n")) >= 2
+
+    assert isinstance(result.layerwise_scores, dict)  # see if layer scores are there

transformer_ranker/datacleaner.py

@@ -40,14 +40,12 @@ class DatasetCleaner:
     def prepare_dataset(
         self, dataset: Union[str, Dataset, DatasetDict]
     ) -> tuple[Union[list[str], list[list[str]]], torch.Tensor, TaskCategory]:
-        """Prepare texts and labels, assign task category.
+        """Prepares texts, labels, and assigns the task category.
 
-        Downsample dataset, find text and label columns, create label map,
-        preprocess labels, pre-tokenize, clean rows, merge text pair columns.
+        Downsamples dataset, finds text and label columns, cleans empty/noisy rows, 
+        pre-tokenizes texts, merges text pair columns, creates label map for classification.
         Returns: (processed texts, label tensor, task category)
         """
-
-        # Verify dataset type
         if not isinstance(dataset, (Dataset, DatasetDict)):
             raise ValueError(f"Unsupported dataset type: {type(dataset)}")
 
@@ -85,14 +83,12 @@ def prepare_dataset(
             if task_category == TaskCategory.TOKEN_CLASSIFICATION and self.remove_bio_encoding:
                 dataset, label_map = self._remove_bio_encoding(dataset, label_column, label_map)
 
-        # Prepare all texts and labels as tensors
         texts = dataset[text_column]
         labels = dataset[label_column]
         if task_category == TaskCategory.TOKEN_CLASSIFICATION:
             labels = [word_label for labels in dataset[label_column] for word_label in labels]
         labels = torch.tensor(labels)
 
-        # Log dataset info
         self._log_dataset_info(
             text_column,
             label_column,

transformer_ranker/ranker.py

@@ -23,15 +23,15 @@ def __init__(
         **kwargs: Any,
     ):
         """
-        Prepares a dataset and compiles metrics to assess transferability.
+        Prepares huggingface text dataset (downsamples and finds the task category).
+        Sets up transferability metrics.
 
         :param dataset: a dataset from huggingface with texts and labels.
-        :param dataset_downsample: a fraction to which the dataset should be reduced.
+        :param dataset_downsample: a fraction to downsample the dataset to.
         :param text_column: the name of the column containing texts.
         :param label_column: the name of the column containing labels.
-        :param kwargs: additional dataset-specific parameters for data cleaning.
+        :param kwargs: additional parameters for dataset preprocessing.
         """
-        # Preprocess and down-sample a dataset
         datacleaner = DatasetCleaner(
             dataset_downsample=dataset_downsample,
             text_column=text_column,
@@ -46,6 +46,7 @@ def __init__(
             "logme": LogME,
             "hscore": HScore,
             "knn": NearestNeighbors,
+            # add more
         }
 
     def run(
@@ -57,23 +58,22 @@ def run(
         **kwargs: Any,
     ):
         """
-        Loads models, collects embeddings, and scores them.
+        Loads language models, collects embedding from each, and scores them.
 
         :param models: A list of model names
         :param estimator: Transferability metric ('hscore', 'logme', 'knn').
-        :param layer_aggregator: Method to aggregate layers ('lastlayer', 'layermean', 'bestlayer').
+        :param layer_aggregator: Method to aggregate layers ('layermean', 'lastlayer', 'bestlayer').
         :param batch_size: Number of samples per batch, defaults to 32.
         :param device: Device for embedding ('cpu', 'cuda', 'cuda:2').
-        :param gpu_estimation: Store and score embeddings on GPU for speedup.
+        :param gpu_estimation: Boolean if to compute transferability estimation using gpu.
         :param kwargs: Additional parameters for embedder class.
         :return: Returns sorted dictionary of model names and their scores
         """
         self._confirm_ranker_setup(estimator=estimator, layer_aggregator=layer_aggregator)
 
-        # Set device for models and the metric
         device = kwargs.pop("device", None)
-        gpu_estimation = kwargs.get("gpu_estimation", True)
-        if gpu_estimation:
+        estimation_using_gpu = kwargs.get("gpu_estimation", True)
+        if estimation_using_gpu:
             self.labels = self.labels.to(device)
 
         # Download models to hf cache
@@ -91,7 +91,6 @@ def run(
                 else kwargs.get("sentence_pooling", "mean")
             )
 
-            # Setup the embedder
             embedder = Embedder(
                 model=model,
                 layer_ids="0" if layer_aggregator == "lastlayer" else "all",
@@ -101,26 +100,27 @@ def run(
                 **kwargs,
             )
 
-            # Collect embeddings
+            # Collect language model embeddings
             embeddings = embedder.embed(
-                self.texts, batch_size=batch_size, unpack_to_cpu=not gpu_estimation, show_progress=True,
+                self.texts, batch_size=batch_size, unpack_to_cpu=not estimation_using_gpu, show_progress=True,
             )  # fmt: skip
 
-            # Flatten them for ner tasks
             if self.task_category == TaskCategory.TOKEN_CLASSIFICATION:
                 embeddings = [word for sentence in embeddings for word in sentence]
 
             model_name = embedder.name
-            del embedder  # remove from memory
+            del embedder  # remove language model from memory
             torch.cuda.empty_cache()
 
             # Compute transferability
             score = self._transferability_score(embeddings, metric, layer_aggregator)
 
-            # Store and log results
             result.add_score(model_name, score)
             logger.info(f"{model_name} {result.metric}: {score:.4f}")
 
+        logger.info(f"Results ▲\n{result}")
+        logger.info(f"Done!")
+
         return result
 
     def _transferability_score(self, embeddings, metric, layer_aggregator, show_progress=True) -> float:
@@ -132,7 +132,7 @@ def _transferability_score(self, embeddings, metric, layer_aggregator, show_prog
             range(num_layers), desc="Transferability score", bar_format=tqdm_bar_format, disable=not show_progress
         )
 
-        # Score each layer separately
+        # Score each hidden layer separately
         for layer_id in transferability_progress:
             layer_embeddings = torch.stack([emb[layer_id] for emb in embeddings])
             score = metric.fit(embeddings=layer_embeddings, labels=self.labels)

transformer_ranker/utils.py

@@ -57,12 +57,12 @@ def prepare_popular_models(model_size="base") -> list[str]:
 
 def configure_logger(name: str, level: int = logging.INFO, log_to_console: bool = True) -> logging.Logger:
     """
-    Configure the package's logger.
+    Configure logger for this transferability NLP framework
 
-    :param name: The name of the logger.
+    :param name: The name of the logger "transformer-ranker".
     :param level: The logging level (default: logging.INFO).
     :param log_to_console: Whether to log to console (default: True)
-    :return: Configured TransformerRanker logger
+    :return: Configured logger
     """
     logger = logging.getLogger(name)
     logger.setLevel(level)
@@ -73,12 +73,12 @@ def configure_logger(name: str, level: int = logging.INFO, log_to_console: bool
         console_handler.setFormatter(logging.Formatter("transformer_ranker:%(message)s"))
         logger.addHandler(console_handler)
 
-    # Suppress future and user warnings
+    # Suppress future and user warnings 
     warnings.simplefilter("ignore", category=FutureWarning)
     warnings.simplefilter("ignore", category=UserWarning)
     transformers_logging.set_verbosity_error()
 
-    # Suppress unused weights messages when loading models
+    # Suppress unused weights messages when loading models (transformers)
     logger.addFilter(
         lambda record: not (
             "Some weights of BertModel were not initialized" in record.getMessage()
@@ -96,21 +96,21 @@ class Result:
     def __init__(self, metric: str):
         self.metric = metric
         self.scores = {}
-        self.layer_scores = {}
+        self.layerwise_scores = {}
 
-    def add_score(self, model_name: str, score: float, layer_scores: list[float] = None) -> None:
-        """Add score for a model."""
+    def add_score(self, model_name: str, score: float) -> None:
         self.scores[model_name] = score
 
-        if layer_scores is not None:  # only used for the bestlayer option
-            self.layer_scores[model_name] = layer_scores
+    def add_layerwise_scores(self, model_name: str, scores: dict) -> None:
+        """Only used for the bestlayer option."""
+        self.layerwise_scores[model_name] = scores
 
     def append(self, other: "Result") -> None:
         """Append scores from multiple runs."""
         if self.metric != other.metric:
             raise ValueError(f"Metrics do not match ({self.metric} vs {other.metric}).")
         self.scores.update(other.scores)
-        self.layer_scores.update(other.layer_scores)
+        self.layerwise_scores.update(other.layerwise_scores)
 
     def best_model(self) -> str:
         """Show the model with the highest score."""
@@ -126,6 +126,6 @@ def __str__(self) -> str:
         sorted_scores = sorted(self.scores.items(), key=lambda item: item[1], reverse=True)
         model_rank = [f"Rank {i+1}. {model}: {score:.4f}" for i, (model, score) in enumerate(sorted_scores)]
         return "\n".join(model_rank)
-    
+
     def __repr__(self) -> str:
         return self.__str__()