nan bug fixed for MedicalTranscriptionsDataset

examples/medical_transcriptions_classification.py

@@ -0,0 +1,49 @@
+import numpy as np
+import torch
+
+from pyhealth.datasets import MedicalTranscriptionsDataset
+from pyhealth.datasets import get_dataloader
+from pyhealth.models import HuggingfaceAutoModel
+from pyhealth.trainer import Trainer
+
+root = "/srv/local/data/zw12/raw_data/MedicalTranscriptions"
+base_dataset = MedicalTranscriptionsDataset(root)
+
+sample_dataset = base_dataset.set_task()
+
+ratios = [0.7, 0.1, 0.2]
+index = np.arange(len(sample_dataset))
+np.random.shuffle(index)
+s1 = int(len(sample_dataset) * ratios[0])
+s2 = int(len(sample_dataset) * (ratios[0] + ratios[1]))
+train_index = index[: s1]
+val_index = index[s1: s2]
+test_index = index[s2:]
+train_dataset = torch.utils.data.Subset(sample_dataset, train_index)
+val_dataset = torch.utils.data.Subset(sample_dataset, val_index)
+test_dataset = torch.utils.data.Subset(sample_dataset, test_index)
+
+train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
+val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
+test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
+
+model = HuggingfaceAutoModel(
+ model_name="emilyalsentzer/Bio_ClinicalBERT",
+ dataset=sample_dataset,
+ feature_keys=["transcription"],
+ label_key="label",
+ mode="multiclass",
+)
+
+trainer = Trainer(model=model)
+
+print(trainer.evaluate(test_dataloader))
+
+trainer.train(
+ train_dataloader=train_dataloader,
+ val_dataloader=val_dataloader,
+ epochs=1,
+ monitor="accuracy"
+)
+
+print(trainer.evaluate(test_dataloader))

pyhealth/models/huggingface.py

@@ -45,6 +45,7 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
  x = self.tokenizer(
  x, return_tensors="pt", padding=True, truncation=True, max_length=256
  )
+ x = x.to(self.device)
  embeddings = self.model(**x).pooler_output
  logits = self.fc(embeddings)
  y_true = self.prepare_labels(kwargs[self.label_key], self.label_tokenizer)

pyhealth/tasks/medical_transcriptions_classification.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field
 from typing import Dict
+import pandas as pd
 
 from pyhealth.tasks import TaskTemplate
 
@@ -11,7 +12,9 @@ class MedicalTranscriptionsClassification(TaskTemplate):
  output_schema: Dict[str, str] = field(default_factory=lambda: {"label": "label"})
 
  def __call__(self, patient):
- if patient["transcription"] is None or patient["medical_specialty"] is None:
+ if patient["transcription"] is None or pd.isna(patient["transcription"]):
+ return []
+ if patient["medical_specialty"] is None or pd.isna(patient["medical_specialty"]):
  return []
  sample = {
  "transcription": patient["transcription"],