Update main.py

main.py

@@ -1,39 +1,79 @@
 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
+import time
+
+# Set up the theme and styling
+CUSTOM_CSS = """
+.gradio-container {
+    font-family: 'Inter', 'Segoe UI', Roboto, sans-serif;
+}
+.gr-prose h1 {
+    font-size: 2.5rem !important;
+    margin-bottom: 0.5rem !important;
+    background: linear-gradient(90deg, #4338ca, #a855f7);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+}
+.gr-prose h2 {
+    font-size: 1.8rem !important;
+    margin-top: 1rem !important;
+}
+.info-box {
+    padding: 1rem;
+    border-radius: 0.5rem;
+    background-color: #f3f4f6;
+    margin-bottom: 1rem;
+    border-left: 4px solid #6366f1;
+}
+.language-badge {
+    display: inline-block;
+    padding: 0.25rem 0.5rem;
+    border-radius: 9999px;
+    font-size: 0.75rem;
+    font-weight: 600;
+    background-color: #e0e7ff;
+    color: #4338ca;
+    margin-right: 0.5rem;
+    margin-bottom: 0.5rem;
+}
+.footer {
+    text-align: center;
+    margin-top: 2rem;
+    padding-top: 1rem;
+    border-top: 1px solid #e2e8f0;
+    font-size: 0.875rem;
+    color: #64748b;
+}
+"""
+
+# Color utilities
+def get_random_light_color():
+    r = random.randint(150, 255)
+    g = random.randint(150, 255)
+    b = random.randint(150, 255)
+    return f"#{r:02x}{g:02x}{b:02x}"
 
-json_example = {'nodes': [{'id': 'Aerosmith', 'type': 'organization', 'detailed_type': 'rock band'}, {'id': 'Steven Tyler', 'type': 'person', 'detailed_type': 'lead singer'}, {'id': 'vocal cord injury', 'type': 'medical condition', 'detailed_type': 'fractured larynx'}, {'id': 'retirement', 'type': 'event', 'detailed_type': 'announcement'}, {'id': 'touring', 'type': 'activity', 'detailed_type': 'musical performance'}, {'id': 'September 2023', 'type': 'date', 'detailed_type': 'specific time'}], 'edges': [{'from': 'Aerosmith', 'to': 'Steven Tyler', 'label': 'led by'}, {'from': 'Steven Tyler', 'to': 'vocal cord injury', 'label': 'suffered'}, {'from': 'vocal cord injury', 'to': 'retirement', 'label': 'caused'}, {'from': 'retirement', 'to': 'touring', 'label': 'ended'}, {'from': 'vocal cord injury', 'to': 'September 2023', 'label': 'occurred in'}]}
+# Text processing helper
+def handle_text(text):
+    return " ".join(text.split())
 
+# Core extraction function
 @spaces.GPU
 def extract(text, model):
     model = Phi3InstructGraph(model=model)    
-    result = model.extract(text)
-    return rapidjson.loads(result)
-
-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}"
+    try:
+        result = model.extract(text)
+        return rapidjson.loads(result)
+    except Exception as e:
+        raise gr.Error(f"🚨 Extraction failed: {str(e)}")
 
 def find_token_indices(doc, substring, text):
     result = []
@@ -50,9 +90,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
@@ -61,36 +99,29 @@ 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)
 
     spans = []
     colors = {}
+    
     for node in data["nodes"]:
-    #   entity_spans = [m.span() for m in re.finditer(re.escape(node["id"]), full_text)]
         entity_spans = find_token_indices(doc, node["id"], full_text)
         for dataentity in entity_spans:
             start = dataentity["start"]
             end = dataentity["end"]
             
-            print("entity spans:", entity_spans)
             if start < len(doc) and end <= len(doc):
-                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()
-
-    for span in spans:
-        print(f"Span: {span.text}, Label: {span.label_}")
+                # 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
@@ -103,108 +134,267 @@ def create_custom_entity_viz(data, full_text):
     }
 
     html = displacy.render(doc, style="span", options=options)
-    return html
-
+    
+    # Add custom styling to the entity visualization
+    styled_html = f"""
+    <div style="border-radius: 0.5rem; padding: 1rem; background-color: white; 
+                border: 1px solid #e2e8f0; box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.1);">
+        <div style="margin-bottom: 0.75rem; font-weight: 500; color: #4b5563;">
+            Entity types found:
+            {' '.join([f'<span style="display: inline-block; margin-right: 0.5rem; margin-bottom: 0.5rem; padding: 0.25rem 0.5rem; border-radius: 9999px; font-size: 0.75rem; background-color: {colors[entity_type]}; color: #1e293b;">{entity_type}</span>' for entity_type in colors.keys()])}
+        </div>
+        {html}
+    </div>
+    """
+    
+    return styled_html
 
 def create_graph(json_data):
