refactor (system): knowledge graph added

README.md

@@ -1,12 +1,12 @@
 ---
-title: Clinical Ner
-emoji: 🏆
+title: Clinical NER / Knowledge Graph
+emoji: 🕸️
 colorFrom: green
 colorTo: green
 sdk: docker
 pinned: false
 license: gpl-3.0
-short_description: Clinical NER
+short_description: Clinical NER / Knowledge Graph
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app/clinical_ner.py

@@ -215,3 +215,105 @@ class ClinicalNERProcessor:
         
         return "\n\n".join(sections)
 
+    def generate_knowledge_graph(self, text, patient_id="patient_1"):
+        """
+        Generates a knowledge graph combining Clinical NER and Anatomy NER.
+        
+        Args:
+            text (str): Input text to analyze
+            patient_id (str): Identifier for the patient node
+            
+        Returns:
+            dict: Dictionary with 'nodes' and 'edges' lists
+        """
+        clinical_entities = self.basic_ner(text)
+        anatomy_entities = self.anatomy_ner(text)
+        
+        nodes = []
+        edges = []
+        node_id = 0
+        
+        # Add patient node
+        patient_node_id = node_id
+        nodes.append({
+            'id': node_id,
+            'label': patient_id,
+            'type': 'patient'
+        })
+        node_id += 1
+        
+        # Add entity type nodes
+        entity_types = ['problem', 'treatment', 'test']
+        type_node_ids = {}
+        
+        for entity_type in entity_types:
+            type_node_ids[entity_type] = node_id
+            nodes.append({
+                'id': node_id,
+                'label': entity_type.capitalize(),
+                'type': 'entity_type'
+            })
+            node_id += 1
+        
+        # Add clinical entity nodes
+        for entity in clinical_entities:
+            entity_type = entity['entity_group'].lower()
+            entity_node_id = node_id
+            
+            nodes.append({
+                'id': node_id,
+                'label': entity['word'],
+                'type': entity_type,
+                'score': entity['score']
+            })
+            
+            # Connect entity to patient
+            edges.append({
+                'source': patient_node_id,
+                'target': entity_node_id,
+                'label': 'has'
+            })
+            
+            # Connect entity to entity type
+            if entity_type in type_node_ids:
+                edges.append({
+                    'source': entity_node_id,
+                    'target': type_node_ids[entity_type],
+                    'label': 'is_a'
+                })
+            
+            node_id += 1
+        
+        # Add Anatomy node
+        anatomy_node_id = node_id
+        nodes.append({
+            'id': node_id,
+            'label': 'Anatomy',
+            'type': 'anatomy_category'
+        })
+        node_id += 1
+        
+        # Add anatomy entity nodes
+        for entity in anatomy_entities:
+            entity_node_id = node_id
+            
+            nodes.append({
+                'id': node_id,
+                'label': entity['word'],
+                'type': 'anatomy',
+                'score': entity['score']
+            })
+            
+            # Connect anatomy entity to Anatomy node
+            edges.append({
+                'source': entity_node_id,
+                'target': anatomy_node_id,
+                'label': 'is_a'
+            })
+            
+            node_id += 1
+        
+        return {
+            'nodes': nodes,
+            'edges': edges
+        }

app/server_clinical_ner.py

@@ -93,6 +93,23 @@ class PrologCombinedResponse(BaseModel):
     anatomy_count: int
     token_count: int
 
+class KnowledgeGraphNode(BaseModel):
+    id: int
+    label: str
+    type: str
+    score: float | None = None
+
+class KnowledgeGraphEdge(BaseModel):
+    source: int
+    target: int
+    label: str
+
+class KnowledgeGraphResponse(BaseModel):
+    nodes: list[KnowledgeGraphNode]
+    edges: list[KnowledgeGraphEdge]
+    node_count: int
+    edge_count: int
+
 @app.get("/")
 async def root():
     return RedirectResponse(url="/browser/")
