ADFLER / app.py
swardiantara's picture
fix highlight issue
da8a100
import gradio as gr
from transformers import pipeline
# --------------------------------------------------------------------------
# 1. Load your NER model from the Hub using a pipeline
# --------------------------------------------------------------------------
# Replace this with your actual model name on the Hugging Face Hub
model_name = "swardiantara/ADFLER-xlnet-base-cased"
ner_pipeline = pipeline("ner", model=model_name, aggregation_strategy="simple")
# --------------------------------------------------------------------------
# 2. Define the prediction function
# --------------------------------------------------------------------------
# This function takes raw text and returns a format that Gradio's HighlightedText component understands.
def recognize_log_events(text):
"""
Performs NER and robustly formats the output for Gradio's HighlightedText.
"""
if not text:
return []
ner_results = ner_pipeline(text)
# Sort entities by their start index to process them in order
ner_results.sort(key=lambda x: x['start'])
highlighted_output = []
last_end = 0
for entity in ner_results:
# Add the text between the last entity and this one (un-highlighted)
if entity['start'] > last_end:
highlighted_output.append((text[last_end:entity['start']], None))
# Add the highlighted entity text
highlighted_output.append((entity['word'], entity['entity_group']))
last_end = entity['end']
# Add any remaining text after the last entity (un-highlighted)
if last_end < len(text):
highlighted_output.append((text[last_end:], None))
return highlighted_output
# --------------------------------------------------------------------------
# 3. Create the Gradio Interface
# --------------------------------------------------------------------------
# A brief description of your project to display in the app
description = """
This demo showcases an NER model for recognizing key events in drone flight logs,
a part of my PhD research in digital forensics at ITS.
Enter a line from a drone log to see the model identify events like 'Takeoff', 'Landing', 'GPS Lock', etc.
"""
# An article providing more context and links
article = """
<div style='text-align: center;'>
<p>For more details, check out the project on GitHub:</p>
<a href='https://dronenlp.github.io/documentation' target='_blank'>DroneNLP Project</a> |
<a href='https://huggingface.co/swardiantara/ADFLER-xlnet-base-cased' target='_blank'>Model Card</a> |
</div>
"""
# <a href='https' target='_blank'>Research Paper (if applicable)</a>
# Example drone log entries for users to try
examples = [
["No image transmission. RTH.; Press Brake button to cancel RTH.; No image transmission. Aircraft returning to home.; Image transmission signal weak. Adjust antennas and make sure they are perpendicular to flight direction of aircraft.; Flight mode changed to Go Home."],
["Battery temperature is below 15 degrees Celsius. Warm up the battery temperature to above 25 degree Celsius to ensure a safe flight."],
["Strong wireless interference. Please fly with caution. Obstacle Avoidance Disabled. Landing gear lowered. Obstacle Avoidance Disabled."],
["Cannot switch flight mode. Turn on 'Multiple Flight Modes' to enable Atti and Sport Modes."]
]
# The main interface
iface = gr.Interface(
fn=recognize_log_events,
inputs=gr.Textbox(
lines=5,
label="Drone Flight Log Entry",
placeholder="Paste a log entry here..."
),
outputs=gr.HighlightedText(
label="Recognized Events",
color_map={"Event": "green"} # Customize these labels and colors!
),
title="🚁 Drone Flight Log Event Recognizer",
description=description,
article=article,
examples=examples
)
# Launch the app!
iface.launch()