optionally allow consolidation via env variables (#288)

Author: de-code

.env

@@ -1,5 +1,9 @@
 COMPOSE_PROJECT_NAME=sciencebeam
 IMAGE_TAG=develop
 VERSION=
-GROBID_TAG=0.5.3
+GROBID_TAG=0.6.0
 SCIENCEBEAM_DOC_CONVERT_PROCESS_TIMEOUT=60
+SCIENCEBEAM_CONSOLIDATE_HEADER=0
+SCIENCEBEAM_CONSOLIDATE_CITATIONS=0
+SCIENCEBEAM_INCLUDE_RAW_AFFILIATIONS=1
+SCIENCEBEAM_INCLUDE_RAW_CITATIONS=1

docker-compose.yml

@@ -14,11 +14,12 @@ services:
             --grobid-url http://grobid:8070/api
         environment:
             - SCIENCEBEAM_DOC_CONVERT_PROCESS_TIMEOUT=${SCIENCEBEAM_DOC_CONVERT_PROCESS_TIMEOUT}
+            - SCIENCEBEAM_CONSOLIDATE_HEADER=${SCIENCEBEAM_CONSOLIDATE_HEADER}
+            - SCIENCEBEAM_CONSOLIDATE_CITATIONS=${SCIENCEBEAM_CONSOLIDATE_CITATIONS}
+            - SCIENCEBEAM_INCLUDE_RAW_AFFILIATIONS=${SCIENCEBEAM_INCLUDE_RAW_AFFILIATIONS}
+            - SCIENCEBEAM_INCLUDE_RAW_CITATIONS=${SCIENCEBEAM_INCLUDE_RAW_CITATIONS}
         ports:
             - "8075:8075"
-        networks:
-            - default
-            - grobid
 
     sciencebeam-base-dev:
         build:
@@ -46,11 +47,3 @@ services:
             - JAVA_OPTS=-Xmx1g
         ports:
             - "8070:8070"
-        networks:
-            - grobid
-
-networks:
-    default:
-        internal: false
-    grobid:
-        internal: true

sciencebeam/transformers/grobid_service.py

@@ -1,6 +1,8 @@
+import os
 from io import BytesIO
 import logging
 from functools import partial
+from typing import Dict
 
 import requests
 
@@ -18,18 +20,80 @@ class GrobidApiPaths:
     PROCESS_FULL_TEXT_DOCUMENT = '/processFulltextDocument'
 
 
+class GrobidServiceConfigEnvVariables:
+    CONSOLIDATE_HEADER = 'SCIENCEBEAM_CONSOLIDATE_HEADER'
+    CONSOLIDATE_CITATIONS = 'SCIENCEBEAM_CONSOLIDATE_CITATIONS'
+    INCLUDE_RAW_AFFILIATIONS = 'SCIENCEBEAM_INCLUDE_RAW_AFFILIATIONS'
+    INCLUDE_RAW_CITATIONS = 'SCIENCEBEAM_INCLUDE_RAW_CITATIONS'
+
+
 service_wrapper = GrobidServiceWrapper()
 
 
 def get_logger():
     return logging.getLogger(__name__)
 
 
+class GrobidServiceConfig:
+    def __init__(
+            self,
+            consolidate_header: bool = False,
+            consolidate_citations: bool = False,
+            include_raw_affiliations: bool = True,
+            include_raw_citations: bool = True):
+        self.consolidate_header = consolidate_header
+        self.consolidate_citations = consolidate_citations
+        self.include_raw_affiliations = include_raw_affiliations
+        self.include_raw_citations = include_raw_citations
+
+
+def get_grobid_service_config() -> GrobidServiceConfig:
+    env_to_prop_map = {
+        GrobidServiceConfigEnvVariables.CONSOLIDATE_HEADER: 'consolidate_header',
+        GrobidServiceConfigEnvVariables.CONSOLIDATE_CITATIONS: 'consolidate_citations',
+        GrobidServiceConfigEnvVariables.INCLUDE_RAW_AFFILIATIONS: 'include_raw_affiliations',
+        GrobidServiceConfigEnvVariables.INCLUDE_RAW_CITATIONS: 'include_raw_citations'
+    }
+    config = GrobidServiceConfig()
+    for env_name, prop_name in env_to_prop_map.items():
+        env_value = os.environ.get(env_name)
+        if env_value:
+            setattr(config, prop_name, env_value == '1')
+    return config
+
+
+def _get_bool_int_param(bool_value: bool) -> str:
+    return '1' if bool_value else '0'
+
+
+def get_request_data_for_config(grobid_service_config: GrobidServiceConfig) -> Dict[str, str]:
+    return {
+        'consolidateHeader': _get_bool_int_param(
+            grobid_service_config.consolidate_header
+        ),
+        'consolidateCitations': _get_bool_int_param(
+            grobid_service_config.consolidate_citations
+        ),
+        'includeRawAffiliations': _get_bool_int_param(
+            grobid_service_config.include_raw_affiliations
+        ),
+        'includeRawCitations': _get_bool_int_param(
+            grobid_service_config.include_raw_citations
+        )
+    }
+
+
 def start_service_if_not_running():
     service_wrapper.start_service_if_not_running()
 
 
-def run_grobid_service(item, base_url, path, start_service=True, field_name=None):
+def run_grobid_service(
+        item,
+        base_url: str,
+        path: str,
+        grobid_service_config: GrobidServiceConfig = None,
+        start_service: bool = True,
+        field_name: str = None):
     """
     Translates PDF content via the GROBID service.
 
@@ -69,12 +133,7 @@ def run_grobid_service(item, base_url, path, start_service=True, field_name=None
         response = requests.post(
             url,
             files={'input': (filename, BytesIO(content))},
-            data={
-                'consolidateHeader': '0',
-                'consolidateCitations': '0',
-                'includeRawAffiliations': '1',
-                'includeRawCitations': '1'
-            }
+            data=get_request_data_for_config(grobid_service_config)
         )
     response.raise_for_status()
     result_content = response.content
@@ -83,11 +142,19 @@ def run_grobid_service(item, base_url, path, start_service=True, field_name=None
     return result_content
 
 
-def grobid_service(base_url, path, start_service=True, field_name=None):
+def grobid_service(
+        base_url: str,
+        path: str,
+        start_service: bool = True,
+        field_name: str = None,
+        grobid_service_config: GrobidServiceConfig = None):
+    if not grobid_service_config:
+        grobid_service_config = get_grobid_service_config()
     return partial(
         run_grobid_service,
         base_url=base_url,
         path=path,
         start_service=start_service,
-        field_name=field_name
+        field_name=field_name,
+        grobid_service_config=grobid_service_config
     )

tests/transformers/grobid_service_test.py

@@ -2,7 +2,14 @@
 
 import pytest
 
-from sciencebeam.transformers.grobid_service import grobid_service as create_grobid_service
+from sciencebeam.transformers.grobid_service import (
+    GrobidServiceConfigEnvVariables,
+    GrobidServiceConfig,
+    get_grobid_service_config,
+    get_request_data_for_config,
+    grobid_service as create_grobid_service
+)
+
 
 BASE_URL = 'http://grobid/api'
 PATH_1 = '/path1'
@@ -21,6 +28,52 @@ def _mock_requests_post():
         yield requests_post
 
 
+@pytest.fixture(name='environ', autouse=True)
+def _mock_environ():
+    with patch('os.environ', {}) as mock:
+        yield mock
+
+
+@pytest.fixture(name='grobid_service_config')
+def _grobid_service_config() -> GrobidServiceConfig:
+    return GrobidServiceConfig()
+
+
+class TestGetGrobidServiceConfig:
+    def test_should_return_default_config(self):
+        config = get_grobid_service_config()
+        assert not config.consolidate_header
+        assert not config.consolidate_citations
+        assert config.include_raw_affiliations
+        assert config.include_raw_citations
+
+    def test_should_be_able_toggle_config(self, environ: dict):
+        environ[GrobidServiceConfigEnvVariables.CONSOLIDATE_HEADER] = '1'
+        environ[GrobidServiceConfigEnvVariables.CONSOLIDATE_CITATIONS] = '1'
+        environ[GrobidServiceConfigEnvVariables.INCLUDE_RAW_AFFILIATIONS] = '0'
+        environ[GrobidServiceConfigEnvVariables.INCLUDE_RAW_CITATIONS] = '0'
+        config = get_grobid_service_config()
+        assert config.consolidate_header
+        assert config.consolidate_citations
+        assert not config.include_raw_affiliations
+        assert not config.include_raw_citations
+
+
+class TestGetRequestDataForConfig:
+    def test_should_generate_dict_for_default_config(
+            self, grobid_service_config: GrobidServiceConfig):
+        grobid_service_config.consolidate_header = False
+        grobid_service_config.consolidate_citations = False
+        grobid_service_config.include_raw_affiliations = True
+        grobid_service_config.include_raw_citations = True
+        assert get_request_data_for_config(grobid_service_config) == {
+            'consolidateHeader': '0',
+            'consolidateCitations': '0',
+            'includeRawAffiliations': '1',
+            'includeRawCitations': '1'
+        }
+
+
 class TestCreateGrobidService:
     def test_should_pass_url_and_data_as_file(self, requests_post):
         create_grobid_service(BASE_URL, PATH_1, start_service=False)(