Loading models hosted on huggingface

brabbit61 · brabbit61 · commit 8042e60f8064 · 2023-06-26T21:24:43.000-07:00
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -14,9 +14,6 @@ services:
     build: 
       dockerfile: ./docker/entity-extraction-pipeline/Dockerfile
       context: .
-      args:
-        HF_NER_MODEL_NAME: "roberta-finetuned-v3"
-        SPACY_NER_MODEL_NAME: "spacy-transformer-v3"
     ports:
       - "5000:5000"
     volumes:
diff --git a/docker/entity-extraction-pipeline/Dockerfile b/docker/entity-extraction-pipeline/Dockerfile
@@ -10,25 +10,18 @@ COPY docker/entity-extraction-pipeline/requirements.txt .
 # Install the required Python packages
 RUN pip install --no-cache-dir -r requirements.txt
 RUN python -m nltk.downloader stopwords
+RUN pip install https://huggingface.co/finding-fossils/metaextractor-spacy/resolve/main/en_metaextractor_spacy-any-py3-none-any.whl
 
 # Copy the entire repository folder into the container
 COPY src ./src
 
-# Build args
-ARG HF_NER_MODEL_NAME
-ARG SPACY_NER_MODEL_NAME
-
 # Set env variables for when running the container
-ENV HF_NER_MODEL_NAME=${HF_NER_MODEL_NAME}
-ENV SPACY_NER_MODEL_NAME=${SPACY_NER_MODEL_NAME}
-ENV USE_NER_MODEL_TYPE=huggingface
+ENV HF_NER_MODEL_NAME="finding-fossils/metaextractor"
+ENV SPACY_NER_MODEL_NAME="en_metaextractor_spacy"
+ENV USE_NER_MODEL_TYPE="huggingface"
 ENV MAX_ARTICLES=-1
 ENV MAX_SENTENCES=-1
 
-# Copy in the model defined by the env variable NER_MODEL_NAME from models folder
-COPY models/ner/${HF_NER_MODEL_NAME} ./models/ner/${HF_NER_MODEL_NAME} 
-COPY models/ner/${SPACY_NER_MODEL_NAME} ./models/ner/${SPACY_NER_MODEL_NAME}
-
 # non-root user control inspired from here: https://stackoverflow.com/questions/66349101/docker-non-root-user-does-not-have-writing-permissions-when-using-volumes
 # Create a non-root user that owns the input/outputs directory by default
 RUN useradd -r extraction-user          # no specific user ID
diff --git a/src/entity_extraction/spacy_entity_extraction.py b/src/entity_extraction/spacy_entity_extraction.py
@@ -32,8 +32,13 @@ def spacy_extract_all(
     """
 
     if ner_model == None:
-        logger.info("Empty model passed, return 0 labels.")
-        return []
+        try:
+            import en_metaextractor_spacy
+            ner_model = en_metaextractor_spacy.load()
+        except:
+            logger.error(f"Spacy model en_metaextractor_spacy not found.")
+            logger.info("Empty model passed, return 0 labels.")
+            return []
 
     entities = []
     doc = ner_model(text)
diff --git a/src/pipeline/entity_extraction_pipeline.py b/src/pipeline/entity_extraction_pipeline.py
@@ -34,8 +34,8 @@
 load_dotenv(find_dotenv())
 
 # get the MODEL_NAME from environment variables
-HF_NER_MODEL_NAME = os.getenv("HF_NER_MODEL_NAME", "roberta-finetuned-v3")
-SPACY_NER_MODEL_NAME = os.getenv("SPACY_NER_MODEL_NAME", "spacy-transformer-v3")
+HF_NER_MODEL_NAME = os.getenv("HF_NER_MODEL_NAME", "finding-fossils/metaextractor")
+SPACY_NER_MODEL_NAME = os.getenv("SPACY_NER_MODEL_NAME", "en_metaextractor_spacy")
 USE_NER_MODEL_TYPE = os.getenv("USE_NER_MODEL_TYPE", "huggingface")
 
 logger = get_logger(__name__)
@@ -286,7 +286,7 @@ def recreate_original_sentences_with_labels(row):
 def extract_entities(
     article_text_data: pd.DataFrame,
     model_type: str = "huggingface",
-    model_path: str = os.path.join("models", "ner", "roberta-finetuned-v3"),
+    model_path: str = "finding-fossils/metaextractor",
 ) -> pd.DataFrame:
     """
     Extracts the entities from the article text data.
@@ -574,10 +574,10 @@ def main():
 
             if USE_NER_MODEL_TYPE == "huggingface":
                 logger.info(f"Using HuggingFace model {HF_NER_MODEL_NAME}")
-                model_path = os.path.join("models", "ner", HF_NER_MODEL_NAME)
+                model_path = HF_NER_MODEL_NAME
             elif USE_NER_MODEL_TYPE == "spacy":
                 logger.info(f"Using Spacy model {SPACY_NER_MODEL_NAME}")
-                model_path = os.path.join("models", "ner", SPACY_NER_MODEL_NAME)
+                model_path = SPACY_NER_MODEL_NAME
             else:
                 raise ValueError(
                     f"Model type {USE_NER_MODEL_TYPE} not supported. Please set MODEL_TYPE to either 'huggingface' or 'spacy'."
