GCP custom entity extractor added (AutoML)

auto_labeling_pipeline/mappings.py

@@ -58,6 +58,9 @@ class GCPEntitiesTemplate(MappingTemplate):
     label_collection = SequenceLabels
     template_file = 'gcp_entities.jinja2'
 
+class GCPCustomEntitiesTemplate(MappingTemplate):
+    label_collection = SequenceLabels
+    template_file = 'gcp_custom_entities.jinja2'
 
 class AmazonRekognitionLabelDetectionTemplate(MappingTemplate):
     label_collection = ClassificationLabels

auto_labeling_pipeline/menu.py

@@ -70,6 +70,11 @@ def register(cls, task: Type[t.Task], model: Type[mo.RequestModel], template: Ty
     mo.GCPEntitiesRequestModel,
     mp.GCPEntitiesTemplate
 )
+Options.register(
+    t.SequenceLabeling,
+    mo.GCPCustomEntitiesRequestModel,
+    mp.GCPCustomEntitiesTemplate
+)
 Options.register(
     t.SequenceLabeling,
     mo.AmazonComprehendEntityRequestModel,

auto_labeling_pipeline/models.py

@@ -107,6 +107,32 @@ def send(self, text: str):
         response = requests.post(url, headers=headers, params=params, json=body).json()
         return response
 
+class GCPCustomEntitiesRequestModel(RequestModel):
+    """
+    This allow you to analyze entities in a text using an AutoML custom entity extractor
+    <a href="https://cloud.google.com/natural-language/automl/docs"> AutoML Natural Language API</a>.
+    """
+    url: str
+    authorization: str
+
+    class Config:
+        title = 'GCP Custom Entity Analysis'
+
+    def send(self, text: str):
+        url = self.url
+        headers = {'Content-Type': 'application/json',
+                   'Authorization': self.authorization}
+        params = {}
+        body = {
+            "payload": {
+                "textSnippet": {
+                    "content": text,
+                    "mime_type": "text/plain"
+                }
+            }
+        }
+        response = requests.post(url, headers=headers, params=params, json=body).json()
+        return response
 
 class AWSMixin(BaseModel):
     aws_access_key: str

auto_labeling_pipeline/templates/gcp_custom_entities.jinja2

@@ -0,0 +1,14 @@
+[
+  {% for entity in input.payload %}
+    {% set outer_loop = loop %}
+    {%- set start_offset = entity.textExtraction.textSegment.startOffset -%}
+    {%- set end_offset = entity.textExtraction.textSegment.endOffset -%}
+      {% if (start_offset) and (end_offset) %}
+        {
+          "label": "{{ entity.displayName }}",
+          "start_offset": {{ start_offset }},
+          "end_offset": {{ end_offset }}
+        }{% if not outer_loop.last %},{% endif %}
+      {% endif %}
+  {% endfor %}
+]