Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import BertForSequenceClassification, BertTokenizer | |
| from sklearn.preprocessing import LabelEncoder | |
| import torch | |
| import numpy as np | |
| # Load model and label encoder | |
| def load_model_and_label_encoder(): | |
| fine_tuned_model = BertForSequenceClassification.from_pretrained('./fine_tuned_model') | |
| label_encoder = LabelEncoder() | |
| label_encoder.classes_ = np.load('./label_encoder_classes.npy', allow_pickle=True) | |
| tokenizer = BertTokenizer.from_pretrained("./tokenizer") # Load tokenizer from local file | |
| return fine_tuned_model, label_encoder, tokenizer | |
| def predict(symptom, fine_tuned_model, label_encoder, tokenizer): | |
| user_input_encoding = tokenizer(symptom, padding=True, truncation=True, return_tensors='pt', max_length=512, return_attention_mask=True, return_token_type_ids=True) | |
| with torch.no_grad(): | |
| logits = fine_tuned_model(**user_input_encoding) | |
| probabilities = torch.nn.functional.softmax(logits.logits, dim=1).numpy()[0] | |
| predicted_labels = np.argsort(-probabilities)[:5] | |
| predicted_diseases = label_encoder.inverse_transform(predicted_labels) | |
| predicted_probabilities = probabilities[predicted_labels] | |
| predictions = [{'disease': disease, 'probability': probability * 100} for disease, probability in zip(predicted_diseases, predicted_probabilities)] | |
| return predictions | |
| def main(): | |
| st.title("Disease Prediction App") | |
| # Load model and label encoder | |
| fine_tuned_model, label_encoder, tokenizer = load_model_and_label_encoder() | |
| # Input | |
| symptom = st.text_input("Enter symptom:", "") | |
| # Predict | |
| if st.button("Predict"): | |
| if symptom: | |
| predictions = predict(symptom, fine_tuned_model, label_encoder, tokenizer) | |
| st.write("Top 5 Predictions:") | |
| for prediction in predictions: | |
| st.write(f"Disease: {prediction['disease']}, Probability: {prediction['probability']:.2f}%") | |
| if __name__ == '__main__': | |
| main() | |