themis-eval 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

themis/datasets/med_qa.py

@@ -0,0 +1,179 @@
+"""Helpers for working with the bigbio/med_qa dataset."""
+
+from __future__ import annotations
+
+import json
+import string
+from pathlib import Path
+from typing import Any, Iterable, Iterator, List, Sequence
+
+from pydantic import BaseModel, Field, ValidationInfo, field_validator
+
+_DATASET_NAME = "bigbio/med_qa"
+_CHOICE_LABELS = tuple(string.ascii_uppercase)
+
+
+class MedQaSample(BaseModel):
+    unique_id: str
+    question: str
+    choices: list[str]
+    answer: str
+    subject: str = Field(default="medicine")
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    choice_labels: list[str] = Field(default_factory=list)
+
+    @field_validator("choices", mode="before")
+    @classmethod
+    def _ensure_choices(cls, value: Any) -> list[str]:
+        if value is None:
+            return []
+        if isinstance(value, dict):
+            return [str(v) for _, v in sorted(value.items())]
+        if isinstance(value, (list, tuple)):
+            return [str(item) for item in value]
+        raise TypeError("choices must be a sequence or mapping")
+
+    @field_validator("choice_labels", mode="before")
+    @classmethod
+    def _build_choice_labels(cls, value: Any, info: ValidationInfo) -> list[str]:
+        if value:
+            return [str(item) for item in value]
+        choices = info.data.get("choices") if hasattr(info, "data") else None
+        total = len(choices) if isinstance(choices, list) else 0
+        return [*_CHOICE_LABELS[:total]]
+
+    def to_generation_example(self) -> dict[str, Any]:
+        effective_labels = (
+            list(self.choice_labels)
+            if self.choice_labels
+            else list(_CHOICE_LABELS[: len(self.choices)])
+        )
+        return {
+            "unique_id": self.unique_id,
+            "question": self.question,
+            "choices": list(self.choices),
+            "choice_labels": effective_labels,
+            "answer": self.answer,
+            "subject": self.subject,
+            "metadata": dict(self.metadata),
+        }
+
+
+def load_med_qa(
+    *,
+    split: str = "test",
+    limit: int | None = None,
+    source: str = "huggingface",
+    data_dir: str | Path | None = None,
+    subset: str = "med_qa_en_bigbio_qa",
+) -> List[MedQaSample]:
+    """Load MedQA samples from Hugging Face or a local directory."""
+
+    if source not in {"huggingface", "local"}:
+        raise ValueError(
+            f"Unsupported source '{source}'. Expected one of: 'huggingface', 'local'."
+        )
+
+    if source == "huggingface":
+        rows = _load_from_huggingface(split=split, subset=subset)
+    else:
+        if data_dir is None:
+            raise ValueError(
+                "data_dir must be provided when source='local'. "
+                "Pass dataset.data_dir in configs or --data-dir on the CLI."
+            )
+        rows = _load_from_local(Path(data_dir))
+
+    samples: list[MedQaSample] = []
+    for index, row in enumerate(rows, start=1):
+        sample = _row_to_sample(row, index=index)
+        samples.append(sample)
+        if limit is not None and len(samples) >= limit:
+            break
+    return samples
+
+
+def _row_to_sample(row: dict[str, Any], *, index: int) -> MedQaSample:
+    unique_id = (
+        row.get("id")
+        or row.get("unique_id")
+        or f"med-qa-{index:05d}"
+    )
+    question = row.get("question") or ""
+    
+    # BigBio MedQA format:
+    # choices: [{'key': 'A', 'text': '...'}, {'key': 'B', 'text': '...'}]
+    # answer: 'A' (or similar)
+    
+    choices_data = row.get("choices") or []
+    choices = []
+    choice_labels = []
+    
+    if isinstance(choices_data, list):
+        # Sort by key to ensure order
+        try:
+            sorted_choices = sorted(choices_data, key=lambda x: x.get("key", ""))
+            for c in sorted_choices:
+                choices.append(str(c.get("text", "")))
+                choice_labels.append(str(c.get("key", "")))
+        except (TypeError, AttributeError):
+            # Fallback if structure is different
+            choices = [str(c) for c in choices_data]
+            
+    answer = ""
+    answer_data = row.get("answer")
+    if isinstance(answer_data, list) and answer_data:
+        answer = str(answer_data[0]) # Usually a list with one element
+    elif isinstance(answer_data, str):
+        answer = answer_data
+
+    metadata_keys = {
+        "question", "choices", "answer", "id"
+    }
+    metadata = {key: value for key, value in row.items() if key not in metadata_keys}
+    
+    return MedQaSample(
+        unique_id=str(unique_id),
+        question=str(question),
+        choices=choices,
+        choice_labels=choice_labels,
+        answer=answer,
+        metadata=metadata,
+    )
+
+
+def _load_from_huggingface(*, split: str, subset: str) -> Iterable[dict[str, Any]]:
+    try:
+        from datasets import load_dataset
+    except ImportError as exc:  # pragma: no cover - optional dependency
+        raise RuntimeError(
+            "datasets is required to load MedQA from Hugging Face. Install it via `uv pip install '.[hf]'`."
+        ) from exc
+
+    dataset = load_dataset(_DATASET_NAME, subset, split=split)
+    for row in dataset:
+        yield dict(row)
+
+
+def _load_from_local(root: Path) -> Iterator[dict[str, Any]]:
+    if not root.exists():
+        raise FileNotFoundError(f"Local dataset directory not found: {root}")
+
+    for path in root.rglob("*"):
+        if path.suffix.lower() == ".json":
+            with path.open("r", encoding="utf-8") as handle:
+                row = json.load(handle)
+                row.setdefault("id", path.stem)
+                yield row
+        elif path.suffix.lower() in {".jsonl", ".ndjson"}:
+            with path.open("r", encoding="utf-8") as handle:
+                for line_num, line in enumerate(handle, start=1):
+                    line = line.strip()
+                    if not line:
+                        continue
+                    row = json.loads(line)
+                    row.setdefault("id", f"{path.stem}-{line_num}")
+                    yield row
+
+
+__all__ = ["MedQaSample", "load_med_qa"]

