Lech-Iyoko/Symptom-to-Condition_Classifier

Symptom-to-Condition Classifier

This repository contains the artefacts for a LightGBM classification model that predicts a likely medical condition based on a user's textual description of their symptoms.

This model is a proof-of-concept for a portfolio project and is NOT a medical diagnostic tool.

Model Details

This is not a standard end-to-end transformer model. It is a classical machine learning pipeline that uses pre-trained transformers for feature extraction.

• Feature Extractor: The model uses embeddings from emilyalsentzer/Bio_ClinicalBERT. Specifically, it generates a 768-dimension vector for each symptom description by applying mean pooling to the last hidden state of the BERT model.
• Classifier: The actual classification is performed by a LightGBM (Light Gradient Boosting Machine) model trained on the embeddings.

Intended Use

This model is intended for educational and demonstrational purposes only. It takes a string of text describing symptoms and outputs a predicted medical condition from a predefined list of 22 classes.

Ethical Considerations & Limitations

• ⚠️ Not for Medical Use: This model should NEVER be used to diagnose, treat, or provide medical advice for real-world health issues. It is not a substitute for consultation with a qualified healthcare professional.
• Data Bias: The model's knowledge is limited to the gretelai/symptom_to_diagnosis dataset. It cannot predict any condition outside of its 22-class training data and may perform poorly on symptom descriptions that are stylistically different from the training set.
• Correlation, Not Causation: The model learns statistical correlations between words and labels. It has no true understanding of biology or medicine.

Training Data

This model was trained on the gretelai/symptom_to_diagnosis dataset, which contains ~1000 symptom descriptions across 22 balanced classes.

Evaluation

The model achieves a Macro F1-score of 0.834 and an Accuracy of 0.835 on the test set.

How to Use

To use this model, you must load the feature extractor (Bio_ClinicalBERT), the LightGBM classifier, and the label encoder.

import torch
import joblib
from transformers import AutoTokenizer, AutoModel
from huggingface_hub import hf_hub_download

# --- CONFIGURATION ---
HF_REPO_ID = "<YourUsername>/Symptom-to-Condition-Classifier" # Replace with your repo ID
LGBM_MODEL_FILENAME = "lgbm_disease_classifier.joblib"
LABEL_ENCODER_FILENAME = "label_encoder.joblib"
BERT_MODEL_NAME = "emilyalsentzer/Bio_ClinicalBERT"

# --- LOAD ARTIFACTS ---
tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME)
bert_model = AutoModel.from_pretrained(BERT_MODEL_NAME)
lgbm_model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=LGBM_MODEL_FILENAME)
label_encoder_path = hf_hub_download(repo_id=HF_REPO_ID, filename=LABEL_ENCODER_FILENAME)
lgbm_model = joblib.load(lgbm_model_path)
label_encoder = joblib.load(label_encoder_path)

# --- INFERENCE PIPELINE ---
def mean_pool(model_output, attention_mask):
    token_embeddings = model_output.last_hidden_state
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
    return sum_embeddings / sum_mask

def predict_condition(text):
    encoded_input = tokenizer(text, padding=True, truncation=True, max_length=256, return_tensors='pt')
    with torch.no_grad():
        model_output = bert_model(**encoded_input)
    embedding = mean_pool(model_output, encoded_input['attention_mask'])
    
    prediction_id = lgbm_model.predict(embedding.cpu().numpy())
    predicted_condition = label_encoder.inverse_transform(prediction_id)[0]
    return predicted_condition

# --- EXAMPLE ---
symptoms = "I have a burning sensation in my stomach that gets worse when I haven't eaten."
prediction = predict_condition(symptoms)
print(f"Predicted Condition: {prediction}")