Create app.py

app.py

@@ -0,0 +1,138 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModel
+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
+
+    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:
+                        parts = line.strip().split("\t")  # Adjust delimiter as needed
+                        if len(parts) >= 2:
+                            code = parts[0].strip()
+                            description = parts[1].strip()
+                            icd_codes[code] = description
+    else:
+        print(f"Directory {directory_path} does not exist!")
+    return 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:
+                        parts = line.strip().split("\t")  # Adjust delimiter as needed
+                        if len(parts) >= 2:
+                            code = parts[0].strip()
+                            description = parts[1].strip()
+                            cpt_codes[code] = description
+    else:
+        print(f"Directory {directory_path} does not exist!")
+    return cpt_codes
+
+# Load ICD and CPT codes dynamically
+ICD_CODES = load_icd_codes_from_files()
+CPT_CODES = load_cpt_codes_from_files()
+
+# 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
+
+# Prediction function
+def predict_codes(text):
+    if not text.strip():
+        return "Please enter a medical summary."
+    
+    # Tokenize input
+    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'])
+    
+    # Get probabilities
+    icd_probs = F.softmax(icd_logits, dim=1)
+    cpt_probs = F.softmax(cpt_logits, dim=1)
+    
+    # Get top 3 predictions
+    top_icd = torch.topk(icd_probs, k=3)
+    top_cpt = torch.topk(cpt_probs, k=3)
+    
+    # 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"
+    
+    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"
+    
+    return result
+
+# Load models globally
+tokenizer, model = load_models()
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=predict_codes,
+    inputs=gr.Textbox(
+        lines=5,
+        placeholder="Enter medical summary here...",
+        label="Medical Summary"
+    ),
+    outputs=gr.Textbox(
+        label="Predicted Codes",
+        lines=8
+    ),
+    title="AutoRCM - Medical Code Predictor",
+    description="Enter a medical summary to get recommended ICD-10 and CPT codes.",
+    examples=[
+        ["Patient presents with blood pressure 150/90. Complains of occasional headaches. History of hypertension."],
+        ["Patient has elevated blood sugar levels. A1C is 7.8. History of type 2 diabetes."],
+        ["Patient complains of chronic lower back pain, worse with movement. No radiation to legs."]
+    ]
+)
+
+# Launch the interface
+iface.launch(share=True)