ShubhamMhaske commited on
Commit
30fc0eb
·
verified ·
1 Parent(s): 9b6b16d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +54 -0
app.py CHANGED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app_offline_ner_min.py
2
+ import os
3
+ os.environ["TRANSFORMERS_OFFLINE"] = "1" # force offline per HF docs
4
+ os.environ["TOKENIZERS_PARALLELISM"] = "false"
5
+
6
+ import torch
7
+ from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
8
+ import gradio as gr
9
+
10
+ # point to your local snapshot downloaded by prepare_model.py
11
+ # path: ./models/biomedical-ner-all
12
+ HERE = os.path.dirname(os.path.abspath(__file__))
13
+ LOCAL_MODEL_DIR = os.path.join(HERE, "models", "biomedical-ner-all")
14
+
15
+ # load strictly from disk
16
+ tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_DIR, local_files_only=True)
17
+ model = AutoModelForTokenClassification.from_pretrained(LOCAL_MODEL_DIR, local_files_only=True)
18
+
19
+ device = 0 if torch.cuda.is_available() else -1
20
+ ner_pipe = pipeline(
21
+ task="token-classification", # NER
22
+ model=model,
23
+ tokenizer=tokenizer,
24
+ aggregation_strategy="simple", # merge subword tokens into entities
25
+ device=device
26
+ )
27
+
28
+ def run_ner(text: str):
29
+ if not text.strip():
30
+ return {"text": "", "entities": []}, []
31
+ out = ner_pipe(text)
32
+ # Gradio HighlightedText format
33
+ highlighted = {
34
+ "text": text,
35
+ "entities": [
36
+ {"entity": r["entity_group"], "start": int(r["start"]), "end": int(r["end"]), "score": float(r["score"])}
37
+ for r in out
38
+ ],
39
+ }
40
+ # also return raw rows for inspection
41
+ rows = [
42
+ {"entity": r["entity_group"], "word": r["word"], "score": float(r["score"]), "start": int(r["start"]), "end": int(r["end"])}
43
+ for r in out
44
+ ]
45
+ return highlighted, rows
46
+
47
+ with gr.Blocks() as demo:
48
+ gr.Markdown("# 🩺 Biomedical NER (offline, local model)")
49
+ inp = gr.Textbox(label="Enter text", value="Patient has a history of asthma treated with albuterol.")
50
+ ner_view = gr.HighlightedText(label="Entities", combine_adjacent=True)
51
+ table = gr.Dataframe(label="Raw predictions", interactive=False)
52
+ inp.change(run_ner, inp, [ner_view, table])
53
+
54
+ demo.launch(debug=True)