Spaces:
Sleeping
Sleeping
File size: 3,568 Bytes
93411f7 04a8bf0 e81f16d 93411f7 e81f16d 04a8bf0 e81f16d 04a8bf0 e81f16d 04a8bf0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
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
@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()
|