Add BHCToAVS model for patient-friendly summaries #730
Conversation
Introduces the BHCToAVS model, which converts clinical Brief Hospital Course (BHC) notes into After-Visit Summaries (AVS) using a fine-tuned Mistral 7B model with LoRA adapters. Adds model implementation, documentation, an example usage script, and unit tests.
There was a problem hiding this comment.
Pull request overview
This PR adds a new clinical summarization model, BHCToAVS, that converts Brief Hospital Course (BHC) notes into patient-friendly After-Visit Summaries (AVS) using a fine-tuned Mistral-7B LoRA adapter. The implementation integrates with PyHealth's model API and includes comprehensive documentation and examples.
- Implements a new text generation model wrapping a Hugging Face LoRA adapter
- Adds unit tests with graceful handling of model download failures
- Provides example usage demonstrating the model's clinical text summarization capabilities
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| pyhealth/models/bhc_to_avs.py | Core model implementation with predict() method for generating patient-friendly summaries |
| pyhealth/models/init.py | Added BHCToAVS to module exports |
| tests/core/test_bhc_to_avs.py | Unit test validating the predict method with error handling for gated models |
| docs/api/models/pyhealth.models.BHCToAVS.rst | Sphinx autodoc configuration for the new model |
| docs/api/models.rst | Updated model index to include BHCToAVS |
| examples/bhc_to_avs_example.py | Example script demonstrating model usage with synthetic clinical text |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
The CI has failed. |
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
Logiquo
added
status: need review
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
Enhanced the BHCToAVS model with improved docstrings, error handling for Hugging Face token requirements, and more robust pipeline initialization. Updated and expanded the test suite to include both unit tests with a mocked pipeline and an optional integration test for real model inference.
Introduces a __post_init__ method in BHCToAVS to call the BaseModel initializer, ensuring proper nn.Module setup. Also updates docstring formatting and attribute documentation for clarity.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 9 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| raise RuntimeError( | ||
| "Hugging Face token not found. This model requires access to a gated repository.\n\n" | ||
| "Set the HF_TOKEN environment variable or pass hf_token=... when initializing BHCToAVS.\n\n" | ||
| "Example:\n" | ||
| " export HF_TOKEN='hf_...'\n" | ||
| " model = BHCToAVS()\n" | ||
| ) |
There was a problem hiding this comment.
The error message spans lines 87-92 and provides good guidance, but the message could be more specific about where to obtain a HuggingFace token. Consider adding a link to the HuggingFace token generation page (https://huggingface.co/settings/tokens) to help users quickly resolve this issue.
| # Prompt used during fine-tuning | ||
| _PROMPT = ( | ||
| "Summarize for the patient what happened during the hospital stay based on this doctor's note:\n" | ||
| "{bhc}\n\n" | ||
| "Summary for the patient:\n" | ||
| ) |
There was a problem hiding this comment.
The comment on line 36 states "Prompt used during fine-tuning" but this prompt is actually used during inference (as seen on line 146). If this prompt was indeed used during fine-tuning and is also being reused during inference, the comment should clarify this. If it's only used during inference, the comment is misleading and should be corrected to "Prompt template used during inference" or similar.
| def _get_pipeline(self): | ||
| """Create and cache the text-generation pipeline.""" | ||
| if not hasattr(self, "_pipeline"): | ||
| # Resolve HuggingFace token | ||
| token = self._resolve_token() | ||
|
|
||
| # Throw RuntimeError if token is not found | ||
| if token is None: | ||
| raise RuntimeError( | ||
| "Hugging Face token not found. This model requires access to a gated repository.\n\n" | ||
| "Set the HF_TOKEN environment variable or pass hf_token=... when initializing BHCToAVS.\n\n" | ||
| "Example:\n" | ||
| " export HF_TOKEN='hf_...'\n" | ||
| " model = BHCToAVS()\n" | ||
| ) | ||
|
|
||
| # Load base model | ||
| base = AutoModelForCausalLM.from_pretrained( | ||
| self.base_model_id, | ||
| torch_dtype=torch.bfloat16, | ||
| device_map="auto", | ||
| token=token, | ||
| ) | ||
|
|
||
| # Load LoRA adapter | ||
| model = PeftModelForCausalLM.from_pretrained( | ||
| base, | ||
| self.adapter_model_id, | ||
| torch_dtype=torch.bfloat16, | ||
| token=token, | ||
| ) | ||
|
|
||
| tokenizer = AutoTokenizer.from_pretrained(self.base_model_id, token=token) | ||
|
|
||
| # Create HF pipeline | ||
| self._pipeline = pipeline( | ||
| "text-generation", | ||
| model=model, | ||
| tokenizer=tokenizer, | ||
| model_kwargs={"torch_dtype": torch.bfloat16}, | ||
| ) | ||
|
|
||
| return self._pipeline |
There was a problem hiding this comment.
The _get_pipeline method loads a large 7B parameter model without any explicit guidance on resource requirements or expected load time. Consider adding documentation (either in the class docstring or method docstring) about: (1) expected memory requirements (GPU/CPU), (2) approximate model loading time, and (3) recommended hardware specifications. This would help users understand the resource implications before attempting to use the model.
| def __post_init__(self): | ||
| # Ensure nn.Module (via BaseModel) is initialized | ||
| super().__init__() | ||
|
|
There was a problem hiding this comment.
The BHCToAVS class uses a dataclass decorator but also inherits from BaseModel (which is an nn.Module). The post_init method calls super().init(), but this pattern is inconsistent with how BaseModel is typically used in the codebase. Looking at SdohClassifier (another dataclass-based model in the same codebase), it does not override post_init. Additionally, BaseModel's init expects an optional dataset parameter, but this implementation calls it without any arguments. This could lead to incorrect initialization. Consider removing post_init entirely or ensuring it properly initializes the BaseModel with appropriate parameters.
| def __post_init__(self): | |
| # Ensure nn.Module (via BaseModel) is initialized | |
| super().__init__() |
| base_model_id: str = field(default="mistralai/Mistral-7B-Instruct-v0.3") | ||
| adapter_model_id: str = field(default="williach31/mistral-7b-bhc-to-avs-lora") | ||
| hf_token: str | None = None |
There was a problem hiding this comment.
The BHCToAVS model deviates from the standard PyHealth model API pattern. Most models in the codebase (e.g., LogisticRegression, AdaCare, GAMENet) accept a dataset parameter in their initialization to query information like tokens and schemas. This model does not accept or use a dataset parameter, which is inconsistent with the BaseModel interface it inherits from. Consider whether this model should follow the standard pattern or if the deviation is intentional for this use case. If intentional, this should be clearly documented.
| self.assertGreater(len(summary.strip()), 0) | ||
|
|
||
| # Output should be different from input | ||
| self.assertNotIn(bhc_text[:40], summary) |
There was a problem hiding this comment.
The test assertion on line 66 uses a weak check that only verifies the first 40 characters of the input are not present in the output. This could pass even if the model is simply copying most of the input text or behaving incorrectly. Consider adding more robust assertions that verify the output actually represents a simplified, patient-friendly summary (e.g., checking for absence of medical jargon, checking for specific expected transformations, or verifying the output differs meaningfully from the input).
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
|
Thank you for the review. I updated the tests by adding a separate optional integration test and modifying the unit test to patch in a dummy pipeline, so CI does not require model downloads or gated access. |
|
Typically a BaseModels should take a SampleDataset to examine a first few batch of data to determine how to construct a models. I don't think your model follows this structure? |
Logiquo
added
status: wait response
2 participants
Contributor: Charan Williams ([email protected])
Contribution Type: New Model
Description:
Added a new clinical summarization model, BHCToAVS, which converts Brief Hospital Course (BHC) notes into patient-friendly After-Visit Summaries (AVS). The model wraps a fine-tuned Mistral-7B LoRA adapter hosted on Hugging Face and integrates with the PyHealth model API. This contribution includes the full model implementation, unit tests, documentation, and an example usage script.
Files to Review:
pyhealth/models/bhc_to_avs.py— Main model implementationpyhealth/models/__init__.py— Added import for the new modeltests/core/test_bhc_to_avs.py— Unit test for the BHCToAVS modeldocs/api/models/pyhealth.models.bhc_to_avs.rst— Sphinx documentation filedocs/api/models.rst— Updated model index to include BHCToAVSexamples/bhc_to_avs_example.py— Example usage demonstrating model prediction