@@ -285,6 +302,33 @@ async def health_check():
         "anatomy_available": ner_model.anatomy_pipeline is not None if ner_model else False
     }
 
+@app.post("/knowledge-graph", response_model=KnowledgeGraphResponse)
+async def generate_knowledge_graph(request: TextRequest):
+    """
+    Generate a knowledge graph combining Clinical NER and Anatomy NER.
+    
+    Returns nodes and edges representing:
+    - Patient node connected to clinical entities
+    - Entity type nodes (Problem, Treatment, Test)
+    - Clinical entities connected to their types
+    - Anatomy category node
+    - Anatomical entities connected to Anatomy node
+    """
+    try:
+        if not request.text.strip():
+            raise HTTPException(status_code=400, detail="Text cannot be empty")
+        
+        graph = ner_model.generate_knowledge_graph(request.text)
+        
+        return {
+            "nodes": graph['nodes'],
+            "edges": graph['edges'],
+            "node_count": len(graph['nodes']),
+            "edge_count": len(graph['edges'])
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing text: {str(e)}")
+
 @app.get("/models/info")
 async def models_info():
     """Get information about loaded models"""

app/static/browser/index.html

@@ -474,6 +474,7 @@
                     <button id="anatomyBtn" class="btn btn-anatomy">🫀 Anatomy NER</button>
                     <button id="posBtn" class="btn btn-secondary">📝 POS Tagging</button>
                     <button id="combinedBtn" class="btn btn-success">🎯 Combined Analysis</button>
+                    <button id="graphBtn" class="btn" style="background: linear-gradient(135deg, #a8edea 0%, #fed6e3 100%); color: #333;">🕸️ Knowledge Graph</button>
                 </div>
             </div>
 
@@ -490,6 +491,7 @@
                     <button class="tab" data-tab="anatomy">Anatomy NER/Visualization</button>
                     <button class="tab" data-tab="pos">POS Results</button>
                     <button class="tab" data-tab="combined">Combined</button>
+                    <button class="tab" data-tab="graph">Knowledge Graph</button>
                 </div>
 
                 <div id="ner" class="tab-content active">
@@ -606,6 +608,44 @@
                         <pre id="prologFactsCombined" class="prolog-output"></pre>
                     </div>
                 </div>
+
+                <div id="graph" class="tab-content">
+                    <div class="result-section">
+                        <h2>🕸️ Knowledge Graph Visualization</h2>
+                        <div id="graph-container"></div>
+                    </div>
+
+                    <div class="graph-tables">
+                        <div class="graph-table">
+                            <h3>📊 Nodes Table</h3>
+                            <table id="nodesTable">
+                                <thead>
+                                    <tr>
+                                        <th>ID</th>
+                                        <th>Label</th>
+                                        <th>Type</th>
+                                        <th>Score</th>
+                                    </tr>
+                                </thead>
+                                <tbody></tbody>
+                            </table>
+                        </div>
+
+                        <div class="graph-table">
+                            <h3>🔗 Edges Table</h3>
+                            <table id="edgesTable">
+                                <thead>
+                                    <tr>
+                                        <th>Source</th>
+                                        <th>Target</th>
+                                        <th>Relationship</th>
+                                    </tr>
+                                </thead>
+                                <tbody></tbody>
+                            </table>
+                        </div>
+                    </div>
+                </div>
             </div>
         </div>
     </div>
@@ -868,6 +908,132 @@
             }
         });
 
