|
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import pandas as pd
|
|
from sklearn.preprocessing import LabelEncoder, StandardScaler
|
|
from sentence_transformers import SentenceTransformer, util
|
|
import json
|
|
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
|
|
from symspellpy import SymSpell, Verbosity
|
|
|
|
device = torch.device("cpu")
|
|
|
|
class DiseaseClassifier(nn.Module):
|
|
def __init__(self, input_size, num_classes, dropout_rate=0.35665610394511454):
|
|
super(DiseaseClassifier, self).__init__()
|
|
self.fc1 = nn.Linear(input_size, 382)
|
|
self.fc2 = nn.Linear(382, 389)
|
|
self.fc3 = nn.Linear(389, 433)
|
|
self.fc4 = nn.Linear(433, num_classes)
|
|
self.activation = nn.LeakyReLU()
|
|
self.dropout = nn.Dropout(dropout_rate)
|
|
|
|
def forward(self, x):
|
|
x = self.activation(self.fc1(x))
|
|
x = self.dropout(x)
|
|
x = self.activation(self.fc2(x))
|
|
x = self.dropout(x)
|
|
x = self.activation(self.fc3(x))
|
|
x = self.dropout(x)
|
|
x = self.fc4(x)
|
|
return x
|
|
|
|
|
|
class DiseasePredictionModel:
|
|
def __init__(self, ai_model_name="model.pth", data_file="data.csv", symptom_json="symptoms.json", dictionary_file="frequency_dictionary_en_82_765.txt"):
|
|
self.df = pd.read_csv(data_file)
|
|
self.symptom_columns = self.load_symptoms(symptom_json)
|
|
self.label_encoder = LabelEncoder()
|
|
self.label_encoder.fit(self.df.iloc[:, 0])
|
|
self.scaler = StandardScaler()
|
|
self.scaler.fit(self.df.iloc[:, 1:].values)
|
|
self.input_size = len(self.symptom_columns)
|
|
self.num_classes = len(self.label_encoder.classes_)
|
|
self.model = self._load_model(ai_model_name)
|
|
self.SYMPTOM_LIST = self.load_symptoms(symptom_json)
|
|
self.sym_spell = SymSpell(max_dictionary_edit_distance=2, prefix_length=7)
|
|
self.sym_spell.load_dictionary(dictionary_file, term_index=0, count_index=1)
|
|
self.tokenizer = AutoTokenizer.from_pretrained("alvaroalon2/biobert_diseases_ner")
|
|
self.nlp_model = AutoModelForTokenClassification.from_pretrained("alvaroalon2/biobert_diseases_ner")
|
|
self.ner_pipeline = pipeline("ner", model=self.nlp_model, tokenizer=self.tokenizer, aggregation_strategy="simple")
|
|
self.semantic_model = SentenceTransformer('all-MiniLM-L6-v2')
|
|
|
|
def _load_model(self, ai_model_name):
|
|
model = DiseaseClassifier(self.input_size, self.num_classes).to(device)
|
|
model.load_state_dict(torch.load(ai_model_name, map_location=device, weights_only=True))
|
|
model.eval()
|
|
return model
|
|
|
|
def predict_disease(self, symptoms):
|
|
input_vector = np.zeros(len(self.symptom_columns))
|
|
for symptom in symptoms:
|
|
if symptom in self.symptom_columns:
|
|
input_vector[list(self.symptom_columns).index(symptom)] = 1
|
|
|
|
input_vector = self.scaler.transform([input_vector])
|
|
|
|
input_tensor = torch.tensor(input_vector, dtype=torch.float32).to(device)
|
|
|
|
with torch.no_grad():
|
|
outputs = self.model(input_tensor)
|
|
_, predicted_class = torch.max(outputs, 1)
|
|
|
|
predicted_disease = self.label_encoder.inverse_transform([predicted_class.cpu().numpy()[0]])[0]
|
|
return predicted_disease
|
|
|
|
def load_symptoms(self, json_file):
|
|
with open(json_file, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
def correct_text(self, text):
|
|
words = text.split()
|
|
corrected_words = []
|
|
|
|
for word in words:
|
|
if word.lower() in [symptom.lower() for symptom in self.SYMPTOM_LIST]:
|
|
corrected_words.append(word)
|
|
else:
|
|
suggestions = self.sym_spell.lookup(word, Verbosity.CLOSEST, max_edit_distance=2)
|
|
if suggestions:
|
|
corrected_words.append(suggestions[0].term)
|
|
else:
|
|
corrected_words.append(word)
|
|
return ' '.join(corrected_words)
|
|
|
|
def extract_symptoms(self, text):
|
|
ner_results = self.ner_pipeline(text)
|
|
symptoms = set()
|
|
for entity in ner_results:
|
|
if entity["entity_group"] == "DISEASE":
|
|
symptoms.add(entity["word"].lower())
|
|
return list(symptoms)
|
|
|
|
def match_symptoms(self, extracted_symptoms):
|
|
matched = {}
|
|
|
|
symptom_embeddings = self.semantic_model.encode(self.SYMPTOM_LIST, convert_to_tensor=True)
|
|
|
|
for symptom in extracted_symptoms:
|
|
symptom_embedding = self.semantic_model.encode(symptom, convert_to_tensor=True)
|
|
|
|
similarities = util.pytorch_cos_sim(symptom_embedding, symptom_embeddings)[0]
|
|
|
|
most_similar_idx = similarities.argmax()
|
|
best_match = self.SYMPTOM_LIST[most_similar_idx]
|
|
matched[symptom] = best_match
|
|
|
|
return matched.values()
|
|
|
|
|
