| # distilbert-disease-specialist-recommendation | |
| ## Model Overview | |
| This is a **Zero-shot Classification Model** designed to classify the appropriate medical department or specialist a patient should consult based on their symptoms. The model is built using **DistilBERT** and trained for multi-class classification. | |
| ### **Supported Classes:** | |
| - **Cardiology** (Heart-related symptoms) | |
| - **Neurology** (Brain and nervous system issues) | |
| - **Orthopedics** (Bone and muscle problems) | |
| - **Dermatology** (Skin-related conditions) | |
| ## Training Data | |
| The model is trained on a curated dataset of patient symptoms and their corresponding medical specialties. The dataset includes textual descriptions of symptoms and their respective labels mapped to one of the four medical departments. | |
| The training data is collected from various medical sources, including: | |
| - **Electronic Health Records (EHRs)**: Anonymized patient records detailing symptoms and diagnoses. | |
| - **Medical Research Papers**: Information extracted from clinical studies. | |
| - **Expert-Labeled Datasets**: Data manually classified by medical professionals. | |
| Preprocessing steps include: | |
| 1. **Tokenization** using DistilBERT’s tokenizer. | |
| 2. **Stopword Removal** to eliminate non-essential words. | |
| 3. **Synonym Mapping** to standardize medical terminology. | |
| 4. **Class Balancing** to ensure equal representation of all departments. | |
| ## Model Configuration | |
| - **Architecture:** DistilBERT for Sequence Classification | |
| - **Hidden Dimension:** 768 | |
| - **Number of Layers:** 6 | |
| - **Attention Heads:** 12 | |
| - **Dropout:** 0.1 | |
| - **Maximum Token Length:** 512 | |
| - **Activation Function:** GELU | |
| - **Optimizer:** AdamW | |
| - **Batch Size:** 16 | |
| - **Epochs:** 3 | |
| - **Transformers Version:** 4.48.3 | |
| ## How to Use the Model | |
| This model can be easily loaded and used with the Hugging Face `transformers` library. | |
| ### **Installation** | |
| ``` | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| import torch | |
| import numpy as np | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| ``` | |
| ### **Loading The Model** | |
| ``` | |
| model = AutoModelForSequenceClassification.from_pretrained("AventIQ-AI/distilbert-disease-specialist-recommendation") | |
| tokenizer = AutoTokenizer.from_pretrained("AventIQ-AI/distilbert-disease-specialist-recommendation") | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model.to(device) | |
| def encode_text(text): | |
| inputs = tokenizer(text, return_tensors="pt", truncation=True, padding="max_length").to(device) | |
| with torch.no_grad(): | |
| outputs = model(**inputs, output_hidden_states=True) | |
| embedding = outputs.hidden_states[-1].mean(dim=1).squeeze(0).cpu().numpy() | |
| return embedding | |
| candidate_labels = ["Cardiology", "Neurology", "Orthopedics", "Dermatology"] | |
| candidate_embeddings = np.array([encode_text(label) for label in candidate_labels]) | |
| def zero_shot_classification_batch(texts): | |
| text_embeddings = np.array([encode_text(text) for text in texts]) | |
| similarities = cosine_similarity(text_embeddings, candidate_embeddings) | |
| ranked_results = [ | |
| sorted(zip(candidate_labels, similarity), key=lambda x: x[1], reverse=True) | |
| for similarity in similarities | |
| ] | |
| return ranked_results | |
| print(zero_shot_classification_batch(["Suffering from High Fever and body Itching"])) | |
| # EXPECTED OUTPUT : [[('Dermatology', 0.55948263), ('Orthopedics', 0.29858905), ('Cardiology', 0.25098807), ('Neurology', 0.17517035)]] | |
| ``` | |
| ## Model Files | |
| - **model.safetensors:** The trained model weights | |
| - **config.json:** Model configuration | |
| - **tokenizer_config.json:** Tokenizer settings | |
| - **special_tokens_map.json:** Special tokens used in the model | |
| - **vocab.txt:** Vocabulary file for tokenization |