Spaces:
Running
on
Zero
Running
on
Zero
| import spaces | |
| import gradio as gr | |
| from phi3_instruct_graph import MODEL_LIST, Phi3InstructGraph | |
| import rapidjson | |
| from pyvis.network import Network | |
| import networkx as nx | |
| import spacy | |
| from spacy import displacy | |
| from spacy.tokens import Span | |
| import random | |
| from tqdm import tqdm | |
| # Constants | |
| TITLE = "π GraphMind: Phi-3 Instruct Graph Explorer" | |
| SUBTITLE = "β¨ Extract and visualize knowledge graphs from any text in multiple languages" | |
| # Custom CSS for styling | |
| CUSTOM_CSS = """ | |
| .gradio-container { | |
| font-family: 'Inter', 'Segoe UI', Roboto, sans-serif; | |
| } | |
| .gr-button-primary { | |
| background-color: #6366f1 !important; | |
| } | |
| .gr-button-secondary { | |
| border-color: #6366f1 !important; | |
| color: #6366f1 !important; | |
| } | |
| """ | |
| # Color utilities | |
| def get_random_light_color(): | |
| r = random.randint(140, 255) | |
| g = random.randint(140, 255) | |
| b = random.randint(140, 255) | |
| return f"#{r:02x}{g:02x}{b:02x}" | |
| # Text preprocessing | |
| def handle_text(text): | |
| return " ".join(text.split()) | |
| # Main processing functions | |
| def extract(text, model): | |
| try: | |
| model = Phi3InstructGraph(model=model) | |
| result = model.extract(text) | |
| return rapidjson.loads(result) | |
| except Exception as e: | |
| raise gr.Error(f"Extraction error: {str(e)}") | |
| def find_token_indices(doc, substring, text): | |
| result = [] | |
| start_index = text.find(substring) | |
| while start_index != -1: | |
| end_index = start_index + len(substring) | |
| start_token = None | |
| end_token = None | |
| for token in doc: | |
| if token.idx == start_index: | |
| start_token = token.i | |
| if token.idx + len(token) == end_index: | |
| end_token = token.i + 1 | |
| if start_token is not None and end_token is not None: | |
| result.append({ | |
| "start": start_token, | |
| "end": end_token | |
| }) | |
| # Search for next occurrence | |
| start_index = text.find(substring, end_index) | |
| return result | |
| def create_custom_entity_viz(data, full_text): | |
| nlp = spacy.blank("xx") | |
| doc = nlp(full_text) | |
| spans = [] | |
| colors = {} | |
| for node in data["nodes"]: | |
| entity_spans = find_token_indices(doc, node["id"], full_text) | |
| for dataentity in entity_spans: | |
| start = dataentity["start"] | |
| end = dataentity["end"] | |
| if start < len(doc) and end <= len(doc): | |
| # Check for overlapping spans | |
| overlapping = any(s.start < end and start < s.end for s in spans) | |
| if not overlapping: | |
| span = Span(doc, start, end, label=node["type"]) | |
| spans.append(span) | |
| if node["type"] not in colors: | |
| colors[node["type"]] = get_random_light_color() | |
| doc.set_ents(spans, default="unmodified") | |
| doc.spans["sc"] = spans | |
| options = { | |
| "colors": colors, | |
| "ents": list(colors.keys()), | |
| "style": "ent", | |
| "manual": True | |
| } | |
| html = displacy.render(doc, style="span", options=options) | |
| return html | |
| def create_graph(json_data): | |
| G = nx.Graph() | |
| # Add nodes with tooltips | |
| for node in json_data['nodes']: | |
| G.add_node(node['id'], title=f"{node['type']}: {node['detailed_type']}") | |
| # Add edges with labels | |
| for edge in json_data['edges']: | |
| G.add_edge(edge['from'], edge['to'], title=edge['label'], label=edge['label']) | |
| # Create network visualization | |
| nt = Network( | |
| width="720px", | |
| height="600px", | |
| directed=True, | |
| notebook=False, | |
| bgcolor="#f8fafc", | |
| font_color="#1e293b" | |
| ) | |
| # Configure network display | |
| nt.from_nx(G) | |
| nt.barnes_hut( | |
| gravity=-3000, | |
| central_gravity=0.3, | |
| spring_length=50, | |
| spring_strength=0.001, | |
| damping=0.09, | |
| overlap=0, | |
| ) | |
| # Customize edge appearance | |
| for edge in nt.edges: | |
| edge['width'] = 2 | |
| edge['arrows'] = {'to': {'enabled': True, 'type': 'arrow'}} | |
| edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'} | |
| edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Arial'} | |
| # Customize node appearance | |
| for node in nt.nodes: | |
| node['color'] = {'background': '#e0e7ff', 'border': '#6366f1', 'highlight': {'background': '#c7d2fe', 'border': '#4f46e5'}} | |
| node['font'] = {'size': 14, 'color': '#1e293b'} | |
| node['shape'] = 'dot' | |
| node['size'] = 25 | |
| # Generate HTML with iframe to isolate styles | |
| html = nt.generate_html() | |
| html = html.replace("'", '"') | |
| return f"""<iframe style="width: 100%; height: 620px; margin: 0 auto; border-radius: 8px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);" | |
| name="result" allow="midi; geolocation; microphone; camera; display-capture; encrypted-media;" | |
| sandbox="allow-modals allow-forms allow-scripts allow-same-origin allow-popups | |
| allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" | |
| allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>""" | |
| def process_and_visualize(text, model, progress=gr.Progress()): | |
| if not text or not model: | |
| raise gr.Error("β οΈ Both text and model must be provided.") | |
| progress(0, desc="Starting extraction...") | |
| json_data = extract(text, model) | |
| progress(0.5, desc="Creating entity visualization...") | |
| entities_viz = create_custom_entity_viz(json_data, text) | |
| progress(0.8, desc="Building knowledge graph...") | |
| graph_html = create_graph(json_data) | |
| node_count = len(json_data["nodes"]) | |
| edge_count = len(json_data["edges"]) | |
| stats = f"π Extracted {node_count} entities and {edge_count} relationships" | |
| progress(1.0, desc="Complete!") | |
| return graph_html, entities_viz, json_data, stats | |
| # Example texts in different languages | |
| EXAMPLES = [ | |
| [handle_text("""Legendary rock band Aerosmith has officially announced their retirement from touring after 54 years, citing | |
| lead singer Steven Tyler's unrecoverable vocal cord injury. | |
| The decision comes after months of unsuccessful treatment for Tyler's fractured larynx, | |
| which he suffered in September 2023.""")], | |
| [handle_text("""Pop star Justin Timberlake, 43, had his driver's license suspended by a New York judge during a virtual | |
| court hearing on August 2, 2024. The suspension follows Timberlake's arrest for driving while intoxicated (DWI) | |
| in Sag Harbor on June 18. Timberlake, who is currently on tour in Europe, | |
| pleaded not guilty to the charges.""")], | |
| [handle_text("""μΈκ³μ μΈ κΈ°μ κΈ°μ μΌμ±μ μλ μλ‘μ΄ μΈκ³΅μ§λ₯ κΈ°λ° μ€λ§νΈν°μ μ¬ν΄ νλ°κΈ°μ μΆμν μμ μ΄λΌκ³ λ°ννλ€. | |
| μ΄ μ€λ§νΈν°μ νμ¬ κ°λ° μ€μΈ κ°€λμ μ리μ¦μ μ΅μ μμΌλ‘, κ°λ ₯ν AI κΈ°λ₯κ³Ό νμ μ μΈ μΉ΄λ©λΌ μμ€ν μ νμ¬ν κ²μΌλ‘ μλ €μ‘λ€. | |
| μΌμ±μ μμ CEOλ μ΄λ² μ μ νμ΄ μ€λ§νΈν° μμ₯μ μλ‘μ΄ νμ μ κ°μ Έμ¬ κ²μ΄λΌκ³ μ λ§νλ€.""")], | |
| [handle_text("""νκ΅ μν 'κΈ°μμΆ©'μ 2020λ μμΉ΄λ°λ―Έ μμμμμ μνμ, κ°λ μ, κ°λ³Έμ, κ΅μ μνμ λ± 4κ° λΆλ¬Έμ μμνλ©° μμ¬λ₯Ό μλ‘ μΌλ€. | |
| λ΄μ€νΈ κ°λ μ΄ μ°μΆν μ΄ μνλ νκ΅ μν μ΅μ΄λ‘ μΉΈ μνμ ν©κΈμ’ λ €μλ μμνμΌλ©°, μ μΈκ³μ μΌλ‘ μμ²λ ν₯νκ³Ό | |
| νλ¨μ νΈνμ λ°μλ€.""")] | |
| ] | |
| def create_ui(): | |
| with gr.Blocks(css=CUSTOM_CSS, title=TITLE) as demo: | |
| # Header | |
| gr.Markdown(f"# {TITLE}") | |
| gr.Markdown(f"{SUBTITLE}") | |
| with gr.Row(): | |
| gr.Markdown("π **Multilingual Support Available** π€") | |
| # Main interface | |
| with gr.Row(): | |
| # Input column | |
| with gr.Column(scale=1): | |
| input_model = gr.Dropdown( | |
| MODEL_LIST, | |
| label="π€ Select Model", | |
| info="Choose a model to process your text", | |
| value=MODEL_LIST[0] if MODEL_LIST else None | |
| ) | |
| input_text = gr.TextArea( | |
| label="π Input Text", | |
| info="Enter text in any language to extract a knowledge graph", | |
| placeholder="Enter text here...", | |
| lines=10 | |
| ) | |
| with gr.Row(): | |
| submit_button = gr.Button("π Extract & Visualize", variant="primary", scale=2) | |
| clear_button = gr.Button("π Clear", variant="secondary", scale=1) | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=input_text, | |
| label="π Example Texts (English & Korean)" | |
| ) | |
| stats_output = gr.Markdown("", label="π Analysis Results") | |
| # Output column | |
| with gr.Column(scale=1): | |
| with gr.Tab("π§© Knowledge Graph"): | |
| output_graph = gr.HTML(label="") | |
| with gr.Tab("π·οΈ Entities"): | |
| output_entity_viz = gr.HTML(label="") | |
| with gr.Tab("π JSON Data"): | |
| output_json = gr.JSON(label="") | |
| # Functionality | |
| submit_button.click( | |
| fn=process_and_visualize, | |
| inputs=[input_text, input_model], | |
| outputs=[output_graph, output_entity_viz, output_json, stats_output] | |
| ) | |
| clear_button.click( | |
| fn=lambda: [None, None, None, ""], | |
| inputs=[], | |
| outputs=[output_graph, output_entity_viz, output_json, stats_output] | |
| ) | |
| # Footer | |
| gr.Markdown("---") | |
| gr.Markdown("π **Instructions:** Enter text in any language, select a model, and click 'Extract & Visualize' to generate a knowledge graph.") | |
| gr.Markdown("π οΈ Powered by Phi-3 Instruct Graph | Emergent Methods") | |
| return demo | |
| demo = create_ui() | |
| demo.launch(share=False) |