Update app.py

app.py

@@ -1,93 +1,38 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch.nn.functional as F
 import os
 
-# Define the model class
-class MedicalCodePredictor(torch.nn.Module):
-    def __init__(self, bert_model):
-        super().__init__()
-        self.bert = bert_model
-        self.dropout = torch.nn.Dropout(0.1)
-        self.icd_classifier = torch.nn.Linear(768, len(ICD_CODES))
-        self.cpt_classifier = torch.nn.Linear(768, len(CPT_CODES))
-        
-    def forward(self, input_ids, attention_mask):
-        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        pooled_output = outputs.last_hidden_state[:, 0, :]
-        pooled_output = self.dropout(pooled_output)
-        
-        icd_logits = self.icd_classifier(pooled_output)
-        cpt_logits = self.cpt_classifier(pooled_output)
-        
-        return icd_logits, cpt_logits
-
-# Load ICD codes from files
-def load_icd_codes_from_files():
-    icd_codes = {}
-    directory_path = "./codes/icd_txt_files/"  # Path to ICD codes directory
-
+# Load ICD and CPT codes from files
+def load_codes_from_files(directory_path, code_type):
+    codes = {}
     if os.path.exists(directory_path):
         for file_name in os.listdir(directory_path):
             if file_name.endswith(".txt"):
                 file_path = os.path.join(directory_path, file_name)
                 with open(file_path, "r", encoding="utf-8") as file:
                     for line in file:
-                        # Skip empty lines
-                        if line.strip():
-                            # Split the line into code and description
-                            parts = line.strip().split(maxsplit=1)
-                            if len(parts) == 2:
-                                code = parts[0].strip()
-                                description = parts[1].strip()
-                                icd_codes[code] = description
-                            else:
-                                print(f"Invalid line format in file {file_name}: {line}")
-    else:
-        print(f"Directory {directory_path} does not exist!")
-
-    if not icd_codes:
-        raise ValueError("No ICD codes were loaded. Please check your files and directory structure.")
-
-    return icd_codes
-
-    ICD_CODES = load_icd_codes_from_files()
-    print(f"Loaded {len(ICD_CODES)} ICD codes.")
-
-# Load CPT codes from files
-def load_cpt_codes_from_files():
-    cpt_codes = {}
-    directory_path = "./codes/cpt_txt_files/"  # Path to CPT codes directory
-
-    if os.path.exists(directory_path):
-        for file_name in os.listdir(directory_path):
-            if file_name.endswith(".txt"):
-                file_path = os.path.join(directory_path, file_name)
-                with open(file_path, "r", encoding="utf-8") as file:
-                    for line in file:
-                        # Split the line into code and description
                         parts = line.strip().split(maxsplit=1)
                         if len(parts) == 2:
                             code = parts[0].strip()
                             description = parts[1].strip()
-                            cpt_codes[code] = description
+                            codes[code] = description
     else:
         print(f"Directory {directory_path} does not exist!")
+    return codes
 
-    return cpt_codes
+# Load ICD and CPT codes
+ICD_CODES = load_codes_from_files("./codes/icd_txt_files/", "ICD")
+CPT_CODES = load_codes_from_files("./codes/cpt_txt_files/", "CPT")
 
-# Load ICD and CPT codes dynamically
-ICD_CODES = load_icd_codes_from_files()
-CPT_CODES = load_cpt_codes_from_files()
+# Check if codes were loaded
+if not ICD_CODES or not CPT_CODES:
+    raise ValueError("No ICD or CPT codes were loaded. Please check your files and directory structure.")
 
-# Load models
-@torch.no_grad()
-def load_models():
-    tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
-    base_model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
-    model = MedicalCodePredictor(base_model)
-    return tokenizer, model
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+model = AutoModelForSequenceClassification.from_pretrained("emilyalsentzer/Bio_ClinicalBERT", num_labels=len(ICD_CODES))
 
 # Prediction function
 def predict_codes(text):
@@ -95,43 +40,42 @@ def predict_codes(text):
         return "Please enter a medical summary."
     
     # Tokenize input
-    inputs = tokenizer(text, 
-                      return_tensors="pt",
-                      max_length=512,
-                      truncation=True,
-                      padding=True)
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        max_length=512,
+        truncation=True,
+        padding=True
+    )
     
     # Get predictions
     model.eval()
-    icd_logits, cpt_logits = model(inputs['input_ids'], inputs['attention_mask'])
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
     
     # Get probabilities
-    icd_probs = F.softmax(icd_logits, dim=1)
-    cpt_probs = F.softmax(cpt_logits, dim=1)
+    probs = F.softmax(logits, dim=1)
     
-    # Get top 3 predictions
-    top_icd = torch.topk(icd_probs, k=3)
-    top_cpt = torch.topk(cpt_probs, k=3)
-
-    # Get top k predictions (limit k to the number of available codes)
+    # Get top 3 predictions for ICD and CPT
     top_k = min(3, len(ICD_CODES))
-    top_icd = torch.topk(icd_probs, k=top_k)
-
+    top_icd = torch.topk(probs, k=top_k)
     
     # Format results
     result = "Recommended ICD-10 Codes:\n"
     for i, (prob, idx) in enumerate(zip(top_icd.values[0], top_icd.indices[0])):
-        result += f"{i+1}. {ICD_CODES.get(idx.item(), 'Unknown')} (Confidence: {prob.item():.2f})\n"
+        code = list(ICD_CODES.keys())[idx.item()]
+        description = ICD_CODES[code]
+        result += f"{i+1}. {code}: {description} (Confidence: {prob.item():.2f})\n"
     
     result += "\nRecommended CPT Codes:\n"
-    for i, (prob, idx) in enumerate(zip(top_cpt.values[0], top_cpt.indices[0])):
-        result += f"{i+1}. {CPT_CODES.get(idx.item(), 'Unknown')} (Confidence: {prob.item():.2f})\n"
+    for i, (prob, idx) in enumerate(zip(top_icd.values[0], top_icd.indices[0])):
+        code = list(CPT_CODES.keys())[idx.item()]
+        description = CPT_CODES[code]
+        result += f"{i+1}. {code}: {description} (Confidence: {prob.item():.2f})\n"
     
     return result
 
-# Load models globally
-tokenizer, model = load_models()
-
 # Create Gradio interface
 iface = gr.Interface(
     fn=predict_codes,
@@ -142,7 +86,7 @@ iface = gr.Interface(
     ),
     outputs=gr.Textbox(
         label="Predicted Codes",
-        lines=8
+        lines=10
     ),
     title="AutoRCM - Medical Code Predictor",
     description="Enter a medical summary to get recommended ICD-10 and CPT codes.",
@@ -154,4 +98,4 @@ iface = gr.Interface(
 )
 
 # Launch the interface
-iface.launch(share=True)
+iface.launch(share=True)