nox changes

src/eva/language/data/datasets/classification/pubmedqa.py

@@ -1,8 +1,8 @@
 """PubMedQA dataset class."""
 
 from typing import Any, Dict, List
-
-from datasets import load_dataset
+import torch
+from datasets import load_dataset, Dataset
 from typing_extensions import override
 
 from eva.language.data.datasets.classification import base
@@ -25,10 +25,18 @@ def __init__(
         """
         super().__init__()
         self._split = split
-        self.dataset = load_dataset(
-            "bigbio/pubmed_qa", name="pubmed_qa_labeled_fold0_source", split=split
+        raw_dataset = load_dataset(
+            "bigbio/pubmed_qa",
+            name="pubmed_qa_labeled_fold0_source",
+            split=split,
+            streaming=False
         )
 
+        if not isinstance(raw_dataset, Dataset):
+            raise TypeError(f"Expected a `Dataset`, but got {type(raw_dataset)}")
+
+        self.dataset: Dataset = raw_dataset
+
     @property
     @override
     def classes(self) -> List[str]:
@@ -41,12 +49,15 @@ def class_to_idx(self) -> Dict[str, int]:
 
     @override
     def load_text(self, index: int) -> str:
-        sample = self.dataset[index]
+        sample = dict(self.dataset[index])
         return f"Question: {sample['QUESTION']}\nContext: {sample['CONTEXTS']}"
 
     @override
-    def load_target(self, index: int) -> int:
-        return self.class_to_idx[self.dataset[index]["final_decision"]]
+    def load_target(self, index: int) -> torch.Tensor:
+        return torch.tensor(
+            self.class_to_idx[self.dataset[index]["final_decision"]],
+            dtype=torch.long
+        )
 
     @override
     def load_metadata(self, index: int) -> Dict[str, Any]: