| --- |
| license: apache-2.0 |
| language: en |
| datasets: |
| - gretelai/symptom_to_diagnosis |
| metrics: |
| - f1 |
| pipeline_tag: text-classification |
| widget: |
| - text: "I have a sharp pain in my chest, difficulty breathing, and a persistent cough." |
| --- |
| |
| # 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. |
|
|
| ```python |
| 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}") |
