Update app.py

app.py

@@ -22,7 +22,7 @@ def classify_text(text):
     # Get token classification results
     result = pipe(text)
     
-    # Format the results into HTML with color highlighting
+    # Format the results into HTML with color highlighting and entity names
     highlighted_text = ""
     last_pos = 0
 
@@ -35,9 +35,9 @@ def classify_text(text):
         # Add text before the entity without highlighting
         highlighted_text += text[last_pos:start]
         
-        # Add highlighted entity text
+        # Add highlighted entity text with the entity name displayed
         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>"
+        highlighted_text += f"<span style='background-color:{color}; padding:2px; border-radius:5px; position:relative;'><strong>{word}</strong> <small style='color:#000; background-color: #fff; padding: 2px 4px; border-radius: 3px; position: absolute; top: -25px; left: 0;'>{entity}</small></span>"
         
         # Update last position
         last_pos = end
@@ -45,15 +45,15 @@ def classify_text(text):
     # Add the rest of the text after the last entity
     highlighted_text += text[last_pos:]
     
-    return highlighted_text
+    return f"<div style='font-family: Arial, sans-serif; line-height: 1.5;'>{highlighted_text}</div>"
 
 # Gradio Interface
 demo = gr.Interface(
     fn=classify_text,
     inputs=gr.Textbox(lines=5, label="Enter Medical Text"),
-    outputs=gr.HTML(label="Entity Classification with Highlighting"),
+    outputs=gr.HTML(label="Entity Classification with Highlighting and Labels"),
     title="Medical Entity Classification",
-    description="Enter medical-related text, and the model will classify medical entities with color highlighting.",
+    description="Enter medical-related text, and the model will classify medical entities with color highlighting and labels.",
     examples=[
         ["45 year old woman diagnosed with CAD"],
         ["A 65-year-old male presents with acute chest pain and a history of hypertension."],