Update app.py

app.py

@@ -34,17 +34,27 @@ def load_icd_codes_from_files():
                 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()
-                            icd_codes[code] = description
+                        # 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 = {}
@@ -102,6 +112,11 @@ def predict_codes(text):
     # 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)
+    top_k = min(3, len(ICD_CODES))
+    top_icd = torch.topk(icd_probs, k=top_k)
+
     
     # Format results
     result = "Recommended ICD-10 Codes:\n"