Spaces:
Runtime error
Runtime error
| import random | |
| import gradio as gr | |
| import networkx as nx | |
| from lib.graph_extract import triplextract, parse_triples | |
| from lib.visualize import create_graph, create_bokeh_plot, create_plotly_plot | |
| from lib.samples import snippets | |
| WORD_LIMIT = 300 | |
| def process_text(text, entity_types, predicates, layout_type, visualization_type): | |
| if not text: | |
| return None, None, "Please enter some text." | |
| words = text.split() | |
| if len(words) > WORD_LIMIT: | |
| return None, None, f"Please limit your input to {WORD_LIMIT} words. Current word count: {len(words)}" | |
| entity_types = [et.strip() for et in entity_types.split(",") if et.strip()] | |
| predicates = [p.strip() for p in predicates.split(",") if p.strip()] | |
| if not entity_types: | |
| return None, None, "Please enter at least one entity type." | |
| if not predicates: | |
| return None, None, "Please enter at least one predicate." | |
| try: | |
| prediction = triplextract(text, entity_types, predicates) | |
| if prediction.startswith("Error"): | |
| return None, None, prediction | |
| entities, relationships = parse_triples(prediction) | |
| if not entities and not relationships: | |
| return None, None, "No entities or relationships found. Try different text or check your input." | |
| G = create_graph(entities, relationships) | |
| if visualization_type == 'Bokeh': | |
| fig = create_bokeh_plot(G, layout_type) | |
| else: | |
| fig = create_plotly_plot(G, layout_type) | |
| output_text = f"Entities: {entities}\nRelationships: {relationships}\n\nRaw output:\n{prediction}" | |
| return G, fig, output_text | |
| except Exception as e: | |
| print(f"Error in process_text: {str(e)}") | |
| return None, None, f"An error occurred: {str(e)}" | |
| def update_graph(G, layout_type, visualization_type): | |
| if G is None: | |
| return None, "Please process text first." | |
| try: | |
| if visualization_type == 'Bokeh': | |
| fig = create_bokeh_plot(G, layout_type) | |
| else: | |
| fig = create_plotly_plot(G, layout_type) | |
| return fig, "" | |
| except Exception as e: | |
| print(f"Error in update_graph: {e}") | |
| return None, f"An error occurred while updating the graph: {str(e)}" | |
| def update_inputs(sample_name): | |
| sample = snippets[sample_name] | |
| return sample.text_input, sample.entity_types, sample.predicates | |
| with gr.Blocks(theme=gr.themes.Monochrome()) as demo: | |
| gr.Markdown("# Knowledge Graph Extractor") | |
| default_sample_name = random.choice(list(snippets.keys())) | |
| default_sample = snippets[default_sample_name] | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| sample_dropdown = gr.Dropdown(choices=list(snippets.keys()), label="Select Sample", value=default_sample_name) | |
| input_text = gr.Textbox(label="Input Text", lines=5, value=default_sample.text_input) | |
| entity_types = gr.Textbox(label="Entity Types", value=default_sample.entity_types) | |
| predicates = gr.Textbox(label="Predicates", value=default_sample.predicates) | |
| layout_type = gr.Dropdown(choices=['spring', 'fruchterman_reingold', 'circular', 'random', 'spectral', 'shell'], | |
| label="Layout Type", value='spring') | |
| visualization_type = gr.Radio(choices=['Bokeh', 'Plotly'], label="Visualization Type", value='Bokeh') | |
| process_btn = gr.Button("Process Text") | |
| with gr.Column(scale=2): | |
| output_graph = gr.Plot(label="Knowledge Graph") | |
| error_message = gr.Textbox(label="Textual Output") | |
| graph_state = gr.State(None) | |
| def process_and_update(text, entity_types, predicates, layout_type, visualization_type): | |
| G, fig, output = process_text(text, entity_types, predicates, layout_type, visualization_type) | |
| return G, fig, output | |
| def update_graph_wrapper(G, layout_type, visualization_type): | |
| if G is not None: | |
| fig, _ = update_graph(G, layout_type, visualization_type) | |
| return fig | |
| sample_dropdown.change(update_inputs, inputs=[sample_dropdown], outputs=[input_text, entity_types, predicates]) | |
| process_btn.click(process_and_update, | |
| inputs=[input_text, entity_types, predicates, layout_type, visualization_type], | |
| outputs=[graph_state, output_graph, error_message]) | |
| layout_type.change(update_graph_wrapper, | |
| inputs=[graph_state, layout_type, visualization_type], | |
| outputs=[output_graph]) | |
| visualization_type.change(update_graph_wrapper, | |
| inputs=[graph_state, layout_type, visualization_type], | |
| outputs=[output_graph]) | |
| if __name__ == "__main__": | |
| demo.launch(share=True) |