-    G = nx.Graph()
+    G = nx.DiGraph()  # Using DiGraph for directed graph
 
+    # Add nodes
     for node in json_data['nodes']:
-        G.add_node(node['id'], title=f"{node['type']}: {node['detailed_type']}")
+        G.add_node(node['id'], 
+                  title=f"{node['type']}: {node['detailed_type']}",
+                  group=node['type'])  # Group nodes by type
 
+    # Add edges
     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",
+        width="100%",
         height="600px",
         directed=True,
         notebook=False,
-        # bgcolor="#111827", 
-        # font_color="white"
-        bgcolor="#FFFFFF",
-        font_color="#111827"
+        bgcolor="#fafafa", 
+        font_color="#1e293b"
     )
+    
+    # Configure network
     nt.from_nx(G)
     nt.barnes_hut(
         gravity=-3000,
         central_gravity=0.3,
-        spring_length=50,
+        spring_length=150,
         spring_strength=0.001,
         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}
-
+    # Create color groups for node types
+    node_types = {node['type'] for node in json_data['nodes']}
+    colors = {}
+    for i, node_type in enumerate(node_types):
+        hue = (i * 137) % 360  # Golden ratio to distribute colors
+        colors[node_type] = f"hsl({hue}, 70%, 70%)"
+
+    # Customize nodes
+    for node in nt.nodes:
+        node_data = next((n for n in json_data['nodes'] if n['id'] == node['id']), None)
+        if node_data:
+            node_type = node_data['type']
+            node['color'] = colors.get(node_type, "#bfdbfe")
+            node['shape'] = 'dot'
+            node['size'] = 20
+            node['borderWidth'] = 2
+            node['borderWidthSelected'] = 4
+            node['font'] = {'size': 14, 'color': '#1e293b', 'face': 'Inter, Arial'}
+    
+    # Customize edges
+    for edge in nt.edges:
+        edge['color'] = {'color': '#94a3b8', 'highlight': '#6366f1', 'hover': '#818cf8'}
+        edge['width'] = 1.5
+        edge['selectionWidth'] = 2
+        edge['hoverWidth'] = 2
+        edge['arrows'] = {'to': {'enabled': True, 'type': 'arrow'}}
+        edge['smooth'] = {'type': 'continuous', 'roundness': 0.2}
+        edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Inter, Arial', 'strokeWidth': 2, 'strokeColor': '#ffffff'}
+
+    # Generate HTML
     html = nt.generate_html()
-    # need to remove ' from HTML
     html = html.replace("'", '"')
-    # return html
+    html = html.replace('height: 600px;', 'height: 600px; border-radius: 8px;')
 
-    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("⚠️ Please provide both text and model")
+    
+    # Progress updates
+    progress(0.1, "Initializing...")
+    time.sleep(0.2)  # Small delay for UI feedback
+    
+    # Extract graph
+    progress(0.2, "Extracting knowledge graph...")
     json_data = extract(text, model)
-    # json_data = json_example
-    print(json_data)
+    
+    # Entity visualization
+    progress(0.6, "Identifying entities...")
     entities_viz = create_custom_entity_viz(json_data, text)
     
+    # Graph visualization
+    progress(0.8, "Building graph visualization...")
     graph_html = create_graph(json_data)
