SOAPgemma_v1: Fine-tuned TxGemma for Medical SOAP Note Generation

SOAPgemma_v1 is a LoRA-fine-tuned version of Google's google/txgemma-2b-predict model, specifically adapted for generating medical SOAP (Subjective, Objective, Assessment, Plan) notes from patient-doctor dialogues.

This model was developed as a collaborative effort by:

• Hossam Shahin
• Khaled Al Raas
• Konstantina Voukelatou
• Shalinda Silva

Model Details

• Base Model: google/txgemma-2b-predict (a Gemma 2 architecture derivative)
• Fine-tuning Technique: Low-Rank Adaptation (LoRA) using PEFT library.
• Quantization: 4-bit quantization (bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16) was applied during fine-tuning for efficiency.
• Frameworks: PyTorch, Transformers, PEFT, TRL.
• Language: English

Model Description

SOAPGemma aims to alleviate the significant documentation burden faced by healthcare professionals. SOAP notes are a structured and essential method for recording patient encounters. This model takes a transcribed patient-doctor dialogue as input and generates a structured SOAP note, summarizing the key information into Subjective, Objective, Assessment, and Plan sections.

The primary goal is to improve clinical efficiency, reduce administrative workload, and allow clinicians to dedicate more time to direct patient care.

Intended Uses & Limitations

Intended Uses

• Automated SOAP Note Generation: To assist healthcare professionals by generating draft SOAP notes from transcribed patient-doctor dialogues.
• Clinical Documentation Support: To serve as a tool to streamline the documentation process in clinical settings.
• Research in Medical NLP: To be used as a baseline or for further research in automated clinical text generation.

Limitations and Potential Bias

• Synthetic Training Data: The model was fine-tuned on the OMI (Open Medical Insight) dataset, which consists of synthetic dialogues and SOAP summaries. While based on real scenarios (PubMed Central), performance on real-world, diverse clinical dialogues may vary.
• Not a Replacement for Clinical Judgment: This model is intended as an assistive tool. All generated notes must be reviewed, edited, and validated by a qualified healthcare professional before being finalized or used for clinical decision-making. The model does not possess clinical expertise.
• Potential for Factual Inaccuracies (Hallucinations): Like all large language models, SOAPGemma may occasionally generate information that is incorrect or not present in the input dialogue. Careful review is crucial.
• Bias in Training Data: The OMI dataset, although synthetic, may inherit biases from its source data (PubMed Central case reports and GPT-4 generation). The model might reflect these biases in its outputs.
• Specificity of Medical Conditions: The model's performance may vary depending on the complexity and rarity of medical conditions discussed in the dialogue.
• Out-of-Distribution Inputs: Performance may degrade if the input dialogues differ significantly in style, vocabulary, or structure from the training data.
• Language Limitation: Currently, the model is designed for English language dialogues only.

How to Use

This model is a PEFT LoRA adapter. To use it, you need to load the base model (google/txgemma-2b-predict) and then apply this adapter.

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# Define the base model ID and the adapter path (your Hugging Face model ID for SOAPgemma_v1)
base_model_id = "google/txgemma-2b-predict"
adapter_model_id = "your_username/SOAPgemma_v1" # Replace with your actual model ID on Hugging Face

# Configure 4-bit quantization
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load the base model with quantization
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    quantization_config=bnb_config,
    device_map="auto", # Automatically maps model to available devices (GPU/CPU)
    trust_remote_code=True # If required by the base model
)

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model_id)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token # Set pad token if not set

# Load the PEFT adapter
model = PeftModel.from_pretrained(base_model, adapter_model_id)
model.eval() # Set the model to evaluation mode

# Example dialogue
dialogue = """
Doctor: Hello, how can I help you today?
Patient: I've been having this persistent cough for about two weeks now. It's dry and sometimes I get a bit short of breath, especially when I climb stairs.
Doctor: Any fever, chills, or body aches?
Patient: No fever, no chills. Just this annoying cough and some fatigue.
Doctor: Okay, let's have a listen to your lungs. (Performs examination) Lungs sound clear. Have you had any recent colds or flu?
Patient: Not really, no.
Doctor: We should probably get a chest X-ray to be sure. In the meantime, try to rest and use a humidifier. I'll also prescribe a cough suppressant.
"""

# Format the input for the model (ensure this matches your fine-tuning format)
# The prompt structure used during fine-tuning was:
# "dialogue: {dialogue}<soap_start> soap_note:{soap_note} <eos>"
# For inference, we provide the dialogue and the start token for the soap_note.
input_text = f"dialogue: {dialogue}<soap_start> soap_note:"

# Tokenize the input
inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=1536).to(model.device)

# Generate the SOAP note
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,  # Adjust as needed
        num_beams=1,         # Using num_beams=1 for greedy decoding, can be increased for beam search
        temperature=0.7,     # Adjust for creativity vs. factuality
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id
    )

# Decode the output
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

# Extract the SOAP note part
# The generated text will include the input prompt, so we need to extract the part after "soap_note:"
soap_note_start_phrase = "soap_note:"
generated_soap_note = generated_text.split(soap_note_start_phrase)[-1].strip()

print("Dialogue:\n", dialogue)
print("\nGenerated SOAP Note:\n", generated_soap_note)