themis/datasets/medmcqa.py

@@ -0,0 +1,169 @@
+"""Helpers for working with the openlifescienceai/medmcqa dataset."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Iterable, Iterator, List, Sequence
+
+from pydantic import BaseModel, Field, field_validator
+
+_DATASET_NAME = "openlifescienceai/medmcqa"
+_CHOICE_LABELS = ["A", "B", "C", "D"]
+
+
+class MedMcqaSample(BaseModel):
+    unique_id: str
+    question: str
+    choices: list[str]
+    answer: str
+    subject: str = Field(default="unknown")
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    choice_labels: list[str] = Field(default_factory=list)
+
+    @field_validator("choices", mode="before")
+    @classmethod
+    def _ensure_choices(cls, value: Any) -> list[str]:
+        if value is None:
+            return []
+        if isinstance(value, dict):
+            return [str(v) for _, v in sorted(value.items())]
+        if isinstance(value, (list, tuple)):
+            return [str(item) for item in value]
+        raise TypeError("choices must be a sequence or mapping")
+
+    def to_generation_example(self) -> dict[str, Any]:
+        effective_labels = (
+            list(self.choice_labels)
+            if self.choice_labels
+            else list(_CHOICE_LABELS[: len(self.choices)])
+        )
+        return {
+            "unique_id": self.unique_id,
+            "question": self.question,
+            "choices": list(self.choices),
+            "choice_labels": effective_labels,
+            "answer": self.answer,
+            "subject": self.subject,
+            "metadata": dict(self.metadata),
+        }
+
+
+def load_medmcqa(
+    *,
+    split: str = "test",
+    limit: int | None = None,
+    source: str = "huggingface",
+    data_dir: str | Path | None = None,
+    subset: str | None = None,
+) -> List[MedMcqaSample]:
+    """Load MedMCQA samples from Hugging Face or a local directory."""
+
+    if source not in {"huggingface", "local"}:
+        raise ValueError(
+            f"Unsupported source '{source}'. Expected one of: 'huggingface', 'local'."
+        )
+
+    if source == "huggingface":
+        rows = _load_from_huggingface(split=split, subset=subset)
+    else:
+        if data_dir is None:
+            raise ValueError(
+                "data_dir must be provided when source='local'. "
+                "Pass dataset.data_dir in configs or --data-dir on the CLI."
+            )
+        rows = _load_from_local(Path(data_dir))
+
+    samples: list[MedMcqaSample] = []
+    for index, row in enumerate(rows, start=1):
+        sample = _row_to_sample(row, index=index)
+        samples.append(sample)
+        if limit is not None and len(samples) >= limit:
+            break
+    return samples
+
+
+def _row_to_sample(row: dict[str, Any], *, index: int) -> MedMcqaSample:
+    unique_id = (
+        row.get("id")
+        or row.get("unique_id")
+        or f"medmcqa-{index:05d}"
+    )
+    question = row.get("question") or ""
+    
+    # MedMCQA has 'opa', 'opb', 'opc', 'opd'
+    choices = [
+        str(row.get("opa") or ""),
+        str(row.get("opb") or ""),
+        str(row.get("opc") or ""),
+        str(row.get("opd") or ""),
+    ]
+    
+    # Answer is an integer 0-3 (sometimes 1-4 depending on version, but usually 0-3 in HF)
+    # Let's check. HF dataset viewer says 'cop': 1 (meaning option B if 1-based, or B if 0-based?)
+    # Usually it's 0-3 or 1-4.
+    # Checking dataset info: "cop: Correct option (1-4)"
+    # So we need to map 1->A, 2->B, 3->C, 4->D
+    
+    cop = row.get("cop")
+    answer = ""
+    if cop is not None:
+        try:
+            cop_int = int(cop)
+            if 1 <= cop_int <= 4:
+                answer = _CHOICE_LABELS[cop_int - 1]
+        except (ValueError, TypeError):
+            pass
+
+    subject = row.get("subject_name") or "medicine"
+
+    metadata_keys = {
+        "question", "opa", "opb", "opc", "opd", "cop", "subject_name", "id"
+    }
+    metadata = {key: value for key, value in row.items() if key not in metadata_keys}
+    
+    return MedMcqaSample(
+        unique_id=str(unique_id),
+        question=str(question),
+        choices=choices,
+        answer=answer,
+        subject=str(subject),
+        metadata=metadata,
+    )
+
+
+def _load_from_huggingface(*, split: str, subset: str | None) -> Iterable[dict[str, Any]]:
+    try:
+        from datasets import load_dataset
+    except ImportError as exc:  # pragma: no cover - optional dependency
+        raise RuntimeError(
+            "datasets is required to load MedMCQA from Hugging Face. Install it via `uv pip install '.[hf]'`."
+        ) from exc
+
+    dataset = load_dataset(_DATASET_NAME, subset, split=split)
+    for row in dataset:
+        yield dict(row)
+
+
+def _load_from_local(root: Path) -> Iterator[dict[str, Any]]:
+    if not root.exists():
+        raise FileNotFoundError(f"Local dataset directory not found: {root}")
+
+    for path in root.rglob("*"):
+        if path.suffix.lower() == ".json":
+            with path.open("r", encoding="utf-8") as handle:
+                row = json.load(handle)
+                row.setdefault("id", path.stem)
+                yield row
+        elif path.suffix.lower() in {".jsonl", ".ndjson"}:
+            with path.open("r", encoding="utf-8") as handle:
+                for line_num, line in enumerate(handle, start=1):
+                    line = line.strip()
+                    if not line:
+                        continue
+                    row = json.loads(line)
+                    row.setdefault("id", f"{path.stem}-{line_num}")
+                    yield row
+
+
+__all__ = ["MedMcqaSample", "load_medmcqa"]

themis/datasets/mmlu_pro.py

@@ -0,0 +1,262 @@
+"""Helpers for working with the TIGER-Lab/MMLU-Pro dataset."""
+
+from __future__ import annotations
+
+import json
+import string
+from pathlib import Path
+from typing import Any, Iterable, Iterator, List, Sequence
+
+from pydantic import BaseModel, Field, ValidationInfo, field_validator
+
+_DATASET_NAME = "TIGER-Lab/MMLU-Pro"
+_CHOICE_LABELS = tuple(string.ascii_uppercase)
+
+
+def _index_to_label(index: int) -> str:
+    if 0 <= index < len(_CHOICE_LABELS):
+        return _CHOICE_LABELS[index]
+    return str(index)
+
+
+def _normalize_answer(value: Any, *, total_choices: int | None = None) -> str:
+    if isinstance(value, int):
+        return _index_to_label(value)
+    if isinstance(value, float):
+        as_int = int(value)
+        if as_int == value:
+            return _index_to_label(as_int)
+    text = str(value or "").strip()
+    if not text:
+        return ""
+    lowered = text.lower()
+    if lowered.startswith("option "):
+        text = text.split(" ", 1)[-1]
+    if lowered.startswith("choice "):
+        text = text.split(" ", 1)[-1]
+    if text.isdigit():
+        index = int(text)
+        if total_choices is None or 0 <= index < total_choices:
+            return _index_to_label(index)
+    text = text.strip().rstrip(".")
+    if len(text) == 1 and text.isalpha():
+        return text.upper()
+    if total_choices is not None:
+        mapping = {str(idx): _index_to_label(idx) for idx in range(total_choices)}
+        normalized = mapping.get(text)
+        if normalized:
+            return normalized
+    return text
+
+
+class MmluProSample(BaseModel):
+    unique_id: str
+    question: str
+    choices: list[str]
+    answer: str
+    subject: str = Field(default="unknown")
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    choice_labels: list[str] = Field(default_factory=list)
+
+    @field_validator("choices", mode="before")
+    @classmethod
+    def _ensure_choices(cls, value: Any) -> list[str]:
+        if value is None:
+            return []
+        if isinstance(value, dict):
+            return [str(v) for _, v in sorted(value.items())]
+        if isinstance(value, (list, tuple)):
+            return [str(item) for item in value]
+        raise TypeError("choices must be a sequence or mapping")
+
+    @field_validator("choice_labels", mode="before")
+    @classmethod
+    def _build_choice_labels(cls, value: Any, info: ValidationInfo) -> list[str]:
+        if value:
+            return [str(item) for item in value]
+        choices = info.data.get("choices") if hasattr(info, "data") else None
+        total = len(choices) if isinstance(choices, list) else 0
+        return [*_CHOICE_LABELS[:total]]
+
+    @field_validator("answer", mode="before")
+    @classmethod
+    def _normalize_answer_field(cls, value: Any, info: ValidationInfo) -> str:
+        choices = info.data.get("choices") if hasattr(info, "data") else None
+        total = len(choices) if isinstance(choices, list) else None
+        return _normalize_answer(value, total_choices=total)
+
+    def to_generation_example(self) -> dict[str, Any]:
+        effective_labels = (
+            list(self.choice_labels)
+            if self.choice_labels
+            else list(_CHOICE_LABELS[: len(self.choices)])
+        )
+        return {
+            "unique_id": self.unique_id,
+            "question": self.question,
+            "choices": list(self.choices),
+            "choice_labels": effective_labels,
+            "answer": self.answer,
+            "subject": self.subject,
+            "metadata": dict(self.metadata),
+        }
+
+
+def load_mmlu_pro(
+    *,
+    split: str = "test",
+    limit: int | None = None,
+    source: str = "huggingface",
+    data_dir: str | Path | None = None,
+    subjects: Sequence[str] | None = None,
+) -> List[MmluProSample]:
+    """Load MMLU-Pro samples from Hugging Face or a local directory."""
+
+    if source not in {"huggingface", "local"}:
+        raise ValueError(
+            f"Unsupported source '{source}'. Expected one of: 'huggingface', 'local'."
+        )
+
+    if source == "huggingface":
+        rows = _load_from_huggingface(split=split)
+    else:
+        if data_dir is None:
+            raise ValueError(
+                "data_dir must be provided when source='local'. "
+                "Pass dataset.data_dir in configs or --data-dir on the CLI."
+            )
+        rows = _load_from_local(Path(data_dir))
+
+    samples: list[MmluProSample] = []
+    selected_subjects = {s.lower() for s in subjects} if subjects else None
+    for index, row in enumerate(rows, start=1):
+        subject = _extract_subject(row)
+        if selected_subjects and subject.lower() not in selected_subjects:
+            continue
+        sample = _row_to_sample(row, index=index, subject=subject)
+        samples.append(sample)
+        if limit is not None and len(samples) >= limit:
+            break
+    return samples
+
+
+def _extract_subject(row: dict[str, Any]) -> str:
+    for key in ("subject", "category", "group", "topic", "field"):
+        value = row.get(key)
+        if value:
+            return str(value)
+    return "unknown"
+
+
+def _row_to_sample(row: dict[str, Any], *, index: int, subject: str) -> MmluProSample:
+    unique_id = (
+        row.get("id")
+        or row.get("question_id")
+        or row.get("unique_id")
+        or f"mmlu-pro-{index:05d}"
+    )
+    question = row.get("question") or row.get("Question") or row.get("prompt") or ""
+    choices = _extract_choices(row)
+    answer = (
+        row.get("answer")
+        or row.get("Answer")
+        or row.get("correct_answer")
+        or row.get("correct")
+        or ""
+    )
+    metadata_keys = {
+        "question",
+        "Question",
+        "prompt",
+        "choices",
+        "options",
+        "answer",
+        "Answer",
+        "correct_answer",
+        "correct",
+    }
+    metadata = {key: value for key, value in row.items() if key not in metadata_keys}
+    sample = MmluProSample.model_validate(
+        {
+            "unique_id": str(unique_id),
+            "question": str(question),
+            "choices": choices,
+            "answer": answer,
+            "subject": str(subject),
+            "metadata": metadata,
+        }
+    )
+    return sample
+
+
+def _extract_choices(row: dict[str, Any]) -> list[str]:
+    candidates = row.get("choices") or row.get("options") or row.get("Choices")
+    if isinstance(candidates, dict):
+        return [str(value) for _, value in sorted(candidates.items())]
+    if isinstance(candidates, (list, tuple)):
+        return [str(item) for item in candidates]
+
+    choice_keys = [
+        "A",
+        "B",
+        "C",
+        "D",
+        "E",
+        "F",
+        "choice_a",
+        "choice_b",
+        "choice_c",
+        "choice_d",
+        "choice_e",
+        "choice_f",
+        "option_a",
+        "option_b",
+        "option_c",
+        "option_d",
+        "option_e",
+        "option_f",
+    ]
+    collected: list[str] = []
+    for key in choice_keys:
+        if key in row:
+            collected.append(str(row[key]))
+    if collected:
+        return collected
+    return []
+
+
+def _load_from_huggingface(*, split: str) -> Iterable[dict[str, Any]]:
+    try:
+        from datasets import load_dataset
+    except ImportError as exc:  # pragma: no cover - optional dependency
+        raise RuntimeError(
+            "datasets is required to load MMLU-Pro from Hugging Face. Install it via `uv pip install '.[hf]'`."
+        ) from exc
+
+    dataset = load_dataset(_DATASET_NAME, split=split)
+    for row in dataset:
+        yield dict(row)
+
+
+def _load_from_local(root: Path) -> Iterator[dict[str, Any]]:
+    if not root.exists():
+        raise FileNotFoundError(f"Local dataset directory not found: {root}")
+
+    for path in root.rglob("*"):
+        if path.suffix.lower() == ".json":
+            with path.open("r", encoding="utf-8") as handle:
+                row = json.load(handle)
+                row.setdefault("id", path.stem)
+                yield row
+        elif path.suffix.lower() in {".jsonl", ".ndjson"}:
+            with path.open("r", encoding="utf-8") as handle:
+                for line_num, line in enumerate(handle, start=1):
+                    line = line.strip()
+                    if not line:
+                        continue
+                    row = json.loads(line)
+                    row.setdefault("id", f"{path.stem}-{line_num}")
+                    yield row
+
+
+__all__ = ["MmluProSample", "load_mmlu_pro"]