yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2025-09-21 06:03:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make EntityExtractor work when loaded from YAML (#1636)

Browse Source 
* Add set_config to EntityExtractor

* Import EntityExtractor in pipeline.py, or it won't be properly registered as a subclass
This commit is contained in:
Sara Zan 2021-10-22 14:41:26 +02:00 committed by GitHub
parent 0aba5ca57d
commit f67b213797
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
haystack/extractor/entity.py
View File

@ -20,6 +20,7 @@ class EntityExtractor(BaseComponent):
def __init__(self, def __init__(self,
model_name_or_path="dslim/bert-base-NER"): model_name_or_path="dslim/bert-base-NER"):
self.set_config(model_name_or_path=model_name_or_path)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
token_classifier = AutoModelForTokenClassification.from_pretrained(model_name_or_path) token_classifier = AutoModelForTokenClassification.from_pretrained(model_name_or_path)
self.model = pipeline("ner", model=token_classifier, tokenizer=tokenizer, aggregation_strategy="simple") self.model = pipeline("ner", model=token_classifier, tokenizer=tokenizer, aggregation_strategy="simple")

2
haystack/pipeline.py
View File

@ -35,6 +35,8 @@ from haystack.summarizer.base import BaseSummarizer
from haystack.translator.base import BaseTranslator from haystack.translator.base import BaseTranslator
from haystack.document_store.base import BaseDocumentStore from haystack.document_store.base import BaseDocumentStore
from haystack.question_generator import QuestionGenerator from haystack.question_generator import QuestionGenerator
from haystack.extractor import EntityExtractor
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)