Merge pull request #3 from plonerma/main

Added Workflow for Testing Contributions to main

.github/workflows/test.yml

@@ -0,0 +1,47 @@
+name: Tests
+
+on:
+  push:
+    branches: [ "main" ]
+    paths-ignore:
+    - "examples/**"
+    - ".gitignore"
+    - "README.md"
+    - "LICENSE"
+  pull_request:
+    branches: [ "main" ]
+    paths-ignore:
+    - "examples/**"
+    - ".gitignore"
+    - "README.md"
+    - "LICENSE"
+jobs:
+  run:
+    name: "Run Tests"
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.8"]
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Set up pip cache
+      if: runner.os == 'Linux'
+      uses: actions/cache@v4
+      with:
+        path: ~/.cache/pip
+        key: ${{ runner.os }}-pip-${{ hashFiles('requirements.txt') }}
+        restore-keys: ${{ runner.os }}-pip-
+    - name: Install package
+      run: pip install -e .
+    - name: Install pytest
+      run: pip install pytest
+    - name: Print Package Versions
+      run: pip freeze
+    - name: Tests
+      run: pytest

.gitignore

@@ -0,0 +1,162 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+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/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.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
+.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/

pyproject.toml

@@ -0,0 +1,9 @@
+[tool.pytest.ini_options]
+pythonpath = [
+  "."
+]
+
+[tool.mypy]
+files="transformer_ranker,tests"
+ignore_missing_imports = true
+check_untyped_defs = true

transformer_ranker/datacleaner.py

@@ -5,7 +5,7 @@
 from .utils import configure_logger
 
 import logging
-from typing import List, Dict, Optional, Union, Tuple, Any
+from typing import List, Dict, Optional, Set, Union, Tuple, Type
 
 
 logger = configure_logger('transformer_ranker', logging.INFO)
