ResearchEngineering commited on
Commit
bd6fb60
·
verified ·
1 Parent(s): 5ba72b0

Update src/streamlit_app.py

Browse files

add model and tokenizer to NER pipeline

Files changed (1) hide show
  1. src/streamlit_app.py +17 -1
src/streamlit_app.py CHANGED
@@ -570,7 +570,23 @@ def load_financial_models():
570
  finbert_model = AutoModelForSequenceClassification.from_pretrained("yiyanghkust/finbert-tone", cache_dir=cache_dir)
571
 
572
  # Financial NER for entity extraction
573
- ner_pipeline = pipeline("ner", model="elastic/distilbert-base-cased-finetuned-conll03-english", aggregation_strategy="simple", cache_dir=cache_dir)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
574
 
575
  return finbert_tokenizer, finbert_model, ner_pipeline, None
576
  except Exception as e:
 
570
  finbert_model = AutoModelForSequenceClassification.from_pretrained("yiyanghkust/finbert-tone", cache_dir=cache_dir)
571
 
572
  # Financial NER for entity extraction
573
+ #ner_pipeline = pipeline("ner", model="elastic/distilbert-base-cased-finetuned-conll03-english", aggregation_strategy="simple", cache_dir=cache_dir)
574
+
575
+ # Load Financial NER model and tokenizer explicitly
576
+ ner_tokenizer = AutoTokenizer.from_pretrained(
577
+ "elastic/distilbert-base-cased-finetuned-conll03-english", cache_dir=cache_dir
578
+ )
579
+ ner_model = AutoModelForTokenClassification.from_pretrained(
580
+ "elastic/distilbert-base-cased-finetuned-conll03-english", cache_dir=cache_dir
581
+ )
582
+
583
+ # Then create pipeline using objects
584
+ ner_pipeline = pipeline(
585
+ "ner",
586
+ model=ner_model,
587
+ tokenizer=ner_tokenizer,
588
+ aggregation_strategy="simple",
589
+ )
590
 
591
  return finbert_tokenizer, finbert_model, ner_pipeline, None
592
  except Exception as e: