File size: 7,583 Bytes
b7f3196 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
"""
Inference module for Healthcare Reason Classification
This module provides inference for the reason classification system,
separate from the medical/insurance classifier.
"""
from ..head import ClassifierHead
from datetime import datetime
import os
import pprint
import torch
from sentence_transformers import SentenceTransformer
# Reason-specific configuration
REASON_CATEGORIES = {
0: "ROUTINE_CARE",
1: "PAIN_CONDITIONS",
2: "INJURIES",
3: "SKIN_CONDITIONS",
4: "STRUCTURAL_ISSUES",
5: "PROCEDURES"
}
REASON_CHECKPOINT_PATH = "classifier/reason_checkpoints"
DATETIME_FORMAT = "%Y%m%d_%H%M%S"
MODEL_NAME = "sentence-transformers/embeddinggemma-300m-medical"
def get_device():
"""Get the best available device for inference."""
if torch.backends.mps.is_available():
return torch.device("mps")
elif torch.cuda.is_available():
return torch.device("cuda")
else:
return torch.device("cpu")
DEVICE = get_device()
def get_reason_models():
"""Get the embedding model and classifier head for reason inference."""
# Load embedding model
embedding_model = SentenceTransformer(
MODEL_NAME,
prompts={
'classification': 'task: healthcare reason classification | query: ',
'retrieval (query)': 'task: search result | query: ',
'retrieval (document)': 'title: {title | "none"} | text: ',
},
default_prompt_name='classification',
)
# Load classifier head (for 6 reason categories)
classifier_head = ClassifierHead(len(REASON_CATEGORIES))
return embedding_model.to(DEVICE), classifier_head.to(DEVICE)
def predict_reason_query(
text: list[str],
embedding_model: SentenceTransformer,
classifier_head: ClassifierHead,
) -> dict:
"""
Runs the full inference pipeline for reason classification: Text -> Embedding -> Classification.
"""
# Set models to evaluation mode
embedding_model.eval()
classifier_head.eval()
with torch.no_grad():
# Embed the text
embeddings = embedding_model.encode(
text,
convert_to_tensor=True,
device=DEVICE
).to(DEVICE)
# Calculate probabilities and prediction
probabilities = classifier_head.predict_proba(embeddings)
# Get the predicted index and confidence
predicted_indices = torch.argmax(probabilities, dim=1)
# Convert tensors to Python types safely
if predicted_indices.dim() == 0: # Single prediction
predicted_indices = [predicted_indices.item()]
else:
predicted_indices = predicted_indices.cpu().tolist()
# Get confidences
confidences = []
for i, idx in enumerate(predicted_indices):
conf = probabilities[i][idx].item() if probabilities.dim() > 1 else probabilities[idx].item()
confidences.append(conf)
# Get the predicted label names
predicted_labels = [REASON_CATEGORIES[i] for i in predicted_indices]
return {
'prediction': predicted_labels,
'confidence': confidences,
'probabilities': probabilities.cpu().tolist()
}
def predict_single_reason(query: str) -> dict:
"""Convenience function to predict a single reason query."""
try:
embedding_model, classifier_head = get_reason_models()
# Try to load the most recent trained checkpoint
if os.path.exists(REASON_CHECKPOINT_PATH):
for d in os.listdir(REASON_CHECKPOINT_PATH):
if d.endswith('.pt'):
checkpoint_path = f"{REASON_CHECKPOINT_PATH}/{d}"
try:
state_dict = torch.load(checkpoint_path, weights_only=True, map_location=DEVICE)
classifier_head.load_state_dict(state_dict)
print(f"Loaded trained weights from {checkpoint_path}")
break
except Exception as e:
print(f"Could not load weights from {checkpoint_path}: {e}")
result = predict_reason_query([query], embedding_model, classifier_head)
# Extract values safely
prediction = result['prediction'][0] if isinstance(result['prediction'], list) else str(result['prediction'])
confidence = result['confidence'] if isinstance(result['confidence'], float) else (result['confidence'][0] if isinstance(result['confidence'], list) else float(result['confidence']))
# Handle probabilities - ensure it's a list
probabilities = result['probabilities']
if isinstance(probabilities, list) and len(probabilities) > 0:
if isinstance(probabilities[0], list):
probabilities = probabilities[0]
# Create probability dictionary
prob_dict = {}
for i, category in REASON_CATEGORIES.items():
if i < len(probabilities):
prob_dict[category] = float(probabilities[i])
else:
prob_dict[category] = 0.0
return {
'query': query,
'category': prediction,
'confidence': confidence,
'probabilities': prob_dict
}
except Exception as e:
# Return a default classification if the model fails
return {
'query': query,
'category': 'GENERAL_MEDICAL',
'confidence': 0.5,
'probabilities': {category: 1.0/len(REASON_CATEGORIES) for category in REASON_CATEGORIES.values()},
'error': str(e)
}
def test_reason_classifier():
"""Test the reason classifier with sample queries."""
latest = None
path = ""
# Try to load the most recent checkpoint
if os.path.exists(REASON_CHECKPOINT_PATH):
for d in os.listdir(REASON_CHECKPOINT_PATH):
if d.endswith('.pt'):
checkpoint_path = f"{REASON_CHECKPOINT_PATH}/{d}"
print(f"Found checkpoint: {checkpoint_path}")
path = checkpoint_path
break
if not path:
print("No trained checkpoints found. Using untrained model.")
else:
print("No checkpoint directory found. Using untrained model.")
embedding_model, classifier = get_reason_models()
# Load trained weights if available
if path and os.path.exists(path):
try:
state_dict = torch.load(path, weights_only=True, map_location=DEVICE)
classifier.load_state_dict(state_dict)
print(f"Loaded trained weights from {path}")
except Exception as e:
print(f"Could not load weights: {e}. Using untrained model.")
# Test queries for reason classification
queries = [
"I have heel pain when I walk",
"My toenail is ingrown and painful",
"I need routine foot care",
"I sprained my ankle playing sports",
"I have plantar fasciitis",
"I need a cortisone injection"
]
print("\nTesting reason classification:")
pred = predict_reason_query(
text=queries,
embedding_model=embedding_model,
classifier_head=classifier,
)
pprint.pprint(pred, indent=4)
if __name__ == "__main__":
test_reason_classifier() |