Update app.py

app.py

@@ -1,37 +1,46 @@
 import spaces
 import gradio as gr
 from phi3_instruct_graph import MODEL_LIST, Phi3InstructGraph
-from textwrap import dedent
 import rapidjson
-import spaces
 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 = "🌐 Phi-3 Instruct Graph Explorer"
+SUBTITLE = "✨ Extract and visualize knowledge graphs from any text in multiple languages"
+THEME = gr.themes.Base().set(
+    primary_hue="indigo",
+    secondary_hue="purple",
+    neutral_hue="slate",
+    radius_size=gr.themes.sizes.radius_sm,
+    shadow_size=gr.themes.sizes.shadow_lg,
+)
+
+# 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}"
 
-@spaces.GPU
-def extract(text, model):
-    model = Phi3InstructGraph(model=model)    
-    result = model.extract(text)
-    return rapidjson.loads(result)
-
+# Text preprocessing
 def handle_text(text):
     return " ".join(text.split())
 
-def get_random_color():
-  return f"#{random.randint(0, 0xFFFFFF):06x}"
-
-def get_random_light_color():
-    # Generate higher RGB values to ensure a lighter color
-    r = random.randint(128, 255)
-    g = random.randint(128, 255)
-    b = random.randint(128, 255)
-    return f"#{r:02x}{g:02x}{b:02x}"
-
-def get_random_color():
-    return f"#{random.randint(0, 0xFFFFFF):06x}"
+# Main processing functions
+@spaces.GPU
+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 = []
@@ -48,9 +57,7 @@ def find_token_indices(doc, substring, text):
             if token.idx + len(token) == end_index:
                 end_token = token.i + 1
 
-        if start_token is None or end_token is None:
-            print(f"Token boundaries not found for '{substring}' at index {start_index}")
-        else:
+        if start_token is not None and end_token is not None:
             result.append({
                 "start": start_token,
                 "end": end_token
@@ -59,12 +66,8 @@ def find_token_indices(doc, substring, text):
         # Search for next occurrence
         start_index = text.find(substring, end_index)
 
-    if not result:
-        print(f"Token boundaries not found for '{substring}'")
-
     return result
 
-
 def create_custom_entity_viz(data, full_text):
     nlp = spacy.blank("xx")
     doc = nlp(full_text)
@@ -82,8 +85,6 @@ def create_custom_entity_viz(data, full_text):
                 overlapping = any(s.start < end and start < s.end for s in spans)
                 if not overlapping:                
                     span = Span(doc, start, end, label=node["type"])
-                    
-                    # print(span)
                     spans.append(span)
                     if node["type"] not in colors:
                         colors[node["type"]] = get_random_light_color()
@@ -101,26 +102,28 @@ def create_custom_entity_viz(data, full_text):
     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="#111827", 
-        font_color="white"
-        # bgcolor="#FFFFFF",
-        # font_color="#111827"
+        bgcolor="#f8fafc", 
+        font_color="#1e293b"
     )
+    
+    # Configure network display
     nt.from_nx(G)
     nt.barnes_hut(
         gravity=-3000,
@@ -130,71 +133,141 @@ def create_graph(json_data):
         damping=0.09,
         overlap=0,
     )
-
+    
     # Customize edge appearance
-    # for edge in nt.edges:
-    #     edge['font'] = {'size': 12, 'color': '#FFD700', 'face': 'Arial'}  # Removed strokeWidth
-    #     edge['color'] = {'color': '#FF4500', 'highlight': '#FF4500'}
-    #     edge['width'] = 1
-    #     edge['arrows'] = {'to': {'enabled': True, 'type': 'arrow'}}
-        # edge['smooth'] = {'type': 'curvedCW', 'roundness': 0.2}
-
+    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: 140%; height: 620px; margin: 0 auto;" name="result" 
-        allow="midi; geolocation; microphone; camera; display-capture; encrypted-media;" 
+    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):
+def process_and_visualize(text, model, progress=gr.Progress()):
     if not text or not model:
-        raise gr.Error("Text and model must be provided.")
+        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)
-    return graph_html, entities_viz, json_data
-
-
-
-with gr.Blocks(title="Phi-3 Instruct Graph (by Emergent Methods") as demo:
-    gr.Markdown("# Phi-3 Instruct Graph (by Emergent Methods)")
-    gr.Markdown("Extract a JSON graph from a text input and visualize it.")
-    with gr.Row():
-        with gr.Column(scale=1):
-            input_model = gr.Dropdown(
-                MODEL_LIST, label="Model",
-            )
-            input_text = gr.TextArea(label="Text", info="The text to be extracted")
-
-            examples = gr.Examples(
-                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."""),
-                ],
-                inputs=input_text
-            )
-
-            submit_button = gr.Button("Extract and Visualize")
     
-        with gr.Column(scale=1):
-            output_entity_viz = gr.HTML(label="Entities Visualization", show_label=True)
-            output_graph = gr.HTML(label="Graph Visualization", show_label=True)            
-
-            submit_button.click(
-                fn=process_and_visualize, 
-                inputs=[input_text, input_model],
-                outputs=[output_graph, output_entity_viz]
-            )
+    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(theme=THEME, 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 | Created by Emergent Methods")
+        
+    return demo
 
+demo = create_ui()
 demo.launch(share=False)