tmp commit

examples/medical_transcriptions_classification.py

@@ -3,7 +3,7 @@
 
 from pyhealth.datasets import MedicalTranscriptionsDataset
 from pyhealth.datasets import get_dataloader
-from pyhealth.models import HuggingfaceAutoModel
+from pyhealth.models import TransformersModel
 from pyhealth.trainer import Trainer
 
 root = "/srv/local/data/zw12/raw_data/MedicalTranscriptions"
@@ -27,7 +27,7 @@
 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 = TransformersModel(
  model_name="emilyalsentzer/Bio_ClinicalBERT",
  dataset=sample_dataset,
  feature_keys=["transcription"],

pyhealth/models/huggingface.py → pyhealth/models/transformers_model.py

@@ -1,30 +1,26 @@
 from typing import List, Dict
 
 import torch
-from transformers import AutoModel, AutoTokenizer
 import torch.nn as nn
-from torchvision import models
+from transformers import AutoModel, AutoTokenizer
 
 from pyhealth.datasets import SampleDataset
 from pyhealth.models import BaseModel
 
 
-class HuggingfaceAutoModel(BaseModel):
- """AutoModel class for Huggingface models.
+class TransformersModel(BaseModel):
+ """Transformers class for Huggingface models.
  """
 
  def __init__(
  self,
- model_name: str,
  dataset: SampleDataset,
  feature_keys: List[str],
  label_key: str,
  mode: str,
- pretrained=False,
- num_layers=18,
- **kwargs,
+ model_name: str,
  ):
- super(HuggingfaceAutoModel, self).__init__(
+ super(TransformersModel, self).__init__(
  dataset=dataset,
  feature_keys=feature_keys,
  label_key=label_key,
@@ -60,7 +56,6 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
 
 if __name__ == "__main__":
  from pyhealth.datasets import MedicalTranscriptionsDataset, get_dataloader
- from torchvision import transforms
 
  base_dataset = MedicalTranscriptionsDataset(
  root="/srv/local/data/zw12/raw_data/MedicalTranscriptions"
@@ -70,14 +65,12 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
 
  train_loader = get_dataloader(sample_dataset, batch_size=16, shuffle=True)
 
- model = HuggingfaceAutoModel(
- model_name="emilyalsentzer/Bio_ClinicalBERT",
+ model = TransformersModel(
  dataset=sample_dataset,
- feature_keys=[
- "transcription",
- ],
+ feature_keys=["transcription"],
  label_key="label",
  mode="multiclass",
+ model_name="emilyalsentzer/Bio_ClinicalBERT",
  )
 
  # data batch
@@ -88,4 +81,4 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
  print(ret)
 
  # try loss backward
- ret["loss"].backward()
+ ret["loss"].backward()