SDOH Coding Model (#538)

* add sentence level sdoh multi-label classification model and test

* revert test mask

* sdoh: llama download HF api key; response parse test; fix inf test

* doc

* just added more details to the docs

* readthedocs updates

---------

Co-authored-by: Paul Landes <landes@mailc.net>
Co-authored-by: John Wu <johnwu3@sunlab-serv-03.cs.illinois.edu>

Author: 3 people

docs/api/models.rst

@@ -31,3 +31,4 @@ We implement the following models for supporting multiple healthcare predictive
     models/pyhealth.models.TCN
     models/pyhealth.models.GAN
     models/pyhealth.models.VAE
+    models/pyhealth.models.SDOH

docs/api/models/pyhealth.models.SDOH

@@ -0,0 +1,10 @@
+pyhealth.models.TransformersModel
+===================================
+
+
+Model for classifying social determininants of health (SDoH)
+
+.. autoclass:: pyhealth.models.sdoh
+    :members:
+    :undoc-members:
+    :show-inheritance:

pixi.lock

@@ -25,3 +25,4 @@
 from .transformer import Transformer, TransformerLayer
 from .transformers_model import TransformersModel
 from .vae import VAE
+from .sdoh import SdohClassifier

pyhealth/models/__init__.py

@@ -0,0 +1,157 @@
+"""Social determinants of health (SDoH) classification.
+
+"""
+__author__ = 'Paul Landes'
+from typing import Dict, Any, Set, ClassVar
+from dataclasses import dataclass, field
+import re
+import torch
+import transformers
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.pipelines import Pipeline
+from peft import PeftModelForCausalLM
+from pyhealth.models.base_model import BaseModel
+
+
+# the prompt and role used to supervised-fine tune the model
+_PROMPT: str = """\
+Classify sentences for social determinants of health (SDOH).
+
+Definitions SDOHs are given with labels in back ticks:
+
+* `housing`: The status of a patient’s housing is a critical SDOH, known to affect the outcome of treatment.
+
+* `transportation`: This SDOH pertains to a patient’s inability to get to/from their healthcare visits.
+
+* `relationship`: Whether or not a patient is in a partnered relationship is an abundant SDOH in the clinical notes.
+
+* `parent`: This SDOH should be used for descriptions of a patient being a parent to at least one child who is a minor (under the age of 18 years old).
+
+* `employment`: This SDOH pertains to expressions of a patient’s employment status. A sentence should be annotated as an Employment Status SDOH if it expresses if the patient is employed (a paid job), unemployed, retired, or a current student.
+
+* `support`: This SDOH is a sentence describes a patient that is actively receiving care support, such as emotional, health, financial support.  This support comes from family and friends but not health care professionals.
+
+* `-`: If no SDOH is found.
+
+Classify sentences for social determinants of health (SDOH) as a list labels in three back ticks. The sentence can be a member of multiple classes so output the labels that are mostly likely to be present.
+
+### Sentence: {sent}
+### SDOH labels:"""
+
+
+@dataclass
+class SdohClassifier(BaseModel):
+    """This predicts sentence level social determinants of health (SDoH) as a
+    multi-label classification from clinical text.  The model was trained from
+    the MIMIC-III derived dataset from `Guevara et al. (2024)`_.
+
+    **Important**: The :obj:`api_key` needs to be populated if the ``Llama 3.1
+    8B Instruct`` (or the setting of :obj:`base_model_id`) has not yet been
+    downloaded.
+
+
+    Example::
+        >>> from pyhealth.models import SdohClassifier
+        >>> sdoh = SdohClassifier()
+        >>> sent = 'Pt is homeless and has no car and has no parents or support'
+        >>> print(sdoh.predict(sent))
+        >>> {'housing', 'transportation'}
+
+
+    Citation:
+
+      `Guevara et al. (2024)`_ Large language models to identify social determinants of
+    health in electronic health records
+
+    .. _Guevara et al. (2024): https://www.nature.com/articles/s41746-023-00970-0
+
+    """
+    _ROLE: ClassVar[str] = 'You are a social determinants of health (SDOH) classifier.'
+    _LABELS: ClassVar[str] = 'transportation housing relationship employment support parent'.split()
+
+    api_key: str = field(default=None)
+    """The API token that starts with ``tf_`` needed to download the Llama
+    model.
+
+    """
+    base_model_id: str = field(default='meta-llama/Llama-3.1-8B-Instruct')
+    """The base model ID, which probably should not be modified."""
+
+    adapter_model_id: str = field(default='plandes/sdoh-llama-3-1-8b')
+    """The LoRA adapter model ID, which probably should not be modified."""
+
+    def _parse_response(self, text: str) -> Set[str]:
+        """Parse the LLM response (also used in the unit test case).."""
+        res_regs = (re.compile(r'(?:.*?`([a-z,` ]{3,}`))', re.DOTALL),
+                    re.compile(r'.*?[`#-]([a-z, \t\n\r]{3,}?)[`-].*', re.DOTALL))
+        matched: str = ''
+        for pat in res_regs:
+            m: re.Match = pat.match(text)
+            if m is not None:
+                matched = m.group(1)
+                break
+        return set(filter(lambda s: matched.find(s) > -1, self._LABELS))
+
+    def _mod_ignore_check_type(self):
+        from transformers.pipelines.text_generation import TextGenerationPipeline
+
+        def noop(*args, **kwargs):
+            pass
+
+        TextGenerationPipeline.check_model_type = noop
+
+    def _get_pipeline(self) -> Pipeline:
+        """Create the text generation pipeline.  The output is parsed by
+        :meth:`_parse_response`."""
+        if not hasattr(self, '_pipeline'):
+            params: Dict[str, Any] = {}
+            if self.api_key is not None:
+                params['token'] = self.api_key
+            base_model = AutoModelForCausalLM.from_pretrained(
+                self.base_model_id, **params)
+            model = PeftModelForCausalLM.from_pretrained(
+                base_model, self.adapter_model_id, **params)
+            tokenizer = AutoTokenizer.from_pretrained(
+                self.base_model_id, **params)
+            # suppress bogus error logging message under transformers 4.53
+            # https://github.com/huggingface/transformers/issues/29395
+            self._mod_ignore_check_type()
+            # create a pipeline for inferencing
+            self._pipeline = transformers.pipeline(
+                'text-generation',
+                framework='pt',
+                model=model,
+                tokenizer=tokenizer,
+                model_kwargs={'torch_dtype': torch.bfloat16},
+                device_map='auto')
+        return self._pipeline
+
+    def predict(self, sent: str) -> Set[str]:
+        """Predict the SDoH labels of ``sent`` (see class docs).
+
+        :param sent: the sentence text used for prediction
+
+        :return: the SDoH labels predicted by the model
+
+        """
+        # prompt used by the chat template
+        messages = [
+            {'role': 'system', 'content': self._ROLE},
+            {'role': 'user', 'content': _PROMPT.format(sent=sent)}]
+
+        pipeline: Pipeline = self._get_pipeline()
+        # inference the LLM
+        outputs = pipeline(
+            messages,
+            max_new_tokens=512,
+            eos_token_id=[
+                pipeline.tokenizer.eos_token_id,
+                pipeline.tokenizer.convert_tokens_to_ids('<|eot_id|>'),
+            ],
+            pad_token_id=pipeline.tokenizer.eos_token_id,
+            do_sample=True,
+            temperature=0.01)
+
+        # the textual LLM output
+        output = outputs[0]['generated_text'][-1]['content']
+        return self._parse_response(output)

pyhealth/models/sdoh.py

@@ -25,6 +25,8 @@ dependencies = [
   "torch~=2.7.1",
   "torchvision",
   "transformers~=4.53.2",
+  "peft",
+  "accelerate",
   "rdkit",
   "scikit-learn~=1.7.0",
   "networkx",