Spaces:
Running
Running
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,577 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import torch
|
3 |
+
from gliner import GLiNER
|
4 |
+
import pandas as pd
|
5 |
+
import warnings
|
6 |
+
import random
|
7 |
+
import re
|
8 |
+
warnings.filterwarnings('ignore')
|
9 |
+
|
10 |
+
# Standard NER entity types
|
11 |
+
STANDARD_ENTITIES = [
|
12 |
+
'DATE', 'EVENT', 'FAC', 'GPE', 'LANG', 'LOC',
|
13 |
+
'MISC', 'NORP', 'ORG', 'PER', 'PRODUCT', 'WORK_OF_ART'
|
14 |
+
]
|
15 |
+
|
16 |
+
# Color schemes
|
17 |
+
STANDARD_COLORS = {
|
18 |
+
'DATE': '#FF6B6B', # Red
|
19 |
+
'EVENT': '#4ECDC4', # Teal
|
20 |
+
'FAC': '#45B7D1', # Blue
|
21 |
+
'GPE': '#F9CA24', # Yellow
|
22 |
+
'LANG': '#6C5CE7', # Purple
|
23 |
+
'LOC': '#A0E7E5', # Light Cyan
|
24 |
+
'MISC': '#FD79A8', # Pink
|
25 |
+
'NORP': '#8E8E93', # Grey
|
26 |
+
'ORG': '#55A3FF', # Light Blue
|
27 |
+
'PER': '#00B894', # Green
|
28 |
+
'PRODUCT': '#E17055', # Orange-Red
|
29 |
+
'WORK_OF_ART': '#DDA0DD' # Plum
|
30 |
+
}
|
31 |
+
|
32 |
+
# Additional colors for custom entities
|
33 |
+
CUSTOM_COLOR_PALETTE = [
|
34 |
+
'#FF9F43', '#10AC84', '#EE5A24', '#0FBC89', '#5F27CD',
|
35 |
+
'#FF3838', '#2F3640', '#3742FA', '#2ED573', '#FFA502',
|
36 |
+
'#FF6348', '#1E90FF', '#FF1493', '#32CD32', '#FFD700',
|
37 |
+
'#FF4500', '#DA70D6', '#00CED1', '#FF69B4', '#7B68EE'
|
38 |
+
]
|
39 |
+
|
40 |
+
class HybridNERManager:
|
41 |
+
def __init__(self):
|
42 |
+
self.gliner_model = None
|
43 |
+
self.spacy_model = None
|
44 |
+
self.all_entity_colors = {}
|
45 |
+
|
46 |
+
def load_gliner_model(self):
|
47 |
+
"""Load GLiNER model for custom entities"""
|
48 |
+
if self.gliner_model is None:
|
49 |
+
try:
|
50 |
+
# Use a more stable model for HF Spaces
|
51 |
+
self.gliner_model = GLiNER.from_pretrained("urchade/gliner_medium-v2.1")
|
52 |
+
print("β GLiNER model loaded successfully")
|
53 |
+
except Exception as e:
|
54 |
+
print(f"Error loading GLiNER model: {str(e)}")
|
55 |
+
return None
|
56 |
+
return self.gliner_model
|
57 |
+
|
58 |
+
def load_spacy_model(self):
|
59 |
+
"""Load spaCy model for standard NER"""
|
60 |
+
if self.spacy_model is None:
|
61 |
+
try:
|
62 |
+
import spacy
|
63 |
+
# Try to load the transformer model first, fallback to smaller model
|
64 |
+
try:
|
65 |
+
self.spacy_model = spacy.load("en_core_web_sm")
|
66 |
+
print("β spaCy model loaded successfully")
|
67 |
+
except OSError:
|
68 |
+
print("spaCy model not found. Using GLiNER for all entity types.")
|
69 |
+
return None
|
70 |
+
except Exception as e:
|
71 |
+
print(f"Error loading spaCy model: {str(e)}")
|
72 |
+
return None
|
73 |
+
return self.spacy_model
|
74 |
+
|
75 |
+
def assign_colors(self, standard_entities, custom_entities):
|
76 |
+
"""Assign colors to all entity types"""
|
77 |
+
self.all_entity_colors = {}
|
78 |
+
|
79 |
+
# Assign standard colors
|
80 |
+
for entity in standard_entities:
|
81 |
+
self.all_entity_colors[entity.upper()] = STANDARD_COLORS.get(entity, '#CCCCCC')
|
82 |
+
|
83 |
+
# Assign custom colors
|
84 |
+
for i, entity in enumerate(custom_entities):
|
85 |
+
if i < len(CUSTOM_COLOR_PALETTE):
|
86 |
+
self.all_entity_colors[entity.upper()] = CUSTOM_COLOR_PALETTE[i]
|
87 |
+
else:
|
88 |
+
# Generate random color if we run out
|
89 |
+
self.all_entity_colors[entity.upper()] = f"#{random.randint(0, 0xFFFFFF):06x}"
|
90 |
+
|
91 |
+
return self.all_entity_colors
|
92 |
+
|
93 |
+
def extract_spacy_entities(self, text, entity_types):
|
94 |
+
"""Extract entities using spaCy"""
|
95 |
+
model = self.load_spacy_model()
|
96 |
+
if model is None:
|
97 |
+
return []
|
98 |
+
|
99 |
+
try:
|
100 |
+
doc = model(text)
|
101 |
+
entities = []
|
102 |
+
for ent in doc.ents:
|
103 |
+
if ent.label_ in entity_types:
|
104 |
+
entities.append({
|
105 |
+
'text': ent.text,
|
106 |
+
'label': ent.label_,
|
107 |
+
'start': ent.start_char,
|
108 |
+
'end': ent.end_char,
|
109 |
+
'confidence': 1.0, # spaCy doesn't provide confidence scores
|
110 |
+
'source': 'spaCy'
|
111 |
+
})
|
112 |
+
return entities
|
113 |
+
except Exception as e:
|
114 |
+
print(f"Error with spaCy extraction: {str(e)}")
|
115 |
+
return []
|
116 |
+
|
117 |
+
def extract_gliner_entities(self, text, entity_types, threshold=0.3, is_custom=True):
|
118 |
+
"""Extract entities using GLiNER"""
|
119 |
+
model = self.load_gliner_model()
|
120 |
+
if model is None:
|
121 |
+
return []
|
122 |
+
|
123 |
+
try:
|
124 |
+
entities = model.predict_entities(text, entity_types, threshold=threshold)
|
125 |
+
result = []
|
126 |
+
for entity in entities:
|
127 |
+
result.append({
|
128 |
+
'text': entity['text'],
|
129 |
+
'label': entity['label'].upper(),
|
130 |
+
'start': entity['start'],
|
131 |
+
'end': entity['end'],
|
132 |
+
'confidence': entity.get('score', 0.0),
|
133 |
+
'source': 'GLiNER-Custom' if is_custom else 'GLiNER-Standard'
|
134 |
+
})
|
135 |
+
return result
|
136 |
+
except Exception as e:
|
137 |
+
print(f"Error with GLiNER extraction: {str(e)}")
|
138 |
+
return []
|
139 |
+
|
140 |
+
def find_overlapping_entities(entities):
|
141 |
+
"""Find and merge overlapping entities"""
|
142 |
+
if not entities:
|
143 |
+
return []
|
144 |
+
|
145 |
+
# Sort entities by start position
|
146 |
+
sorted_entities = sorted(entities, key=lambda x: x['start'])
|
147 |
+
merged_entities = []
|
148 |
+
|
149 |
+
i = 0
|
150 |
+
while i < len(sorted_entities):
|
151 |
+
current_entity = sorted_entities[i]
|
152 |
+
overlapping_entities = [current_entity]
|
153 |
+
|
154 |
+
# Find all entities that overlap with current entity
|
155 |
+
j = i + 1
|
156 |
+
while j < len(sorted_entities):
|
157 |
+
next_entity = sorted_entities[j]
|
158 |
+
|
159 |
+
# Check if entities overlap
|
160 |
+
if (current_entity['start'] <= next_entity['start'] < current_entity['end'] or
|
161 |
+
next_entity['start'] <= current_entity['start'] < next_entity['end'] or
|
162 |
+
current_entity['text'].lower() == next_entity['text'].lower()):
|
163 |
+
overlapping_entities.append(next_entity)
|
164 |
+
sorted_entities.pop(j)
|
165 |
+
else:
|
166 |
+
j += 1
|
167 |
+
|
168 |
+
# Create merged entity
|
169 |
+
if len(overlapping_entities) == 1:
|
170 |
+
merged_entities.append(overlapping_entities[0])
|
171 |
+
else:
|
172 |
+
merged_entity = merge_entities(overlapping_entities)
|
173 |
+
merged_entities.append(merged_entity)
|
174 |
+
|
175 |
+
i += 1
|
176 |
+
|
177 |
+
return merged_entities
|
178 |
+
|
179 |
+
def merge_entities(entity_list):
|
180 |
+
"""Merge multiple overlapping entities into one"""
|
181 |
+
if len(entity_list) == 1:
|
182 |
+
return entity_list[0]
|
183 |
+
|
184 |
+
# Use the entity with the longest text span as the base
|
185 |
+
base_entity = max(entity_list, key=lambda x: len(x['text']))
|
186 |
+
|
187 |
+
# Collect all labels and sources
|
188 |
+
labels = [entity['label'] for entity in entity_list]
|
189 |
+
sources = [entity['source'] for entity in entity_list]
|
190 |
+
confidences = [entity['confidence'] for entity in entity_list]
|
191 |
+
|
192 |
+
return {
|
193 |
+
'text': base_entity['text'],
|
194 |
+
'start': base_entity['start'],
|
195 |
+
'end': base_entity['end'],
|
196 |
+
'labels': labels,
|
197 |
+
'sources': sources,
|
198 |
+
'confidences': confidences,
|
199 |
+
'is_merged': True,
|
200 |
+
'entity_count': len(entity_list)
|
201 |
+
}
|
202 |
+
|
203 |
+
def create_highlighted_html(text, entities, entity_colors):
|
204 |
+
"""Create HTML with highlighted entities"""
|
205 |
+
if not entities:
|
206 |
+
return f"<div style='padding: 15px; border: 1px solid #ddd; border-radius: 5px; background-color: #fafafa;'><p>{text}</p></div>"
|
207 |
+
|
208 |
+
# Find and merge overlapping entities
|
209 |
+
merged_entities = find_overlapping_entities(entities)
|
210 |
+
|
211 |
+
# Sort by start position
|
212 |
+
sorted_entities = sorted(merged_entities, key=lambda x: x['start'])
|
213 |
+
|
214 |
+
# Create HTML with highlighting
|
215 |
+
html_parts = []
|
216 |
+
last_end = 0
|
217 |
+
|
218 |
+
for entity in sorted_entities:
|
219 |
+
# Add text before entity
|
220 |
+
html_parts.append(text[last_end:entity['start']])
|
221 |
+
|
222 |
+
if entity.get('is_merged', False):
|
223 |
+
# Handle merged entity with multiple colors
|
224 |
+
html_parts.append(create_merged_entity_html(entity, entity_colors))
|
225 |
+
else:
|
226 |
+
# Handle single entity
|
227 |
+
html_parts.append(create_single_entity_html(entity, entity_colors))
|
228 |
+
|
229 |
+
last_end = entity['end']
|
230 |
+
|
231 |
+
# Add remaining text
|
232 |
+
html_parts.append(text[last_end:])
|
233 |
+
|
234 |
+
highlighted_text = ''.join(html_parts)
|
235 |
+
|
236 |
+
return f"""
|
237 |
+
<div style='padding: 15px; border: 2px solid #ddd; border-radius: 8px; background-color: #fafafa; margin: 10px 0;'>
|
238 |
+
<h4 style='margin: 0 0 15px 0; color: #333;'>π Text with Highlighted Entities</h4>
|
239 |
+
<div style='line-height: 1.8; font-size: 16px; background-color: white; padding: 15px; border-radius: 5px;'>{highlighted_text}</div>
|
240 |
+
</div>
|
241 |
+
"""
|
242 |
+
|
243 |
+
def create_single_entity_html(entity, entity_colors):
|
244 |
+
"""Create HTML for a single entity"""
|
245 |
+
label = entity['label']
|
246 |
+
color = entity_colors.get(label.upper(), '#CCCCCC')
|
247 |
+
confidence = entity.get('confidence', 0.0)
|
248 |
+
source = entity.get('source', 'Unknown')
|
249 |
+
|
250 |
+
return (f'<span style="background-color: {color}; padding: 2px 4px; '
|
251 |
+
f'border-radius: 3px; margin: 0 1px; '
|
252 |
+
f'border: 1px solid {color}; color: white; font-weight: bold;" '
|
253 |
+
f'title="{label} ({source}) - confidence: {confidence:.2f}">'
|
254 |
+
f'{entity["text"]}</span>')
|
255 |
+
|
256 |
+
def create_merged_entity_html(entity, entity_colors):
|
257 |
+
"""Create HTML for a merged entity with multiple colors"""
|
258 |
+
labels = entity['labels']
|
259 |
+
sources = entity['sources']
|
260 |
+
confidences = entity['confidences']
|
261 |
+
|
262 |
+
# Get colors for each label
|
263 |
+
colors = []
|
264 |
+
for label in labels:
|
265 |
+
color = entity_colors.get(label.upper(), '#CCCCCC')
|
266 |
+
colors.append(color)
|
267 |
+
|
268 |
+
# Create gradient background
|
269 |
+
if len(colors) == 2:
|
270 |
+
gradient = f"linear-gradient(to right, {colors[0]} 50%, {colors[1]} 50%)"
|
271 |
+
else:
|
272 |
+
# For more colors, create equal segments
|
273 |
+
segment_size = 100 / len(colors)
|
274 |
+
gradient_parts = []
|
275 |
+
for i, color in enumerate(colors):
|
276 |
+
start = i * segment_size
|
277 |
+
end = (i + 1) * segment_size
|
278 |
+
gradient_parts.append(f"{color} {start}%, {color} {end}%")
|
279 |
+
gradient = f"linear-gradient(to right, {', '.join(gradient_parts)})"
|
280 |
+
|
281 |
+
# Create tooltip
|
282 |
+
tooltip_parts = []
|
283 |
+
for i, label in enumerate(labels):
|
284 |
+
tooltip_parts.append(f"{label} ({sources[i]}) - {confidences[i]:.2f}")
|
285 |
+
tooltip = " | ".join(tooltip_parts)
|
286 |
+
|
287 |
+
return (f'<span style="background: {gradient}; padding: 2px 4px; '
|
288 |
+
f'border-radius: 3px; margin: 0 1px; '
|
289 |
+
f'border: 2px solid #333; color: white; font-weight: bold;" '
|
290 |
+
f'title="MERGED: {tooltip}">'
|
291 |
+
f'{entity["text"]} π</span>')
|
292 |
+
|
293 |
+
def create_entity_table_html(entities, entity_colors):
|
294 |
+
"""Create HTML table of entities"""
|
295 |
+
if not entities:
|
296 |
+
return "<p>No entities found.</p>"
|
297 |
+
|
298 |
+
# Merge overlapping entities
|
299 |
+
merged_entities = find_overlapping_entities(entities)
|
300 |
+
|
301 |
+
# Group entities by type
|
302 |
+
entity_groups = {}
|
303 |
+
for entity in merged_entities:
|
304 |
+
if entity.get('is_merged', False):
|
305 |
+
key = 'MERGED_ENTITIES'
|
306 |
+
else:
|
307 |
+
key = entity['label']
|
308 |
+
|
309 |
+
if key not in entity_groups:
|
310 |
+
entity_groups[key] = []
|
311 |
+
entity_groups[key].append(entity)
|
312 |
+
|
313 |
+
# Create HTML table
|
314 |
+
html = "<div style='margin: 15px 0;'>"
|
315 |
+
|
316 |
+
for entity_type, entities_of_type in entity_groups.items():
|
317 |
+
if entity_type == 'MERGED_ENTITIES':
|
318 |
+
color = '#666666'
|
319 |
+
header = f"π Merged Entities ({len(entities_of_type)})"
|
320 |
+
else:
|
321 |
+
color = entity_colors.get(entity_type.upper(), '#CCCCCC')
|
322 |
+
header = f"{entity_type} ({len(entities_of_type)})"
|
323 |
+
|
324 |
+
html += f"""
|
325 |
+
<h4 style="color: {color}; margin: 15px 0 10px 0;">{header}</h4>
|
326 |
+
<table style="width: 100%; border-collapse: collapse; margin-bottom: 20px; border: 1px solid #ddd;">
|
327 |
+
<thead>
|
328 |
+
<tr style="background-color: {color}; color: white;">
|
329 |
+
<th style="padding: 10px; text-align: left; border: 1px solid #ddd;">Entity Text</th>
|
330 |
+
<th style="padding: 10px; text-align: left; border: 1px solid #ddd;">Label(s)</th>
|
331 |
+
<th style="padding: 10px; text-align: left; border: 1px solid #ddd;">Source(s)</th>
|
332 |
+
<th style="padding: 10px; text-align: left; border: 1px solid #ddd;">Confidence</th>
|
333 |
+
</tr>
|
334 |
+
</thead>
|
335 |
+
<tbody>
|
336 |
+
"""
|
337 |
+
|
338 |
+
for entity in entities_of_type:
|
339 |
+
if entity.get('is_merged', False):
|
340 |
+
labels_text = " | ".join(entity['labels'])
|
341 |
+
sources_text = " | ".join(entity['sources'])
|
342 |
+
conf_text = " | ".join([f"{c:.2f}" for c in entity['confidences']])
|
343 |
+
else:
|
344 |
+
labels_text = entity['label']
|
345 |
+
sources_text = entity['source']
|
346 |
+
conf_text = f"{entity['confidence']:.2f}"
|
347 |
+
|
348 |
+
html += f"""
|
349 |
+
<tr style="background-color: #fff;">
|
350 |
+
<td style="padding: 8px; border: 1px solid #ddd; font-weight: bold;">{entity['text']}</td>
|
351 |
+
<td style="padding: 8px; border: 1px solid #ddd;">{labels_text}</td>
|
352 |
+
<td style="padding: 8px; border: 1px solid #ddd;">{sources_text}</td>
|
353 |
+
<td style="padding: 8px; border: 1px solid #ddd;">{conf_text}</td>
|
354 |
+
</tr>
|
355 |
+
"""
|
356 |
+
|
357 |
+
html += "</tbody></table>"
|
358 |
+
|
359 |
+
html += "</div>"
|
360 |
+
return html
|
361 |
+
|
362 |
+
def create_legend_html(entity_colors, standard_entities, custom_entities):
|
363 |
+
"""Create a legend showing entity colors"""
|
364 |
+
if not entity_colors:
|
365 |
+
return ""
|
366 |
+
|
367 |
+
html = "<div style='margin: 15px 0; padding: 15px; background-color: #f8f9fa; border-radius: 8px;'>"
|
368 |
+
html += "<h4 style='margin: 0 0 15px 0;'>π¨ Entity Type Legend</h4>"
|
369 |
+
|
370 |
+
if standard_entities:
|
371 |
+
html += "<div style='margin-bottom: 15px;'>"
|
372 |
+
html += "<h5 style='margin: 0 0 8px 0;'>π― Standard Entities:</h5>"
|
373 |
+
html += "<div style='display: flex; flex-wrap: wrap; gap: 8px;'>"
|
374 |
+
for entity_type in standard_entities:
|
375 |
+
color = entity_colors.get(entity_type.upper(), '#ccc')
|
376 |
+
html += f"<span style='background-color: {color}; padding: 4px 8px; border-radius: 15px; color: white; font-weight: bold; font-size: 12px;'>{entity_type}</span>"
|
377 |
+
html += "</div></div>"
|
378 |
+
|
379 |
+
if custom_entities:
|
380 |
+
html += "<div>"
|
381 |
+
html += "<h5 style='margin: 0 0 8px 0;'>β¨ Custom Entities:</h5>"
|
382 |
+
html += "<div style='display: flex; flex-wrap: wrap; gap: 8px;'>"
|
383 |
+
for entity_type in custom_entities:
|
384 |
+
color = entity_colors.get(entity_type.upper(), '#ccc')
|
385 |
+
html += f"<span style='background-color: {color}; padding: 4px 8px; border-radius: 15px; color: white; font-weight: bold; font-size: 12px;'>{entity_type}</span>"
|
386 |
+
html += "</div></div>"
|
387 |
+
|
388 |
+
html += "</div>"
|
389 |
+
return html
|
390 |
+
|
391 |
+
# Initialize the NER manager
|
392 |
+
ner_manager = HybridNERManager()
|
393 |
+
|
394 |
+
def process_text(text, standard_entities, custom_entities_str, confidence_threshold, use_spacy, use_gliner_standard):
|
395 |
+
"""Main processing function for Gradio interface"""
|
396 |
+
if not text.strip():
|
397 |
+
return "β Please enter some text to analyze", "", ""
|
398 |
+
|
399 |
+
# Parse custom entities
|
400 |
+
custom_entities = []
|
401 |
+
if custom_entities_str.strip():
|
402 |
+
custom_entities = [entity.strip() for entity in custom_entities_str.split(',') if entity.strip()]
|
403 |
+
|
404 |
+
# Parse standard entities
|
405 |
+
selected_standard = [entity for entity in standard_entities if entity]
|
406 |
+
|
407 |
+
if not selected_standard and not custom_entities:
|
408 |
+
return "β Please select at least one standard entity type OR enter custom entity types", "", ""
|
409 |
+
|
410 |
+
all_entities = []
|
411 |
+
|
412 |
+
# Extract standard entities
|
413 |
+
if selected_standard:
|
414 |
+
if use_spacy:
|
415 |
+
spacy_entities = ner_manager.extract_spacy_entities(text, selected_standard)
|
416 |
+
all_entities.extend(spacy_entities)
|
417 |
+
|
418 |
+
if use_gliner_standard:
|
419 |
+
gliner_standard_entities = ner_manager.extract_gliner_entities(text, selected_standard, confidence_threshold, is_custom=False)
|
420 |
+
all_entities.extend(gliner_standard_entities)
|
421 |
+
|
422 |
+
# Extract custom entities
|
423 |
+
if custom_entities:
|
424 |
+
custom_entity_results = ner_manager.extract_gliner_entities(text, custom_entities, confidence_threshold, is_custom=True)
|
425 |
+
all_entities.extend(custom_entity_results)
|
426 |
+
|
427 |
+
if not all_entities:
|
428 |
+
return "β No entities found. Try lowering the confidence threshold or using different entity types.", "", ""
|
429 |
+
|
430 |
+
# Assign colors
|
431 |
+
entity_colors = ner_manager.assign_colors(selected_standard, custom_entities)
|
432 |
+
|
433 |
+
# Create outputs
|
434 |
+
legend_html = create_legend_html(entity_colors, selected_standard, custom_entities)
|
435 |
+
highlighted_html = create_highlighted_html(text, all_entities, entity_colors)
|
436 |
+
table_html = create_entity_table_html(all_entities, entity_colors)
|
437 |
+
|
438 |
+
# Create summary
|
439 |
+
total_entities = len(all_entities)
|
440 |
+
merged_entities = find_overlapping_entities(all_entities)
|
441 |
+
final_count = len(merged_entities)
|
442 |
+
merged_count = sum(1 for e in merged_entities if e.get('is_merged', False))
|
443 |
+
|
444 |
+
summary = f"""
|
445 |
+
## π Analysis Summary
|
446 |
+
- **Total entities found:** {total_entities}
|
447 |
+
- **Final entities displayed:** {final_count}
|
448 |
+
- **Merged entities:** {merged_count}
|
449 |
+
- **Average confidence:** {sum(e.get('confidence', 0) for e in all_entities) / total_entities:.3f}
|
450 |
+
"""
|
451 |
+
|
452 |
+
return summary, legend_html + highlighted_html, table_html
|
453 |
+
|
454 |
+
# Create Gradio interface
|
455 |
+
def create_interface():
|
456 |
+
with gr.Blocks(title="Hybrid NER + GLiNER Tool", theme=gr.themes.Soft()) as demo:
|
457 |
+
gr.Markdown("""
|
458 |
+
# π― Hybrid NER + Custom GLiNER Entity Recognition Tool
|
459 |
+
|
460 |
+
Combine standard NER categories with your own custom entity types! This tool uses both traditional NER models and GLiNER for comprehensive entity extraction.
|
461 |
+
|
462 |
+
## π NEW: Overlapping entities are automatically merged with split-color highlighting!
|
463 |
+
|
464 |
+
### How to use:
|
465 |
+
1. **π Enter your text** in the text area below
|
466 |
+
2. **π― Select standard entities** you want to find (PER, ORG, LOC, etc.)
|
467 |
+
3. **β¨ Add custom entities** (comma-separated) like "relationships, occupations, skills"
|
468 |
+
4. **βοΈ Choose models** and adjust confidence threshold
|
469 |
+
5. **π Click "Analyze Text"** to see results
|
470 |
+
""")
|
471 |
+
|
472 |
+
with gr.Row():
|
473 |
+
with gr.Column(scale=2):
|
474 |
+
text_input = gr.Textbox(
|
475 |
+
label="π Text to Analyze",
|
476 |
+
placeholder="Enter your text here...",
|
477 |
+
lines=6,
|
478 |
+
max_lines=10
|
479 |
+
)
|
480 |
+
|
481 |
+
with gr.Column(scale=1):
|
482 |
+
confidence_threshold = gr.Slider(
|
483 |
+
minimum=0.1,
|
484 |
+
maximum=0.9,
|
485 |
+
value=0.3,
|
486 |
+
step=0.1,
|
487 |
+
label="ποΈ Confidence Threshold"
|
488 |
+
)
|
489 |
+
|
490 |
+
with gr.Row():
|
491 |
+
with gr.Column():
|
492 |
+
gr.Markdown("### π― Standard Entity Types")
|
493 |
+
standard_entities = gr.CheckboxGroup(
|
494 |
+
choices=STANDARD_ENTITIES,
|
495 |
+
value=['PER', 'ORG', 'LOC', 'MISC'], # Default selection
|
496 |
+
label="Select Standard Entities"
|
497 |
+
)
|
498 |
+
|
499 |
+
with gr.Row():
|
500 |
+
use_spacy = gr.Checkbox(label="Use spaCy", value=True)
|
501 |
+
use_gliner_standard = gr.Checkbox(label="Use GLiNER for Standard", value=False)
|
502 |
+
|
503 |
+
with gr.Column():
|
504 |
+
gr.Markdown("### β¨ Custom Entity Types")
|
505 |
+
custom_entities = gr.Textbox(
|
506 |
+
label="Custom Entities (comma-separated)",
|
507 |
+
placeholder="e.g. relationships, occupations, skills, emotions",
|
508 |
+
lines=3
|
509 |
+
)
|
510 |
+
gr.Markdown("""
|
511 |
+
**Examples:**
|
512 |
+
- relationships, occupations, skills
|
513 |
+
- emotions, actions, objects
|
514 |
+
- medical conditions, treatments
|
515 |
+
""")
|
516 |
+
|
517 |
+
analyze_btn = gr.Button("π Analyze Text", variant="primary", size="lg")
|
518 |
+
|
519 |
+
# Output sections
|
520 |
+
with gr.Row():
|
521 |
+
summary_output = gr.Markdown(label="Summary")
|
522 |
+
|
523 |
+
with gr.Row():
|
524 |
+
highlighted_output = gr.HTML(label="Highlighted Text")
|
525 |
+
|
526 |
+
with gr.Row():
|
527 |
+
table_output = gr.HTML(label="Detailed Results")
|
528 |
+
|
529 |
+
# Connect the button to the processing function
|
530 |
+
analyze_btn.click(
|
531 |
+
fn=process_text,
|
532 |
+
inputs=[
|
533 |
+
text_input,
|
534 |
+
standard_entities,
|
535 |
+
custom_entities,
|
536 |
+
confidence_threshold,
|
537 |
+
use_spacy,
|
538 |
+
use_gliner_standard
|
539 |
+
],
|
540 |
+
outputs=[summary_output, highlighted_output, table_output]
|
541 |
+
)
|
542 |
+
|
543 |
+
# Add examples
|
544 |
+
gr.Examples(
|
545 |
+
examples=[
|
546 |
+
[
|
547 |
+
"John Smith works at Google in New York. He graduated from Stanford University in 2015 and specializes in artificial intelligence research. His wife Sarah is a doctor at Mount Sinai Hospital.",
|
548 |
+
["PER", "ORG", "LOC", "DATE"],
|
549 |
+
"relationships, occupations, educational background",
|
550 |
+
0.3,
|
551 |
+
True,
|
552 |
+
False
|
553 |
+
],
|
554 |
+
[
|
555 |
+
"The meeting between CEO Jane Doe and the board of directors at Microsoft headquarters in Seattle discussed the Q4 financial results and the new AI strategy for 2024.",
|
556 |
+
["PER", "ORG", "LOC", "DATE"],
|
557 |
+
"corporate roles, business events, financial terms",
|
558 |
+
0.4,
|
559 |
+
True,
|
560 |
+
True
|
561 |
+
]
|
562 |
+
],
|
563 |
+
inputs=[
|
564 |
+
text_input,
|
565 |
+
standard_entities,
|
566 |
+
custom_entities,
|
567 |
+
confidence_threshold,
|
568 |
+
use_spacy,
|
569 |
+
use_gliner_standard
|
570 |
+
]
|
571 |
+
)
|
572 |
+
|
573 |
+
return demo
|
574 |
+
|
575 |
+
if __name__ == "__main__":
|
576 |
+
demo = create_interface()
|
577 |
+
demo.launch()
|