-    return graph_html, entities_viz, json_data
-
-
-
-with gr.Blocks(title="Phi-3 Mini 4k Instruct Graph (by Emergent Methods") as demo:
-    gr.Markdown("# Phi-3 Mini 4k Instruct Graph (by Emergent Methods)")
-    gr.Markdown("Extract a JSON graph from a text input and visualize it.")
-
+    
+    # Statistics
+    entity_types = {}
+    for node in json_data['nodes']:
+        entity_type = node['type']
+        if entity_type in entity_types:
+            entity_types[entity_type] += 1
+        else:
+            entity_types[entity_type] = 1
+    
+    stats_html = f"""
+    <div class="info-box">
+        <h3 style="margin-top: 0;">📊 Extraction Results</h3>
+        <p>✅ Successfully extracted <b>{len(json_data['nodes'])}</b> entities and <b>{len(json_data['edges'])}</b> relationships.</p>
+        
+        <div>
+            <h4>Entity Types:</h4>
+            <div>
+                {''.join([f'<span class="language-badge">{entity_type}: {count}</span>' for entity_type, count in entity_types.items()])}
+            </div>
+        </div>
+    </div>
+    """
+    
+    progress(1.0, "Done!")
+    return graph_html, entities_viz, json_data, stats_html
+
+def language_info():
+    return """
+    <div class="info-box">
+        <h3 style="margin-top: 0;">🌍 Multilingual Support</h3>
+        <p>This application supports text analysis in multiple languages, including:</p>
+        <div>
+            <span class="language-badge">English 🇬🇧</span>
+            <span class="language-badge">Korean 🇰🇷</span>
+            <span class="language-badge">Spanish 🇪🇸</span>
+            <span class="language-badge">French 🇫🇷</span>
+            <span class="language-badge">German 🇩🇪</span>
+            <span class="language-badge">Japanese 🇯🇵</span>
+            <span class="language-badge">Chinese 🇨🇳</span>
+            <span class="language-badge">And more...</span>
+        </div>
+    </div>
+    """
+
+def tips_html():
+    return """
+    <div class="info-box">
+        <h3 style="margin-top: 0;">💡 Tips for Best Results</h3>
+        <ul>
+            <li>Use clear, descriptive sentences with well-defined relationships</li>
+            <li>Include specific entities, events, dates, and locations for better extraction</li>
+            <li>Longer texts provide more context for relationship identification</li>
+            <li>Try different models to compare extraction results</li>
+        </ul>
+    </div>
+    """
+
+# Examples in multiple 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개 부문을 수상하며 역사를 새로 썼다. 
+    봉준호 감독이 연출한 이 영화는 한국 영화 최초로 칸 영화제 황금종려상도 수상했으며, 전 세계적으로 엄청난 흥행과 
+    평단의 호평을 받았다.""")]
+]
+
+# Main UI
+with gr.Blocks(css=CUSTOM_CSS, title="🧠 Phi-3 Knowledge Graph Explorer") as demo:
+    # Header
+    gr.Markdown("# 🧠 Phi-3 Knowledge Graph Explorer")
+    gr.Markdown("### ✨ Extract and visualize knowledge graphs from text in any language")
+    
     with gr.Row():
-        with gr.Column(scale=1):
-            input_model = gr.Dropdown(
-                MODEL_LIST, label="Model",
-                # value=MODEL_LIST[0]
+        with gr.Column(scale=2):
+            input_text = gr.TextArea(
+                label="📝 Enter your text",
+                placeholder="Paste or type your text here...",
+                lines=10
             )
-            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
+            
+            with gr.Row():
+                input_model = gr.Dropdown(
+                    MODEL_LIST,
+                    label="🤖 Model",
+                    value=MODEL_LIST[0] if MODEL_LIST else None,
+                    info="Select the model to use for extraction"
+                )
+                
+                with gr.Column():
+                    submit_button = gr.Button("🔍 Extract & Visualize", variant="primary")
+                    clear_button = gr.Button("🔄 Clear", variant="secondary")
+            
+            # Multilingual support info
+            gr.HTML(language_info())
+            
+            # Examples section
+            gr.Examples(
+                examples=EXAMPLES,
+                inputs=input_text,
+                label="📚 Example Texts (English & Korean)"
             )
-
-            submit_button = gr.Button("Extract and Visualize")
+            
+            # Tips
+            gr.HTML(tips_html())
+            
+        with gr.Column(scale=3):
+            # Stats output
+            stats_output = gr.HTML(label="")
+            
+            # Tabs for different visualizations
+            with gr.Tabs():
+                with gr.TabItem("🔄 Knowledge Graph"):
+                    output_graph = gr.HTML()
+                
+                with gr.TabItem("🏷️ Entity Recognition"):
+                    output_entity_viz = gr.HTML()
+                
+                with gr.TabItem("📊 JSON Data"):
+                    output_json = gr.JSON()
     
-        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)            
-            # output_json = gr.JSON(label="JSON Graph")
-
-            submit_button.click(
-                fn=process_and_visualize, 
-                inputs=[input_text, input_model],
-                outputs=[output_graph, output_entity_viz]
-            )
+    # Footer
+    gr.HTML("""
+    <div class="footer">
+        <p>🌐 Powered by Phi-3 Instruct Graph | Created by Emergent Methods</p>
+        <p>© 2025 | Knowledge Graph Explorer</p>
+    </div>
+    """)
+    
+    # Set up event handlers
+    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]
+    )
 
+# Launch the app
 demo.launch(share=False)