ayushsinha commited on
Commit
7375de2
Β·
verified Β·
1 Parent(s): f909de5

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +57 -0
  2. 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