Update app.py

app.py

@@ -4,29 +4,56 @@ from transformers import pipeline
 # Load the token classification model
 pipe = pipeline("token-classification", model="Clinical-AI-Apollo/Medical-NER", aggregation_strategy='simple')
 
+# Define colors for different entity types
+entity_colors = {
+    "AGE": "#ffadad",
+    "SEX": "#ffd6a5",
+    "DISEASE_DISORDER": "#caffbf",
+    "SIGN_SYMPTOM": "#9bf6ff",
+    "LAB_VALUE": "#a0c4ff",
+    "THERAPEUTIC_PROCEDURE": "#bdb2ff",
+    "CLINICAL_EVENT": "#ffc6ff",
+    "DIAGNOSTIC_PROCEDURE": "#fffffc",
+    "DETAILED_DESCRIPTION": "#fdffb6",
+    "BIOLOGICAL_STRUCTURE": "#ffb5a7"
+}
+
 def classify_text(text):
     # Get token classification results
     result = pipe(text)
     
-    # Format the results to resemble the UI shown in the image
-    formatted_output = ""
+    # Format the results into HTML with color highlighting
+    highlighted_text = ""
+    last_pos = 0
+
     for res in result:
         entity = res['entity_group']
         word = res['word']
-        score = res['score']
         start = res['start']
         end = res['end']
-        formatted_output += f"Entity: {entity}, Word: {word}, Score: {score:.4f}, Span: [{start}:{end}]\n"
+        
+        # Add text before the entity without highlighting
+        highlighted_text += text[last_pos:start]
+        
+        # Add highlighted entity text
+        color = entity_colors.get(entity, "#e0e0e0")  # Default to gray if entity type not defined
+        highlighted_text += f"<span style='background-color:{color}; padding:2px; border-radius:5px;'>{word}</span>"
+        
+        # Update last position
+        last_pos = end
+    
+    # Add the rest of the text after the last entity
+    highlighted_text += text[last_pos:]
     
-    return formatted_output
+    return highlighted_text
 
 # Gradio Interface
 demo = gr.Interface(
     fn=classify_text,
     inputs=gr.Textbox(lines=5, label="Enter Medical Text"),
-    outputs=gr.Textbox(label="Entity Classification"),
+    outputs=gr.HTML(label="Entity Classification with Highlighting"),
     title="Medical Entity Classification",
-    description="Enter medical-related text, and the model will classify medical entities.",
+    description="Enter medical-related text, and the model will classify medical entities with color highlighting.",
     examples=[
         ["45 year old woman diagnosed with CAD"],
         ["A 65-year-old male presents with acute chest pain and a history of hypertension."],