Graded tasks integrations

CHANGELOG.md

@@ -2,6 +2,23 @@
 
 ### Feat
 
+- ``HouseGradedSkillExtractRanking``: graded-relevance skill-extraction ranking
+  task on the HOUSE subset of CAREER, re-annotated against the full ESCO v1.1.0
+  taxonomy with a 0-4 score scale (HF: ``TechWolf/Skill-extraction-House-graded``,
+  BEIR layout, validation split only).
+- ``TechGradedSkillExtractRanking``: graded-relevance skill-extraction ranking
+  task on the TECH subset of CAREER, re-annotated against the full ESCO v1.1.0
+  taxonomy with a 0-4 score scale (HF: ``TechWolf/Skill-extraction-Tech-graded``,
+  BEIR layout, validation split only).
+- ``SkillSkapeGradedSkillExtractRanking``: graded-relevance skill-extraction
+  ranking task on SkillSkape, re-annotated against the full ESCO v1.1.0 taxonomy
+  with a 0-4 score scale (HF: ``TechWolf/Skill-extraction-SkillSkape-graded``,
+  BEIR layout, validation split only).
+- ``ESCOGradedSkillNormRanking``: graded-relevance skill-normalization ranking
+  task that maps surface skill terms (ESCO alt-labels) to canonical ESCO skills,
+  annotated against the full ESCO v1.1.0 taxonomy with a 0-4 score scale
+  (HF: ``TechWolf/Skill-normalisation-ESCO-graded``, BEIR layout, validation
+  split only).
 - graded relevance support for ranking metrics: ``RankingDataset`` accepts an
   optional ``target_relevance`` field aligned 1-to-1 with ``target_indices``.
   ``ndcg@k`` uses a (2^rel - 1) gain when graded labels are provided; binary

README.md

@@ -85,12 +85,16 @@ pip install workrb
 | Job Normalization MELO         | `MELORanking` | multi_label | 633 queries x 33813 targets         | 21        |
 | Skill to Job WorkBench         | `ESCOSkill2JobRanking` | multi_label | 13492 queries x 3039 targets        | 28        |
 | Skill Extraction House         | `HouseSkillExtractRanking` | multi_label | 262 queries x 13891 targets         | 28        |
+| Skill Extraction House Graded  | `HouseGradedSkillExtractRanking` | multi_label | 61 queries x 13891 targets          | 1         |
 | Skill Extraction Tech          | `TechSkillExtractRanking` | multi_label | 338 queries x 13891 targets         | 28        |
+| Skill Extraction Tech Graded   | `TechGradedSkillExtractRanking` | multi_label | 75 queries x 13891 targets          | 1         |
 | Skill Extraction SkillSkape    | `SkillSkapeExtractRanking` | multi_label | 1191 queries x 13891 targets        | 28        |
+| Skill Extraction SkillSkape Graded | `SkillSkapeGradedSkillExtractRanking` | multi_label | 100 queries x 13891 targets         | 1         |
 | Skill Extraction TechWolf      | `TechWolfSkillExtractRanking` | multi_label | 326 queries x 13891 targets         | 28        |
 | Skill Extraction SkillXL       | `SkillXLSkillExtractRanking` | multi_label | 944 queries x 13891 targets         | 28        |
 | Skill Similarity SkillMatch-1K | `SkillMatch1kSkillSimilarityRanking` | single_label | 900 queries x 2648 targets          | 1         |
 | Skill Normalization ESCO       | `ESCOSkillNormRanking` | multi_label | 72008 queries x 13939 targets       | 28        |
+| Skill Normalization ESCO Graded | `ESCOGradedSkillNormRanking` | multi_label | 50 queries x 13891 targets          | 1         |
 | Skill Normalization MELS       | `MELSRanking` | multi_label | 1722 queries x 19466 targets        | 5         |
 | Query-Candidate Matching       | `SearchQueryCandidateRanking` | multi_label | 200 queries x 4019 (x-lang) targets | 5         |
 | Project-Candidate Matching     | `ProjectCandidateRanking` | multi_label | 200 queries x 4019 (x-lang) targets | 5         |

pyproject.toml

