hassoudi commited on
Commit
9dcc0d8
·
verified ·
1 Parent(s): 56f96a2

Upload model.py

Browse files
Files changed (1) hide show
  1. model.py +6 -2
model.py CHANGED
@@ -1,10 +1,14 @@
1
  import os
2
  from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
3
 
 
 
 
4
  HF_TOKEN = os.getenv("HF_TOKEN")
5
  MODEL_ID = "TypicaAI/magbert-ner"
6
 
7
- tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
8
- model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
 
9
 
10
  ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="first")
 
1
  import os
2
  from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
3
 
4
+ os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
5
+
6
+
7
  HF_TOKEN = os.getenv("HF_TOKEN")
8
  MODEL_ID = "TypicaAI/magbert-ner"
9
 
10
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
11
+ model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, token=HF_TOKEN)
12
+
13
 
14
  ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="first")