hassoudi commited on
Commit
4a49fad
·
verified ·
1 Parent(s): 38e01ae

Upload model.py

Browse files
Files changed (1) hide show
  1. model.py +41 -9
model.py CHANGED
@@ -1,18 +1,50 @@
1
  import os
2
- from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 
 
3
 
4
- # Safe writable cache dirs in Hugging Face Spaces
5
- os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
6
- os.environ["HF_HOME"] = "/tmp/hf_home"
7
- os.environ["HF_HUB_CACHE"] = "/tmp/hf_hub"
8
 
9
 
10
 
11
- HF_TOKEN = os.getenv("HF_TOKEN")
12
  MODEL_ID = "TypicaAI/magbert-ner"
13
 
14
- tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
15
- model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, token=HF_TOKEN)
 
 
16
 
17
 
18
- ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="first")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import os
2
+ #from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
3
+ from transformers import pipeline
4
+ from huggingface_hub import login
5
 
6
+ # Safe writable cache dirs in Hugging Face Spaces move to space settings vars`
7
+ #os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
8
+ #os.environ["HF_HOME"] = "/tmp/hf_home"
9
+ #os.environ["HF_HUB_CACHE"] = "/tmp/hf_hub"
10
 
11
 
12
 
 
13
  MODEL_ID = "TypicaAI/magbert-ner"
14
 
15
+ #HF_TOKEN = os.getenv("HF_TOKEN")
16
+ #tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
17
+ #model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, token=HF_TOKEN)
18
+ #ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="first")
19
 
20
 
21
+ # Authenticate using the secret `HFTOKEN`
22
+ def authenticate_with_token():
23
+ """Authenticate with the Hugging Face API using the HFTOKEN secret."""
24
+ hf_token = os.getenv("HF_TOKEN") # Retrieve the token from environment variables
25
+ if not hf_token:
26
+ raise ValueError("HF_TOKEN is not set. Please add it to the Secrets in your Space settings.")
27
+ login(token=hf_token)
28
+
29
+
30
+ def load_healthcare_ner_pipeline():
31
+ """Load the Hugging Face pipeline for Healthcare NER."""
32
+ global ner_pipeline
33
+ if ner_pipeline is None:
34
+
35
+ # Authenticate and initialize pipeline
36
+ authenticate_with_token()
37
+
38
+ ner_pipeline = pipeline(
39
+ "token-classification",
40
+ model=MODEL_ID,
41
+ aggregation_strategy="first" # Groups B- and I- tokens into entities
42
+ )
43
+ return ner_pipeline
44
+
45
+
46
+
47
+
48
+ # Initialize global pipeline
49
+ ner_pipeline = load_healthcare_ner_pipeline()
50
+