+        // Knowledge Graph button handler
+        graphBtn.addEventListener('click', async () => {
+            const text = clinicalText.value.trim();
+            
+            if (!text) {
+                showError('Please enter some text to analyze.');
+                return;
+            }
+
+            hideError();
+            hideResults();
+            showLoading();
+            graphBtn.disabled = true;
+
+            try {
+                const response = await fetch(`${API_BASE_URL}/knowledge-graph`, {
+                    method: 'POST',
+                    headers: { 'Content-Type': 'application/json' },
+                    body: JSON.stringify({ text })
+                });
+
+                if (!response.ok) {
+                    throw new Error('API request failed');
+                }
+
+                const data = await response.json();
+                displayKnowledgeGraph(data.nodes, data.edges);
+                switchToTab('graph');
+                showResults();
+            } catch (err) {
+                showError(`Error: ${err.message}. Please check if the API service is running.`);
+            } finally {
+                hideLoading();
+                graphBtn.disabled = false;
+            }
+        });
+
+        function displayKnowledgeGraph(nodes, edges) {
+            // Display tables
+            const nodesTableBody = document.querySelector('#nodesTable tbody');
+            nodesTableBody.innerHTML = '';
+            nodes.forEach(node => {
+                const row = nodesTableBody.insertRow();
+                row.innerHTML = `
+                    <td>${node.id}</td>
+                    <td class="node-${node.type}">${node.label}</td>
+                    <td>${node.type}</td>
+                    <td>${node.score ? node.score.toFixed(4) : '-'}</td>
+                `;
+            });
+
+            const edgesTableBody = document.querySelector('#edgesTable tbody');
+            edgesTableBody.innerHTML = '';
+            edges.forEach(edge => {
+                const sourceNode = nodes.find(n => n.id === edge.source);
+                const targetNode = nodes.find(n => n.id === edge.target);
+                const row = edgesTableBody.insertRow();
+                row.innerHTML = `
+                    <td>${sourceNode ? sourceNode.label : edge.source}</td>
+                    <td>${targetNode ? targetNode.label : edge.target}</td>
+                    <td>${edge.label}</td>
+                `;
+            });
+
+            // Visualize graph
+            const container = document.getElementById('graph-container');
+            
+            // Define colors for different node types
+            const nodeColors = {
+                patient: { background: '#e74c3c', border: '#c0392b' },
+                entity_type: { background: '#3498db', border: '#2980b9' },
+                problem: { background: '#ffcccb', border: '#8b0000' },
+                treatment: { background: '#c7f0c7', border: '#006400' },
+                test: { background: '#cce5ff', border: '#004085' },
+                anatomy_category: { background: '#9b59b6', border: '#8e44ad' },
+                anatomy: { background: '#e8daef', border: '#9b59b6' }
+            };
+
+            // Prepare nodes for vis.js
+            const visNodes = nodes.map(node => ({
+                id: node.id,
+                label: node.label,
+                color: nodeColors[node.type] || { background: '#95a5a6', border: '#7f8c8d' },
+                font: { color: '#333', size: 14 },
+                shape: node.type === 'patient' ? 'diamond' : 
+                       (node.type === 'entity_type' || node.type === 'anatomy_category') ? 'box' : 'ellipse'
+            }));
+
+            // Prepare edges for vis.js
+            const visEdges = edges.map(edge => ({
+                from: edge.source,
+                to: edge.target,
+                label: edge.label,
+                arrows: 'to',
+                font: { size: 10, align: 'middle' }
+            }));
+
+            // Create network
+            const data = {
+                nodes: new vis.DataSet(visNodes),
+                edges: new vis.DataSet(visEdges)
+            };
+
+            const options = {
+                layout: {
+                    hierarchical: {
+                        enabled: true,
+                        direction: 'UD',
+                        sortMethod: 'directed',
+                        levelSeparation: 150,
+                        nodeSpacing: 200
+                    }
+                },
+                physics: {
+                    enabled: false
+                },
+                interaction: {
+                    dragNodes: true,
+                    dragView: true,
+                    zoomView: true
+                }
+            };
+
+            new vis.Network(container, data, options);
+        }
+
         function displayNERResults(originalText, entities, prolog) {
             const sortedEntities = [...entities].sort((a, b) => a.start - b.start);