Spaces:
Runtime error
Runtime error
Upload 2 files
Browse files- app.py +57 -0
- requirements.txt +4 -0
app.py
ADDED
@@ -0,0 +1,57 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import torch
|
3 |
+
from transformers import pipeline
|
4 |
+
|
5 |
+
# Load the NER model
|
6 |
+
model_name = "AventIQ-AI/bert-medical-entity-extraction"
|
7 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
8 |
+
|
9 |
+
print("Loading model...")
|
10 |
+
ner_pipeline = pipeline("ner", model=model_name, tokenizer=model_name, aggregation_strategy="simple", device=0 if torch.cuda.is_available() else -1)
|
11 |
+
|
12 |
+
# Define entity mapping based on README
|
13 |
+
entity_mapping = {
|
14 |
+
"LABEL_1": "Symptom",
|
15 |
+
"LABEL_2": "Disease",
|
16 |
+
"LABEL_3": "Medication",
|
17 |
+
"LABEL_4": "Treatment",
|
18 |
+
"LABEL_5": "Anatomy",
|
19 |
+
"LABEL_6": "Medical Procedure"
|
20 |
+
}
|
21 |
+
|
22 |
+
def extract_medical_entities(text):
|
23 |
+
"""Extract relevant medical entities from the input text."""
|
24 |
+
if not text.strip():
|
25 |
+
return "β οΈ Please enter a valid medical text."
|
26 |
+
|
27 |
+
print(f"Processing: {text}")
|
28 |
+
entities = ner_pipeline(text)
|
29 |
+
|
30 |
+
# Filter out non-entity labels (e.g., "O" or punctuation)
|
31 |
+
relevant_entities = [
|
32 |
+
f"π **{entity['word'].replace('##', '')}** β `{entity_mapping.get(entity['entity_group'], entity['entity_group'])}`"
|
33 |
+
for entity in entities if entity['entity_group'] in entity_mapping
|
34 |
+
]
|
35 |
+
|
36 |
+
response = "\n".join(relevant_entities) if relevant_entities else "β οΈ No relevant medical entities detected."
|
37 |
+
print(f"Response: {response}")
|
38 |
+
return response
|
39 |
+
|
40 |
+
# Create Gradio Interface
|
41 |
+
iface = gr.Interface(
|
42 |
+
fn=extract_medical_entities,
|
43 |
+
inputs=gr.Textbox(label="π Enter Medical Text", placeholder="Type or paste a medical report...", lines=3),
|
44 |
+
outputs=gr.Textbox(label="π₯ Extracted Medical Entities", placeholder="Detected medical terms will appear here...", lines=5),
|
45 |
+
title="π¬ Medical Entity Extraction",
|
46 |
+
description="π Enter a medical-related text, and the AI will extract **diseases, symptoms, medications, and treatments.**",
|
47 |
+
theme="compact",
|
48 |
+
allow_flagging="never",
|
49 |
+
examples=[
|
50 |
+
["The patient is diagnosed with diabetes and prescribed metformin."],
|
51 |
+
["Symptoms include fever, sore throat, and fatigue."],
|
52 |
+
["He underwent a knee replacement surgery at Mayo Clinic."]
|
53 |
+
],
|
54 |
+
)
|
55 |
+
|
56 |
+
if __name__ == "__main__":
|
57 |
+
iface.launch()
|
requirements.txt
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
torch
|
2 |
+
transformers
|
3 |
+
gradio
|
4 |
+
sentencepiece
|