@@ -123,7 +123,7 @@ def prepare_dataset(self, dataset: Union[str, DatasetDict, Dataset]) -> Union[Da
         self.log_dataset_info(dataset)
 
         # Simplify the dataset: keep only relevant columns
-        keep_columns = [self.text_column, self.text_pair_column, self.label_column]
+        keep_columns = [col for col in (self.text_column, self.text_pair_column, self.label_column) if col is not None]
         dataset = self._remove_columns(dataset, keep_columns=keep_columns)
 
         return dataset
@@ -147,7 +147,7 @@ def _downsample(dataset: Dataset, ratio: float) -> Dataset:
 
     @staticmethod
     def _find_text_and_label_columns(dataset: Dataset, text_column: Optional[str] = None,
-                                     label_column: Optional[str] = None) -> Tuple[str, str, type[Any]]:
+                                     label_column: Optional[str] = None) -> Tuple[str, str, Type]:
         """Determine text and label columns in hf datasets based on popular keywords"""
         # A list of mostly used column names for texts
         text_columns = [
@@ -186,7 +186,7 @@ def merge_texts(example: Dict[str, str]) -> Dict[str, str]:
         return dataset
 
     @staticmethod
-    def _find_task_type(label_column: str, label_type: Union[type(int), type(str), type(list), type(float)]) -> str:
+    def _find_task_type(label_column: str, label_type: Union[Type[int], Type[str], Type[list], Type[float]]) -> str:
         """Determine task type based on the label column's data type."""
         label_type_to_task_type = {
             int: "sentence classification",  # labels can be integers
@@ -214,11 +214,9 @@ def pre_tokenize(example):
         return dataset
 
     @staticmethod
-    def _merge_data_splits(dataset: Union[DatasetDict, List[Dataset]]) -> Dataset:
+    def _merge_data_splits(dataset: DatasetDict) -> Dataset:
         """Merge DatasetDict into a single dataset."""
-        datasets_to_merge = [dataset[split] for split in dataset.keys()]
-        merged_dataset = datasets.concatenate_datasets(datasets_to_merge)
-        return merged_dataset
+        return datasets.concatenate_datasets(list(dataset.values()))
 
     @staticmethod
     def _remove_empty_rows(dataset: Dataset, text_column: str, label_column: str) -> Dataset:
@@ -285,7 +283,7 @@ def _change_to_span_encoding(
             else:
                 # Create label map manually if not found
                 logger.info('Label map not found. Creating manually...')
-                unique_labels = set()
+                unique_labels: Set[str] = set()
                 label_data = dataset[label_column] if isinstance(dataset, Dataset) else [dataset[split][label_column]
                                                                                          for split in dataset]
                 for label_list in label_data:
@@ -296,7 +294,7 @@ def _change_to_span_encoding(
         logger.info(f"Label map: {label_map}")
 
         # Remove BIO encoding from the label map
-        span_label_map = {}
+        span_label_map: Dict[str, int] = {}
         for label in label_map:
             main_label = label.split('-')[-1] if isinstance(label, str) else label
             if main_label not in span_label_map:

transformer_ranker/embedder.py

@@ -5,12 +5,14 @@
 from tqdm import tqdm
 from typing import Optional, List, Union
 
+from transformers import PreTrainedTokenizerFast
+
 
 class Embedder:
     def __init__(
         self,
         model: Union[str, torch.nn.Module],
-        tokenizer: Union[str, AutoTokenizer] = None,
+        tokenizer: Union[str, PreTrainedTokenizerFast, None] = None,
         layer_ids: str = "all",
         subword_pooling: str = "mean",
         layer_pooling: Optional[str] = None,
@@ -22,7 +24,7 @@ def __init__(
         """
         Embed sentences using a pre-trained transformer model. It works at the word level, meaning each sentence
         is represented by a list of word vectors. You can pool these into a single sentence embedding if needed.
-        ♻️ Feel free to use it if you ever need a simple implementation for transformer embeddings.
+        ♻️  Feel free to use it if you ever need a simple implementation for transformer embeddings.
 
         :param model: Name of the model to be used. Either a model handle (e.g. 'bert-base-uncased')
         or a loaded model e.g. AutoModel('bert-base-uncased').
@@ -46,7 +48,16 @@ def __init__(
             self.model_name = model
 
         # Load a model-specific tokenizer
-        self.tokenizer = tokenizer or AutoTokenizer.from_pretrained(self.model_name, add_prefix_space=True)
+        self.tokenizer: PreTrainedTokenizerFast
+
+        if tokenizer is None:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, add_prefix_space=True)
+
+        elif isinstance(tokenizer, str):
+            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer, add_prefix_space=True)
+
+        else:
+            self.tokenizer = tokenizer
 
         # Add padding token for models that do not have it (e.g. GPT2)
         if self.tokenizer.pad_token is None:
@@ -70,8 +81,12 @@ def __init__(
         self.sentence_pooling = sentence_pooling
 
         # Set cpu or gpu device
-        self.device = (device if device else
-                       torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        if device is None:
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        else:
+            self.device = torch.device(device)
+
         self.model = self.model.to(self.device)
 
     def tokenize(self, sentences):
@@ -153,10 +168,10 @@ def embed_batch(self, sentences, move_embeddings_to_cpu: bool = True) -> List[to
         for subword_embeddings, word_ids in zip(embeddings, word_ids):
 
             # Pool sub-words to get word-level embeddings
-            word_embeddings = self._pool_subwords(subword_embeddings, word_ids)
+            word_embedding_list = self._pool_subwords(subword_embeddings, word_ids)
 
             # Stack all word-level embeddings that represent a sentence
-            word_embeddings = torch.stack(word_embeddings, dim=0)
+            word_embeddings = torch.stack(word_embedding_list, dim=0)
 
             # Pool word-level embeddings into a single sentence vector if specified
             sentence_embedding = self._pool_words(word_embeddings) if self.sentence_pooling else word_embeddings
@@ -170,11 +185,11 @@ def embed_batch(self, sentences, move_embeddings_to_cpu: bool = True) -> List[to
 
         return sentence_embeddings
 
-    def _filter_layer_ids(self, layer_ids):
+    def _filter_layer_ids(self, layer_ids) -> List[int]:
         """Transform a string with layer ids into a list of ints and
          remove ids that are out of bound of the actual transformer size"""
         if layer_ids == "all":
-            layer_ids = ", ".join([str(-1 * (i + 1)) for i in range(self.num_transformer_layers)])
+            return [-i for i in range(1, self.num_transformer_layers + 1)]
 
         layer_ids = [int(number) for number in layer_ids.split(",")]
 
@@ -199,9 +214,9 @@ def _extract_relevant_layers(self, batched_embeddings: torch.Tensor) -> torch.Te
     def _pool_subwords(self, sentence_embedding, sentence_word_ids) -> List[torch.Tensor]:
         """Pool sub-word embeddings into word embeddings for a single sentence.
         Subword pooling methods: 'first', 'last', 'mean'"""
-        word_embeddings = []
-        subword_embeddings = []
-        previous_word_id = 0
+        word_embeddings: List[torch.Tensor] = []
+        subword_embeddings: List[torch.Tensor] = []
+        previous_word_id: int = 0
 
         # Gather word-level embeddings as lists of subwords
         for token_embedding, word_id in zip(sentence_embedding, sentence_word_ids):

transformer_ranker/estimators/hscore.py

@@ -58,7 +58,7 @@ def fit(self, embeddings: torch.Tensor, labels: torch.Tensor) -> float:
             class_means[i] = class_features * torch.sqrt(mask.sum())
 
         # Covariance for class-conditioned means
-        covg = torch.mm(class_means.T, class_means) / num_samples
+        covg = torch.mm(class_means.T, class_means) / (num_samples - 1)
 
         # Shrinkage-based H-score
         hscore = torch.trace(torch.mm(pinv_covf_alpha, (1 - shrinkage) * covg)).item()