wisent 0.5.12__py3-none-any.whl → 0.5.13__py3-none-any.whl

wisent/core/contrastive_pairs/lm_eval_pairs/lm_task_extractors/mutual.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from wisent.core.contrastive_pairs.core.pair import ContrastivePair
+from wisent.core.contrastive_pairs.core.response import NegativeResponse, PositiveResponse
+from wisent.core.contrastive_pairs.lm_eval_pairs.atoms import LMEvalBenchmarkExtractor
+from wisent.core.cli_logger import setup_logger, bind
+
+if TYPE_CHECKING:
+    from lm_eval.api.task import ConfigurableTask
+
+
+__all__ = ["MutualExtractor"]
+_LOG = setup_logger(__name__)
+
+
+class MutualExtractor(LMEvalBenchmarkExtractor):
+    """Extractor for the Mutual benchmark."""
+
+    def extract_contrastive_pairs(
+        self,
+        lm_eval_task_data: ConfigurableTask,
+        limit: int | None = None,
+    ) -> list[ContrastivePair]:
+        """
+        Build contrastive pairs from Mutual docs.
+
+        Mutual schema:
+            - article: str
+            - options: []
+            - answers: str
+            
+        Args:
+            lm_eval_task_data: lm-eval task instance for Mutual.
+            limit: Optional maximum number of pairs to produce.
+
+        Returns:
+            A list of ContrastivePair objects.
+        """
+        log = bind(_LOG, task=getattr(lm_eval_task_data, "NAME", "unknown"))
+
+        max_items = self._normalize_limit(limit)
+        docs = self.load_docs(lm_eval_task_data, max_items)
+
+        pairs: list[ContrastivePair] = []
+
+        log.info("Extracting contrastive pairs", extra={"doc_count": len(docs)})
+
+        for doc in docs:
+            pair = self._extract_pair_from_doc(doc)
+            if pair is not None:
+                pairs.append(pair)
+                if max_items is not None and len(pairs) >= max_items:
+                    break
+
+        if not pairs:
+            task_name = getattr(lm_eval_task_data, "NAME", type(lm_eval_task_data).__name__)
+            log.warning("No valid Mutual pairs extracted", extra={"task": task_name})
+
+        return pairs
+    
+    def _extract_pair_from_doc(self, doc: dict[str, Any]) -> ContrastivePair | None:
+        """
+        Convert a single Mutual doc into a ContrastivePair, if possible.
+        Returns None when required fields are missing or malformed.
+        """
+        log = bind(_LOG, doc_id=doc.get("id", "unknown"))
+
+        try:
+            article = str(doc.get("article", "")).strip()
+            options = doc.get("options", [])
+            answer = str(doc.get("answers", "")).strip()
+            answer_idx = int(ord(answer) - ord("A"))
+            
+
+            if not article or not options or not (0 <= answer_idx < len(options)):
+                log.debug(
+                    "Skipping doc due to missing/invalid fields",
+                    extra={"doc": doc},
+                )
+                return None
+            
+            correct = options[answer_idx]
+            incorrect = options[(answer_idx+1)%len(options)]
+        
+            formatted_question = (f"{article}\nA. {incorrect}\nB. {correct}")
+
+            metadata = {
+                "label": "mutual",
+            }
+
+            return self._build_pair(
+                question=formatted_question,
+                correct=correct,
+                incorrect=incorrect,
+                metadata=metadata,
+            )
+
+        except Exception as exc:  
+            log.error("Error extracting pair from doc", exc_info=exc, extra={"doc": doc})
+            return None
+
+    @staticmethod
+    def _build_pair(
+        question: str,
+        correct: str,
+        incorrect: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> ContrastivePair:
+        positive_response = PositiveResponse(model_response=correct)
+        negative_response = NegativeResponse(model_response=incorrect)
+        return ContrastivePair(prompt=question, positive_response=positive_response, negative_response=negative_response, label=metadata.get("label"))

wisent/core/contrastive_pairs/lm_eval_pairs/lm_task_extractors/openbookqa.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from wisent.core.contrastive_pairs.core.pair import ContrastivePair
+from wisent.core.contrastive_pairs.core.response import NegativeResponse, PositiveResponse
+from wisent.core.contrastive_pairs.lm_eval_pairs.atoms import LMEvalBenchmarkExtractor
+from wisent.core.cli_logger import setup_logger, bind
+
+if TYPE_CHECKING:
+    from lm_eval.api.task import ConfigurableTask
+
+
+__all__ = ["OpenBookQAExtractor"]
+_LOG = setup_logger(__name__)
+
+
+class OpenBookQAExtractor(LMEvalBenchmarkExtractor):
+    """Extractor for the OpenBookQA benchmark."""
+
+    def extract_contrastive_pairs(
+        self,
+        lm_eval_task_data: ConfigurableTask,
+        limit: int | None = None,
+    ) -> list[ContrastivePair]:
+        """
+        Build contrastive pairs from OpenBookQA docs.
+
+        OpenBookQA schema:
+            - question_stem: str
+            - choices: dict,
+            - choices["text"]: list with possible endings strings
+            - answerKey: str
+
+        Args:
+            lm_eval_task_data: lm-eval task instance for OpenBookQA.
+            limit: Optional maximum number of pairs to produce.
+
+        Returns:
+            A list of ContrastivePair objects.
+        """
+        log = bind(_LOG, task=getattr(lm_eval_task_data, "NAME", "unknown"))
+
+        max_items = self._normalize_limit(limit)
+        docs = self.load_docs(lm_eval_task_data, max_items)
+
+        pairs: list[ContrastivePair] = []
+
+        log.info("Extracting contrastive pairs", extra={"doc_count": len(docs)})
+
+        for doc in docs:
+            pair = self._extract_pair_from_doc(doc)
+            if pair is not None:
+                pairs.append(pair)
+                if max_items is not None and len(pairs) >= max_items:
+                    break
+
+        if not pairs:
+            task_name = getattr(lm_eval_task_data, "NAME", type(lm_eval_task_data).__name__)
+            log.warning("No valid OpenBookQA pairs extracted", extra={"task": task_name})
+
+        return pairs
+    
+    def _extract_pair_from_doc(self, doc: dict[str, Any]) -> ContrastivePair | None:
+        """
+        Convert a single OpenBookQA doc into a ContrastivePair, if possible.
+        Returns None when required fields are missing or malformed.
+        """
+        log = bind(_LOG, doc_id=doc.get("id", "unknown"))
+
+        try:
+            question_stem = str(doc.get("question_stem", "")).strip()
+            choices = doc.get("choices", {})
+            endings = choices["text"]
+            answer = str(doc.get("answerKey", "")).strip()
+            answer_idx = int(ord(answer) - ord('A'))
+
+            if not question_stem or not endings or not (0 <= answer_idx < len(endings)):
+                log.debug(
+                    "Skipping doc due to missing/invalid fields",
+                    extra={"doc": doc},
+                )
+                return None
+            
+            correct = endings[answer_idx]
+            incorrect = endings[(answer_idx+1)%len(endings)]
+
+            question = f"{question_stem}"
+            formatted_question = f"{question}\nA. {incorrect}\nB. {correct}"
+
+            metadata = {
+                "label": "openbookqa",
+            }
+
+            return self._build_pair(
+                question=formatted_question,
+                correct=correct,
+                incorrect=incorrect,
+                metadata=metadata,
+            )
+
+        except Exception as exc:  
+            log.error("Error extracting pair from doc", exc_info=exc, extra={"doc": doc})
+            return None
+
+    @staticmethod
+    def _build_pair(
+        question: str,
+        correct: str,
+        incorrect: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> ContrastivePair:
+        positive_response = PositiveResponse(model_response=correct)
+        negative_response = NegativeResponse(model_response=incorrect)
+        return ContrastivePair(prompt=question, positive_response=positive_response, negative_response=negative_response, label=metadata.get("label"))

wisent/core/contrastive_pairs/lm_eval_pairs/lm_task_extractors/pawsx.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from wisent.core.contrastive_pairs.core.pair import ContrastivePair
+from wisent.core.contrastive_pairs.core.response import NegativeResponse, PositiveResponse
+from wisent.core.contrastive_pairs.lm_eval_pairs.atoms import LMEvalBenchmarkExtractor
+from wisent.core.cli_logger import setup_logger, bind
+
+if TYPE_CHECKING:
+    from lm_eval.api.task import ConfigurableTask
+
+
+__all__ = ["PAWSXExtractor"]
+_LOG = setup_logger(__name__)
+
+
+class PAWSXExtractor(LMEvalBenchmarkExtractor):
+    """Extractor for the PAWS-X benchmark."""
+
+    def extract_contrastive_pairs(
+        self,
+        lm_eval_task_data: ConfigurableTask,
+        limit: int | None = None,
+    ) -> list[ContrastivePair]:
+        """
+        Build contrastive pairs from PAWS-X docs.
+
+        PAWS-X schema:
+            - sentence1: str
+            - sentence2: str
+            - label: 0 or 1
+            
+        Args:
+            lm_eval_task_data: lm-eval task instance for PAWS-X.
+            limit: Optional maximum number of pairs to produce.
+
+        Returns:
+            A list of ContrastivePair objects.
+        """
+        log = bind(_LOG, task=getattr(lm_eval_task_data, "NAME", "unknown"))
+
+        max_items = self._normalize_limit(limit)
+        docs = self.load_docs(lm_eval_task_data, max_items)
+
+        pairs: list[ContrastivePair] = []
+
+        log.info("Extracting contrastive pairs", extra={"doc_count": len(docs)})
+
+        for doc in docs:
+            pair = self._extract_pair_from_doc(doc)
+            if pair is not None:
+                pairs.append(pair)
+                if max_items is not None and len(pairs) >= max_items:
+                    break
+
+        if not pairs:
+            task_name = getattr(lm_eval_task_data, "NAME", type(lm_eval_task_data).__name__)
+            log.warning("No valid PAWS-X pairs extracted", extra={"task": task_name})
+
+        return pairs
+    
+    def _extract_pair_from_doc(self, doc: dict[str, Any]) -> ContrastivePair | None:
+        """
+        Convert a single PAWS-X doc into a ContrastivePair, if possible.
+        Returns None when required fields are missing or malformed.
+        """
+        log = bind(_LOG, doc_id=doc.get("id", "unknown"))
+
+        try:
+            sentence1 = str(doc.get("sentence1", "")).strip()
+            sentence2 = str(doc.get("sentence2", "")).strip()
+            label = doc.get("label")
+
+            if not sentence1 or not sentence2 or label not in {0, 1}:
+                log.debug(
+                    "Skipping doc due to missing/invalid fields",
+                    extra={"doc": doc},
+                )
+                return None
+            
+            formatted_question = f"Is sentence '{sentence1}' paraphrase of sentence '{sentence2}'?\nA. Yes\nB. No"
+
+            correct = "Yes" if label == 0 else "No"
+            incorrect = "No" if label == 0 else "Yes"
+
+            metadata = {
+                "label": "paws-x",
+            }
+
+            return self._build_pair(
+                question=formatted_question,
+                correct=correct,
+                incorrect=incorrect,
+                metadata=metadata,
+            )
+
+        except Exception as exc:  
+            log.error("Error extracting pair from doc", exc_info=exc, extra={"doc": doc})
+            return None
+
+    @staticmethod
+    def _build_pair(
+        question: str,
+        correct: str,
+        incorrect: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> ContrastivePair:
+        positive_response = PositiveResponse(model_response=correct)
+        negative_response = NegativeResponse(model_response=incorrect)
+        return ContrastivePair(prompt=question, positive_response=positive_response, negative_response=negative_response, label=metadata.get("label"))

wisent/core/contrastive_pairs/lm_eval_pairs/lm_task_extractors/piqa.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from wisent.core.contrastive_pairs.core.pair import ContrastivePair
+from wisent.core.contrastive_pairs.core.response import NegativeResponse, PositiveResponse
+from wisent.core.contrastive_pairs.lm_eval_pairs.atoms import LMEvalBenchmarkExtractor
+from wisent.core.cli_logger import setup_logger, bind
+
+if TYPE_CHECKING:
+    from lm_eval.api.task import ConfigurableTask
+
+
+__all__ = ["PIQAExtractor"]
+_LOG = setup_logger(__name__)
+
+
+class PIQAExtractor(LMEvalBenchmarkExtractor):
+    """Extractor for the PIQA benchmark."""
+
+    def extract_contrastive_pairs(
+        self,
+        lm_eval_task_data: ConfigurableTask,
+        limit: int | None = None,
+    ) -> list[ContrastivePair]:
+        """
+        Build contrastive pairs from PIQA docs.
+        
+        PIQA schema:
+            - goal: str 
+            - sol1, sol2: str
+            - label: 0 or 1
+
+        Args:
+            lm_eval_task_data: lm-eval task instance for PIQA.
+            limit: Optional maximum number of pairs to produce.
+
+        Returns:
+            A list of ContrastivePair objects.
+        """
+        log = bind(_LOG, task=getattr(lm_eval_task_data, "NAME", "unknown"))
+
+        max_items = self._normalize_limit(limit)
+        docs = self.load_docs(lm_eval_task_data, max_items)
+
+        pairs: list[ContrastivePair] = []
+
+        log.info("Extracting contrastive pairs", extra={"doc_count": len(docs)})
+
+        for doc in docs:
+            pair = self._extract_pair_from_doc(doc)
+            if pair is not None:
+                pairs.append(pair)
+                if max_items is not None and len(pairs) >= max_items:
+                    break
+
+        if not pairs:
+            task_name = getattr(lm_eval_task_data, "NAME", type(lm_eval_task_data).__name__)
+            log.warning("No valid PIQA pairs extracted", extra={"task": task_name})
+
+        return pairs
+    
+    def _extract_pair_from_doc(self, doc: dict[str, Any]) -> ContrastivePair | None:
+        """
+        Convert a single PIQA doc into a ContrastivePair, if possible.
+        Returns None when required fields are missing or malformed.
+        """
+        log = bind(_LOG, doc_id=doc.get("id", "unknown"))
+
+        try:
+            goal = str(doc.get("goal", "")).strip()
+            sol1 = str(doc.get("sol1", "")).strip()
+            sol2 = str(doc.get("sol2", "")).strip()
+            label = doc.get("label")
+
+            if not goal or not sol1 or not sol2 or label not in {0, 1}:
+                log.debug(
+                    "Skipping doc due to missing/invalid fields",
+                    extra={"doc": doc},
+                )
+                return None
+
+            question = f"Question: {goal}\nAnswer:"
+            formatted_question = f"{question}\nA. {sol1}\nB. {sol2}"
+
+            correct = sol1 if label == 0 else sol2
+            incorrect = sol2 if label == 0 else sol1
+
+            metadata = {
+                "label": "piqa",
+            }
+
+            return self._build_pair(
+                question=formatted_question,
+                correct=correct,
+                incorrect=incorrect,
+                metadata=metadata,
+            )
+
+        except Exception as exc:  
+            log.error("Error extracting pair from doc", exc_info=exc, extra={"doc": doc})
+            return None
+
+    @staticmethod
+    def _build_pair(
+        question: str,
+        correct: str,
+        incorrect: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> ContrastivePair:
+        positive_response = PositiveResponse(model_response=correct)
+        negative_response = NegativeResponse(model_response=incorrect)
+        return ContrastivePair(prompt=question, positive_response=positive_response, negative_response=negative_response, label=metadata.get("label"))

wisent/core/contrastive_pairs/lm_eval_pairs/lm_task_extractors/prost.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from wisent.core.contrastive_pairs.core.pair import ContrastivePair
+from wisent.core.contrastive_pairs.core.response import NegativeResponse, PositiveResponse
+from wisent.core.contrastive_pairs.lm_eval_pairs.atoms import LMEvalBenchmarkExtractor
+from wisent.core.cli_logger import setup_logger, bind
+
+if TYPE_CHECKING:
+    from lm_eval.api.task import ConfigurableTask
+
+
+__all__ = ["PROSTExtractor"]
+_LOG = setup_logger(__name__)
+
+
+class PROSTExtractor(LMEvalBenchmarkExtractor):
+    """Extractor for the PROST benchmark."""
+
+    def extract_contrastive_pairs(
+        self,
+        lm_eval_task_data: ConfigurableTask,
+        limit: int | None = None,
+    ) -> list[ContrastivePair]:
+        """
+        Build contrastive pairs from PROST docs.
+
+        PROST schema:
+            - context: str
+            - ex_question: str
+            - A, B, C, D: possible answers, str
+            - label: 0 or 1 or 2 or 3
+
+        Args:
+            lm_eval_task_data: lm-eval task instance for PROST.
+            limit: Optional maximum number of pairs to produce.
+
+        Returns:
+            A list of ContrastivePair objects.
+        """
+        log = bind(_LOG, task=getattr(lm_eval_task_data, "NAME", "unknown"))
+
+        max_items = self._normalize_limit(limit)
+        docs = self.load_docs(lm_eval_task_data, max_items)
+
+        pairs: list[ContrastivePair] = []
+
+        log.info("Extracting contrastive pairs", extra={"doc_count": len(docs)})
+
+        for doc in docs:
+            pair = self._extract_pair_from_doc(doc)
+            if pair is not None:
+                pairs.append(pair)
+                if max_items is not None and len(pairs) >= max_items:
+                    break
+
+        if not pairs:
+            task_name = getattr(lm_eval_task_data, "NAME", type(lm_eval_task_data).__name__)
+            log.warning("No valid PROST pairs extracted", extra={"task": task_name})
+
+        return pairs
+    
+    def _extract_pair_from_doc(self, doc: dict[str, Any]) -> ContrastivePair | None:
+        """
+        Convert a single PROST doc into a ContrastivePair, if possible.
+        Returns None when required fields are missing or malformed.
+        """
+        log = bind(_LOG, doc_id=doc.get("id", "unknown"))
+
+        try:
+            context= str(doc.get("context", "")).strip()
+            question = str(doc.get("ex_question", "")).strip()
+            answers = [str(doc.get("A", "")).strip(), str(doc.get("B", "")).strip(), str(doc.get("C", "")).strip(), str(doc.get("D", "")).strip()]
+            label = doc.get("label")
+
+            if not context or not question or not answers or label not in {0, 1, 2, 3}:
+                log.debug(
+                    "Skipping doc due to missing/invalid fields",
+                    extra={"doc": doc},
+                )
+                return None
+            
+            correct = answers[label]
+            incorrect = answers[(label+1)%len(answers)]
+
+            formatted_question = f"{context}\nQuestion: {question}\nAnswer:\nA. {incorrect}\nB. {correct}"
+
+            metadata = {
+                "label": "prost",
+            }
+
+            return self._build_pair(
+                question=formatted_question,
+                correct=correct,
+                incorrect=incorrect,
+                metadata=metadata,
+            )
+
+        except Exception as exc:  
+            log.error("Error extracting pair from doc", exc_info=exc, extra={"doc": doc})
+            return None
+
+    @staticmethod
+    def _build_pair(
+        question: str,
+        correct: str,
+        incorrect: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> ContrastivePair:
+        positive_response = PositiveResponse(model_response=correct)
+        negative_response = NegativeResponse(model_response=incorrect)
+        return ContrastivePair(prompt=question, positive_response=positive_response, negative_response=negative_response, label=metadata.get("label"))

wisent/core/contrastive_pairs/lm_eval_pairs/lm_task_extractors/pubmedqa.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from wisent.core.contrastive_pairs.core.pair import ContrastivePair
+from wisent.core.contrastive_pairs.core.response import NegativeResponse, PositiveResponse
+from wisent.core.contrastive_pairs.lm_eval_pairs.atoms import LMEvalBenchmarkExtractor
+from wisent.core.cli_logger import setup_logger, bind
+
+if TYPE_CHECKING:
+    from lm_eval.api.task import ConfigurableTask
+
+
+__all__ = ["PubMedQAExtractor"]
+_LOG = setup_logger(__name__)
+
+
+class PubMedQAExtractor(LMEvalBenchmarkExtractor):
+    """Extractor for the PubMedQA benchmark."""
+
+    def extract_contrastive_pairs(
+        self,
+        lm_eval_task_data: ConfigurableTask,
+        limit: int | None = None,
+    ) -> list[ContrastivePair]:
+        """
+        Build contrastive pairs from PubMedQA docs.
+
+        PubMedQA schema:
+            - CONTEXTS: list of strings to concatenate
+            - QUESTION: str
+            - final_decision: "yes" or "no"
+
+        Args:
+            lm_eval_task_data: lm-eval task instance for PubMedQA.
+            limit: Optional maximum number of pairs to produce.
+
+        Returns:
+            A list of ContrastivePair objects.
+        """
+        log = bind(_LOG, task=getattr(lm_eval_task_data, "NAME", "unknown"))
+
+        max_items = self._normalize_limit(limit)
+        docs = self.load_docs(lm_eval_task_data, max_items)
+
+        pairs: list[ContrastivePair] = []
+
+        log.info("Extracting contrastive pairs", extra={"doc_count": len(docs)})
+
+        for doc in docs:
+            pair = self._extract_pair_from_doc(doc)
+            if pair is not None:
+                pairs.append(pair)
+                if max_items is not None and len(pairs) >= max_items:
+                    break
+
+        if not pairs:
+            task_name = getattr(lm_eval_task_data, "NAME", type(lm_eval_task_data).__name__)
+            log.warning("No valid PubMedQA pairs extracted", extra={"task": task_name})
+
+        return pairs
+    
+    def _extract_pair_from_doc(self, doc: dict[str, Any]) -> ContrastivePair | None:
+        """
+        Convert a single PubMedQA doc into a ContrastivePair, if possible.
+        Returns None when required fields are missing or malformed.
+        """
+        log = bind(_LOG, doc_id=doc.get("id", "unknown"))
+
+        try:
+            contexts = doc.get("CONTEXTS", [])
+            question = str(doc.get("QUESTION", "")).strip()
+            final_decision = str(doc.get("final_decision", "")).strip()
+
+            if not contexts or not question or not final_decision:
+                log.debug(
+                    "Skipping doc due to missing/invalid fields",
+                    extra={"doc": doc},
+                )
+                return None
+        
+            formatted_context = " ".join(s.strip() for s in contexts if isinstance(s, str) and s.strip())
+            formatted_question = f"Abstract: {formatted_context}\nQuestion: {question}\nAnswer:\nA. yes\nB. no"
+
+            correct = final_decision
+            incorrect = "yes" if correct == "no" else "no"
+
+            metadata = {
+                "label": "pubmedqa",
+            }
+
+            return self._build_pair(
+                question=formatted_question,
+                correct=correct,
+                incorrect=incorrect,
+                metadata=metadata,
+            )
+
+        except Exception as exc:  
+            log.error("Error extracting pair from doc", exc_info=exc, extra={"doc": doc})
+            return None
+
+    @staticmethod
+    def _build_pair(
+        question: str,
+        correct: str,
+        incorrect: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> ContrastivePair:
+        positive_response = PositiveResponse(model_response=correct)
+        negative_response = NegativeResponse(model_response=incorrect)
+        return ContrastivePair(prompt=question, positive_response=positive_response, negative_response=negative_response, label=metadata.get("label"))