Update app.py

app.py

@@ -1,38 +1,40 @@
+import requests
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch.nn.functional as F
-import os
 
-# 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:
-                        parts = line.strip().split(maxsplit=1)
-                        if len(parts) == 2:
-                            code = parts[0].strip()
-                            description = parts[1].strip()
-                            codes[code] = description
-    else:
-        print(f"Directory {directory_path} does not exist!")
-    return 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")
-
-# 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.")
+# API endpoints (replace with actual API endpoints)
+ICD_API_URL = "https://icd10api.com/"
+CPT_API_URL = "https://api.ama-assn.org/cpt/codes"  # Example, requires AMA license
 
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
-model = AutoModelForSequenceClassification.from_pretrained("emilyalsentzer/Bio_ClinicalBERT", num_labels=len(ICD_CODES))
+model = AutoModelForSequenceClassification.from_pretrained("emilyalsentzer/Bio_ClinicalBERT", num_labels=1000)  # Adjust num_labels as needed
+
+# Function to fetch ICD codes from API
+def fetch_icd_codes(query):
+    try:
+        response = requests.get(f"{ICD_API_URL}?desc={query}")
+        if response.status_code == 200:
+            return response.json()  # Adjust based on API response format
+        else:
+            return []
+    except Exception as e:
+        print(f"Error fetching ICD codes: {e}")
+        return []
+
+# Function to fetch CPT codes from API
+def fetch_cpt_codes(query):
+    try:
+        response = requests.get(f"{CPT_API_URL}?desc={query}")
+        if response.status_code == 200:
+            return response.json()  # Adjust based on API response format
+        else:
+            return []
+    except Exception as e:
+        print(f"Error fetching CPT codes: {e}")
+        return []
 
 # Prediction function
 def predict_codes(text):
@@ -57,22 +59,21 @@ def predict_codes(text):
     # Get probabilities
     probs = F.softmax(logits, dim=1)
     
-    # Get top 3 predictions for ICD and CPT
-    top_k = min(3, len(ICD_CODES))
-    top_icd = torch.topk(probs, k=top_k)
+    # Get top 3 predictions
+    top_k = torch.topk(probs, k=3)
+    
+    # Fetch ICD and CPT codes using APIs
+    icd_results = fetch_icd_codes(text)
+    cpt_results = fetch_cpt_codes(text)
     
     # Format results
     result = "Recommended ICD-10 Codes:\n"
-    for i, (prob, idx) in enumerate(zip(top_icd.values[0], top_icd.indices[0])):
-        code = list(ICD_CODES.keys())[idx.item()]
-        description = ICD_CODES[code]
-        result += f"{i+1}. {code}: {description} (Confidence: {prob.item():.2f})\n"
+    for i, code in enumerate(icd_results[:3]):  # Show top 3 ICD codes
+        result += f"{i+1}. {code.get('code', 'Unknown')}: {code.get('description', 'No description')}\n"
     
     result += "\nRecommended CPT Codes:\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"
+    for i, code in enumerate(cpt_results[:3]):  # Show top 3 CPT codes
+        result += f"{i+1}. {code.get('code', 'Unknown')}: {code.get('description', 'No description')}\n"
     
     return result