File size: 9,553 Bytes
ea99abb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 |
import gradio as gr
import pandas as pd
from utils.ner_helpers import is_llm_model
from typing import Dict, List, Any, Tuple
from tasks.knowledge_graph import build_knowledge_graph, visualize_knowledge_graph_interactive
import base64
from io import BytesIO
def kg_ui():
"""Knowledge Graph UI component"""
# Define models
KG_MODELS = [
"gemini-2.0-flash",
"gpt-4",
"claude-2",
"en_core_web_sm",
"en_core_web_md"
]
DEFAULT_MODEL = "gemini-2.0-flash"
def build_kg(text, model, custom_instructions, interactive=False):
"""Process text for knowledge graph generation"""
import gradio as gr
if not text.strip():
# Trả về các giá trị rỗng cho tất cả các tab
return (
"<div style='text-align: center; color: #666; padding: 20px;'>No text provided</div>",
pd.DataFrame(),
pd.DataFrame(),
False, True, False, True, False, True
)
use_llm = is_llm_model(model)
result = build_knowledge_graph(
text=text,
model_name=model,
use_llm=use_llm
)
entities = result.get("entities", [])
relations = result.get("relations", [])
visualization = result.get("visualization")
# DataFrames
if entities:
entities_df = pd.DataFrame(entities)
entities_df = entities_df.rename(columns={
"text": "Entity",
"label": "Type",
"start": "Start Position",
"end": "End Position"
})
else:
entities_df = pd.DataFrame()
if relations:
relations_df = pd.DataFrame(relations)
relations_df = relations_df.rename(columns={
"subject": "Subject",
"relation": "Relation",
"object": "Object"
})
else:
relations_df = pd.DataFrame()
# Visualization
if interactive and entities and relations:
try:
interactive_html = visualize_knowledge_graph_interactive(entities, relations)
visualization_html = f"<div style='width:100%;overflow-x:auto'>{interactive_html}</div>"
viz_vis = True
no_viz_vis = False
except Exception as e:
visualization_html = f"<div style='color:#d32f2f;padding:20px;'>Error rendering interactive graph: {e}</div>"
viz_vis = True
no_viz_vis = False
elif visualization:
visualization_html = f"<img src='data:image/png;base64,{visualization}' style='max-width:100%;height:auto;'/>"
viz_vis = True
no_viz_vis = False
else:
visualization_html = ""
viz_vis = False
no_viz_vis = True
# Visibility flags
entities_vis = not entities_df.empty
no_entities_vis = not entities_vis
relations_vis = not relations_df.empty
no_relations_vis = not relations_vis
# Return
return (
visualization_html,
entities_df,
relations_df,
viz_vis,
no_viz_vis,
entities_vis,
no_entities_vis,
relations_vis,
no_relations_vis
)
# UI Components
with gr.Row():
with gr.Column(scale=2):
input_text = gr.Textbox(
label="Input Text",
lines=8,
placeholder="Enter text to extract knowledge graph...",
elem_id="kg-input-text"
)
gr.Examples(
examples=[
["Elon Musk founded SpaceX and Tesla in the United States."],
["Amazon acquired Whole Foods in 2017."]
],
inputs=[input_text],
label="Examples"
)
# Model selection
model = gr.Dropdown(
KG_MODELS,
value=DEFAULT_MODEL,
label="Model",
interactive=True,
elem_id="kg-model-dropdown"
)
with gr.Accordion("Advanced Options", open=False, elem_id="kg-advanced-options"):
custom_instructions = gr.Textbox(
label="Custom Instructions",
lines=2,
placeholder="(Optional) Add specific instructions for knowledge graph generation...",
elem_id="kg-custom-instructions"
)
btn = gr.Button("Generate Knowledge Graph", elem_id="kg-btn")
with gr.Column(scale=3):
# Results container with tabs
with gr.Tabs() as output_tabs:
with gr.Tab("Graph Visualization", id="kg-viz-tab"):
no_viz_html = gr.HTML(
"<div style='text-align: center; color: #666; padding: 20px;'>"
"Generate a knowledge graph to visualize relationships.</div>",
visible=True,
elem_id="kg-no-viz"
)
viz_html = gr.HTML(
label="Knowledge Graph Visualization",
visible=False,
elem_id="kg-viz-html"
)
with gr.Tab("Entities", id="kg-entities-tab"):
no_entities_html = gr.HTML(
"<div style='text-align: center; color: #666; padding: 20px;'>"
"No entities found. Try generating a knowledge graph first.</div>",
visible=True,
elem_id="kg-no-entities"
)
entities_table = gr.DataFrame(
headers=["Entity", "Type", "Start Position", "End Position"],
datatype=["str", "str", "number", "number"],
visible=False,
elem_id="kg-entities-table"
)
with gr.Tab("Relationships", id="kg-relations-tab"):
no_relations_html = gr.HTML(
"<div style='text-align: center; color: #666; padding: 20px;'>"
"No relationships found. Try generating a knowledge graph first.</div>",
visible=True,
elem_id="kg-no-relations"
)
relations_table = gr.DataFrame(
headers=["Subject", "Relation", "Object"],
datatype=["str", "str", "str"],
visible=False,
elem_id="kg-relations-table"
)
with gr.Accordion("About Knowledge Graphs", open=False):
gr.Markdown("""
## Knowledge Graphs
Knowledge graphs represent relationships between entities in text as a network. This tool:
- **Extracts entities**: Identifies people, places, organizations, and concepts
- **Maps relationships**: Shows how entities are connected to each other
- **Visualizes connections**: Creates an interactive graph you can explore
### How it works
- **LLM models** can understand complex relationships in text
- **Traditional models** use pattern matching and syntactic parsing
Knowledge graphs are particularly useful for:
- Research and analysis
- Content exploration
- Understanding complex narratives
- Discovering hidden connections
Try it with news articles, scientific papers, or story excerpts to see different types of relationships.
""")
# Toggle for interactive/static visualization
with gr.Row():
interactive_toggle = gr.Checkbox(
label="Interactive Graph (pyvis)",
value=True,
elem_id="kg-interactive-toggle"
)
# Event handler: use build_kg for all outputs
def process_and_update_ui(text, model, custom_instructions, interactive):
return build_kg(text, model, custom_instructions, interactive)
# Wire button to unified handler
def gradio_output_adapter(visualization_html, entities_df, relations_df, viz_vis, no_viz_vis, entities_vis, no_entities_vis, relations_vis, no_relations_vis):
return [
gr.update(value=visualization_html, visible=viz_vis),
gr.update(value=entities_df, visible=entities_vis),
gr.update(value=relations_df, visible=relations_vis),
gr.update(visible=no_viz_vis),
gr.update(visible=no_entities_vis),
gr.update(visible=no_relations_vis),
]
btn.click(
fn=lambda text, model, custom_instructions, interactive: gradio_output_adapter(*build_kg(text, model, custom_instructions, interactive)),
inputs=[input_text, model, custom_instructions, interactive_toggle],
outputs=[
viz_html, entities_table, relations_table,
no_viz_html, no_entities_html, no_relations_html,
]
)
return None
|