@@ -13,8 +13,7 @@ readme = "README.md"
 requires-python = ">=3.10,<4.0"
 license = { text = "Apache 2.0" }
 authors = [
-    { name = "TechWolf Research Team" },
-    { name = "Matthias De Lange", email = "matthias.delange@techwolf.ai" },
+    { name = "TechWolf Research Team", email = "workrb@techwolf.ai" },
 ]
 keywords = [
     "nlp",

src/workrb/rankings.py

@@ -140,9 +140,7 @@ def materialize_prediction_matrix(
     matrix = np.zeros((num_queries, num_targets), dtype=np.float32)
     for q_idx, row in scores.items():
         if not 0 <= q_idx < num_queries:
-            raise IndexError(
-                f"query_index {q_idx} out of bounds for num_queries={num_queries}"
-            )
+            raise IndexError(f"query_index {q_idx} out of bounds for num_queries={num_queries}")
         for t_idx, score in row.items():
             if not 0 <= t_idx < num_targets:
                 raise IndexError(

src/workrb/run.py

@@ -783,5 +783,11 @@ def _record_dataset_result(
     if config:
         config.save_results_checkpoint(results)
 
+    # Prefer the task's first default metric, but fall back to whatever was
+    # actually computed: a caller-supplied `metrics` set need not include the
+    # default key metric, and indexing it directly would raise KeyError.
     key_metric = task.default_metrics[0]
-    logger.info(f"\t{key_metric}: {metrics_dict[key_metric]:.3f}")
+    if key_metric not in metrics_dict:
+        key_metric = next(iter(metrics_dict), None)
+    if key_metric is not None:
+        logger.info(f"\t{key_metric}: {metrics_dict[key_metric]:.3f}")

src/workrb/tasks/__init__.py

@@ -19,14 +19,18 @@
 from .ranking.mels import MELSRanking
 from .ranking.skill2job import ESCOSkill2JobRanking
 from .ranking.skill_extraction import (
+    HouseGradedSkillExtractRanking,
     HouseSkillExtractRanking,
     SkillSkapeExtractRanking,
+    SkillSkapeGradedSkillExtractRanking,
     SkillXLSkillExtractRanking,
+    TechGradedSkillExtractRanking,
     TechSkillExtractRanking,
+    TechWolfGradedSkillExtractRanking,
     TechWolfSkillExtractRanking,
 )
 from .ranking.skill_similarity import SkillMatch1kSkillSimilarityRanking
-from .ranking.skillnorm import ESCOSkillNormRanking
+from .ranking.skillnorm import ESCOGradedSkillNormRanking, ESCOSkillNormRanking
 
 __all__ = [
     # Abstract classes
@@ -39,17 +43,22 @@
     # Classification tasks
     "ESCOJob2SkillClassification",
     # Ranking tasks
+    "ESCOGradedSkillNormRanking",
     "ESCOJob2SkillRanking",
     "ESCOSkill2JobRanking",
     "ESCOSkillNormRanking",
     "JobBERTJobNormRanking",
     "JobTitleSimilarityRanking",
     "MELORanking",
     "MELSRanking",
+    "HouseGradedSkillExtractRanking",
     "HouseSkillExtractRanking",
+    "TechGradedSkillExtractRanking",
     "TechSkillExtractRanking",
+    "TechWolfGradedSkillExtractRanking",
     "TechWolfSkillExtractRanking",
     "SkillSkapeExtractRanking",
+    "SkillSkapeGradedSkillExtractRanking",
     "SkillXLSkillExtractRanking",
     "SkillMatch1kSkillSimilarityRanking",
     "ProjectCandidateRanking",

src/workrb/tasks/abstract/ranking_base.py

@@ -476,4 +476,3 @@ def compute_metrics_from_prediction_matrix(
             pos_label_relevance=dataset.target_relevance,
             binary_relevance_threshold=self.binary_relevance_threshold,
         )
-

src/workrb/tasks/ranking/__init__.py

@@ -18,19 +18,25 @@
 from workrb.tasks.ranking.mels import MELSRanking
 from workrb.tasks.ranking.skill2job import ESCOSkill2JobRanking
 from workrb.tasks.ranking.skill_extraction import (
+    HouseGradedSkillExtractRanking,
     HouseSkillExtractRanking,
     SkillSkapeExtractRanking,
+    SkillSkapeGradedSkillExtractRanking,
     SkillXLSkillExtractRanking,
+    TechGradedSkillExtractRanking,
     TechSkillExtractRanking,
+    TechWolfGradedSkillExtractRanking,
     TechWolfSkillExtractRanking,
 )
 from workrb.tasks.ranking.skill_similarity import SkillMatch1kSkillSimilarityRanking
-from workrb.tasks.ranking.skillnorm import ESCOSkillNormRanking
+from workrb.tasks.ranking.skillnorm import ESCOGradedSkillNormRanking, ESCOSkillNormRanking
 
 __all__ = [
+    "ESCOGradedSkillNormRanking",
     "ESCOJob2SkillRanking",
     "ESCOSkill2JobRanking",
     "ESCOSkillNormRanking",
+    "HouseGradedSkillExtractRanking",
     "HouseSkillExtractRanking",
     "JobBERTJobNormRanking",
     "JobTitleSimilarityRanking",
@@ -40,7 +46,10 @@
     "SearchQueryCandidateRanking",
     "SkillMatch1kSkillSimilarityRanking",
     "SkillSkapeExtractRanking",
+    "SkillSkapeGradedSkillExtractRanking",
     "SkillXLSkillExtractRanking",
+    "TechGradedSkillExtractRanking",
     "TechSkillExtractRanking",
+    "TechWolfGradedSkillExtractRanking",
     "TechWolfSkillExtractRanking",
 ]

src/workrb/tasks/ranking/graded_beir.py

@@ -0,0 +1,146 @@
+"""Shared base for BEIR-layout graded-relevance ranking tasks.
+
+Several tasks publish graded (0-4) relevance annotations on the Hugging Face
+Hub following the BEIR convention (``queries``, ``corpus``, ``qrels`` configs).
+They differ only in their task group and query input type (e.g. skill
+extraction from sentences vs skill normalization from surface terms), so the
+data-loading logic lives here and concrete tasks supply the task-specific
+properties.
+"""
+
+import pandas as pd
+from datasets import Dataset, load_dataset
+
+from workrb.tasks.abstract.base import DatasetSplit, LabelType, Language
+from workrb.tasks.abstract.ranking_base import RankingDataset, RankingTask
+from workrb.types import ModelInputType
+
+
+class GradedBEIRRankingTask(RankingTask):
+    """Base class for BEIR-layout graded ranking tasks.
+
+    Reads the ``queries``, ``corpus`` and ``qrels`` configs published on the
+    Hugging Face Hub via ``load_dataset``. The target_space is the corpus's
+    ``title`` column (ESCO preferred labels, in corpus order); qrels are
+    expected to contain only non-zero judgments (absent items are implicit
+    grade 0).
+
+    Which splits a task exposes depends on what the underlying dataset
+    publishes: some release both a validation and a test split, others only
+    one. Because several tasks share this loader, the supported splits are
+    declared per subclass via :attr:`split_to_hf_split` (which maps each
+    supported :class:`DatasetSplit` to the HF split name its
+    ``queries``/``qrels`` configs live under) rather than hardcoded with an
+    inline guard as in the single-split tasks (e.g. ``MELORanking``). The
+    default exposes only the validation split, the common case for in-progress
+    benchmark datasets whose test split is withheld.
+
+    Concrete subclasses set ``hf_name`` via ``__init__`` and supply the
+    task-specific ``task_group``, ``query_input_type``, ``name``,
+    ``description`` and ``citation``.
+    """
+
+    def __init__(self, hf_name: str, **kwargs):
+        """Initialize the task.
+
+        Args:
+            hf_name: Name of the Hugging Face dataset (BEIR layout).
+            **kwargs: Additional arguments for the base class.
+        """
+        self.hf_name = hf_name
+        super().__init__(**kwargs)
+
+    @property
+    def split_to_hf_split(self) -> dict[DatasetSplit, str]:
+        """Map each supported split to the HF split name backing it.
+
+        The corpus config is always loaded from the ``corpus`` split; this
+        mapping only governs the ``queries`` and ``qrels`` configs. Override to
+        expose more or fewer splits, e.g. ``{DatasetSplit.VAL: "validation",
+        DatasetSplit.TEST: "test"}`` for a dataset that releases both.
+        """
+        return {DatasetSplit.VAL: "validation"}
+
+    @property
+    def supported_query_languages(self) -> list[Language]:
+        """Annotations are released in English only at this stage."""
+        return [Language.EN]
+
+    @property
+    def supported_target_languages(self) -> list[Language]:
+        """The corpus titles are released in English only at this stage."""
+        return [Language.EN]
+
+    @property
+    def label_type(self) -> LabelType:
+        """Label type is multi-label."""
+        return LabelType.MULTI_LABEL
+
+    @property
+    def target_input_type(self) -> ModelInputType:
+        """Target input type for ESCO skills."""
+        return ModelInputType.SKILL_NAME
+
+    @property
+    def default_metrics(self) -> list[str]:
+        """Default metrics include nDCG to leverage the graded labels.
+
+        ``ndcg`` without a cutoff scores the full ranking (k = |target_space|).
+        """
+        return ["ndcg", "ndcg@5", "ndcg@10", "map", "rp@10", "mrr"]
+
+    def load_dataset(self, dataset_id: str, split: DatasetSplit) -> RankingDataset:
+        """Load BEIR-style graded annotations and convert to a RankingDataset."""
+        hf_split = self.split_to_hf_split.get(split)
+        if hf_split is None:
+            supported = ", ".join(sorted(s.value for s in self.split_to_hf_split))
+            raise ValueError(
+                f"Split '{split.value}' not supported for {type(self).__name__}: "
+                f"only [{supported}] {'is' if len(self.split_to_hf_split) == 1 else 'are'} "
+                f"annotated for this dataset."
+            )
+
+        queries_ds = load_dataset(self.hf_name, "queries", split=hf_split)
+        corpus_ds = load_dataset(self.hf_name, "corpus", split="corpus")
+        qrels_ds = load_dataset(self.hf_name, "qrels", split=hf_split)
+        assert isinstance(queries_ds, Dataset)
+        assert isinstance(corpus_ds, Dataset)
+        assert isinstance(qrels_ds, Dataset)
+        queries_df = queries_ds.to_pandas()
+        corpus_df = corpus_ds.to_pandas()
+        qrels_df = qrels_ds.to_pandas()
+        assert isinstance(queries_df, pd.DataFrame)
+        assert isinstance(corpus_df, pd.DataFrame)
+        assert isinstance(qrels_df, pd.DataFrame)
+
+        # target_space is the corpus titles in corpus order; URIs map by row index.
+        target_space = corpus_df["title"].tolist()
+        uri_to_idx = {uri: i for i, uri in enumerate(corpus_df["_id"])}
+
+        qrels_df["target_idx"] = qrels_df["corpus-id"].map(uri_to_idx)
+        # Every qrel should resolve; if any don't, surface the issue rather than silently dropping.
+        unresolved = qrels_df["target_idx"].isna().sum()
+        assert unresolved == 0, (
+            f"{unresolved} qrel rows reference corpus-ids not present in the corpus config"
+        )
+        qrels_df["target_idx"] = qrels_df["target_idx"].astype(int)
+
+        id_to_query = dict(zip(queries_df["_id"], queries_df["text"], strict=True))
+        qrels_df["sentence"] = qrels_df["query-id"].map(id_to_query)
+
+        grouped = qrels_df.groupby("sentence")
+        filtered_queries: list[str] = []
+        filtered_indices: list[list[int]] = []
+        filtered_relevance: list[list[float]] = []
+        for sentence, group in grouped:
+            filtered_queries.append(str(sentence))
+            filtered_indices.append(group["target_idx"].tolist())
+            filtered_relevance.append([float(s) for s in group["score"].tolist()])
+
+        return RankingDataset(
+            query_texts=filtered_queries,
+            target_indices=filtered_indices,
+            target_space=target_space,
+            dataset_id=dataset_id,
+            target_relevance=filtered_relevance,
+        )