Spaces:
Running
Running
| import streamlit as st | |
| from streamlit_text_label import label_select, Selection | |
| from streamlit_agraph import agraph, Node, Edge, Config | |
| # Import Relation class | |
| from dataclasses import dataclass | |
| class Relation: | |
| source: Selection | |
| target: Selection | |
| label: str | |
| def mock_entity_recognition(text): | |
| # Simulate entity recognition functionality | |
| entities = [ | |
| Selection(start=0, end=12, text="FINAL REPORT", | |
| labels=["REPORT_TYPE"]), | |
| Selection(start=22, end=33, text="Pseudomonas", labels=["DISEASE"]), | |
| Selection(start=39, end=49, text="intubation", labels=["PROCEDURE"]), | |
| Selection(start=62, end=67, text="study", labels=["EXAM"]), | |
| Selection(start=116, end=126, text="Substantial", labels=["SEVERITY"]), | |
| Selection(start=127, end=145, text="bilateral pleural", | |
| labels=["ANATOMY"]), | |
| Selection(start=146, end=155, text="effusions", labels=["OBSERVATION"]) | |
| ] | |
| return entities | |
| def mock_relation_extraction(entities): | |
| # Simulate relation extraction functionality | |
| relations = [ | |
| Relation(source=entities[0], | |
| target=entities[1], label="DISEASE_CAUSE"), | |
| Relation(source=entities[1], target=entities[2], | |
| label="PROCEDURE_EFFECT"), | |
| Relation(source=entities[2], target=entities[3], label="EXAM_RESULT"), | |
| Relation(source=entities[3], target=entities[4], | |
| label="SEVERITY_LEVEL"), | |
| Relation(source=entities[4], target=entities[5], | |
| label="ANATOMY_LOCATION"), | |
| ] | |
| return relations | |
| def create_graph(entities, relations): | |
| nodes = [Node(id=e.text, label=e.text, size=25, color=f"#{hash(e.labels[0]) % 0xFFFFFF:06x}") for e in entities] | |
| edges = [Edge(source=r.source.text, target=r.target.text, label=r.label) for r in relations] | |
| config = Config(width=750, height=500, directed=True, physics=True, hierarchical=False) | |
| return agraph(nodes=nodes, edges=edges, config=config) | |
| def main(): | |
| st.title("Medical Text Entity Recognition") | |
| # 1. Text input | |
| text = st.text_area("Enter medical text:", value="FINAL REPORT HISTORY : Pseudomonas with intubation . FINDINGS : In comparison with the study of ___ , there is little change in the monitoring and support devices . Substantial bilateral pleural effusions , more prominent on the right with bibasilar atelectasis .") | |
| if st.button("Recognize Entities"): | |
| # 2. Call the simulated entity recognition function | |
| entities = mock_entity_recognition(text) | |
| # 3. Use streamlit_text_label to highlight recognized entities | |
| st.subheader("Recognition Results:") | |
| labeled_text = label_select( | |
| body=text, | |
| labels=["REPORT_TYPE", "DISEASE", "PROCEDURE", | |
| "EXAM", "SEVERITY", "ANATOMY", "OBSERVATION"], | |
| selections=entities | |
| ) | |
| st.write("Recognized entities:") | |
| for entity in entities: | |
| st.write(f"{entity.text} ({entity.labels[0]})") | |
| # 4. Call the simulated relation extraction function | |
| relations = mock_relation_extraction(entities) | |
| # 5. Display relations | |
| st.subheader("Extracted Relations:") | |
| for relation in relations: | |
| st.write(f"{relation.source.text} --{relation.label}--> {relation.target.text}") | |
| # 6. Create and display graph using streamlit-agraph | |
| st.subheader("Entity Relationship Graph:") | |
| create_graph(entities, relations) | |
| if __name__ == "__main__": | |
| main() | |