Skip to content

Commit 6e96d54

Browse files
dizcologydizcology
committed
add label_text test
1 parent 364765b commit 6e96d54

File tree

2 files changed

+115
-23
lines changed

2 files changed

+115
-23
lines changed

‎datalabeling/label_text.py

Lines changed: 28 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -18,53 +18,58 @@
1818

1919

2020
# [START datalabeling_label_text_beta]
21-
def label_text(dataset_resource_name, instruction_resource_name,
22-
annotation_spec_set_resource_name):
21+
def label_text(dataset_resource_name, instruction_resource_name, annotation_spec_set_resource_name):
2322
"""Labels a text dataset."""
2423
from google.cloud import datalabeling_v1beta1 as datalabeling
2524
client = datalabeling.DataLabelingServiceClient()
2625

2726
basic_config = datalabeling.types.HumanAnnotationConfig(
28-
instruction=instruction_resource_name,
29-
annotated_dataset_display_name='YOUR_ANNOTATED_DATASET_DISPLAY_NAME',
30-
label_group='YOUR_LABEL_GROUP',
31-
replica_count=1)
27+
instruction=instruction_resource_name,
28+
annotated_dataset_display_name='YOUR_ANNOTATED_DATASET_DISPLAY_NAME',
29+
label_group='YOUR_LABEL_GROUP',
30+
replica_count=1
31+
)
3232

3333
feature = datalabeling.enums.LabelTextRequest.Feature.TEXT_ENTITY_EXTRACTION
3434

3535
text_entity_extraction_config = datalabeling.types.TextEntityExtractionConfig(
36-
annotation_spec_set=annotation_spec_set_resource_name)
36+
annotation_spec_set=annotation_spec_set_resource_name)
3737

3838
response = client.label_text(
39-
dataset_resource_name,
40-
basic_config,
41-
feature,
42-
text_entity_extraction_config=text_entity_extraction_config)
39+
dataset_resource_name,
40+
basic_config,
41+
feature,
42+
text_entity_extraction_config=text_entity_extraction_config
43+
)
4344

4445
print('Label_text operation name: {}'.format(response.operation.name))
4546
return response
4647
# [END datalabeling_label_text_beta]
4748

4849
if __name__ == '__main__':
4950
parser = argparse.ArgumentParser(
50-
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
51+
description=__doc__,
52+
formatter_class=argparse.RawDescriptionHelpFormatter
53+
)
5154

5255
parser.add_argument(
53-
'--dataset-resource-name',
54-
help='Dataset resource name. Required.',
55-
required=True)
56+
'--dataset-resource-name',
57+
help='Dataset resource name. Required.',
58+
required=True
59+
)
5660

5761
parser.add_argument(
58-
'--instruction-resource-name',
59-
help='Instruction resource name. Required.',
60-
required=True)
62+
'--instruction-resource-name',
63+
help='Instruction resource name. Required.',
64+
required=True
65+
)
6166

6267
parser.add_argument(
63-
'--annotation-spec-set-resource-name',
64-
help='Annotation spec set resource name. Required.',
65-
required=True)
68+
'--annotation-spec-set-resource-name',
69+
help='Annotation spec set resource name. Required.',
70+
required=True
71+
)
6672

6773
args = parser.parse_args()
6874

69-
label_text(args.dataset_resource_name, args.instruction_resource_name,
70-
args.annotation_spec_set_resource_name)
75+
label_text(args.dataset_resource_name, args.instruction_resource_name, args.annotation_spec_set_resource_name)

‎datalabeling/label_text_test.py

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
#!/usr/bin/env python
2+
3+
# Copyright 2019 Google, Inc
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
import os
18+
import pytest
19+
20+
from google.cloud import datalabeling_v1beta1 as datalabeling
21+
import create_annotation_spec_set
22+
import create_instruction
23+
import import_data
24+
import label_text
25+
import manage_dataset
26+
27+
PROJECT_ID = os.getenv('GCLOUD_PROJECT')
28+
29+
30+
@pytest.fixture(scope='function')
31+
def dataset():
32+
# create a temporary dataset
33+
dataset = manage_dataset.create_dataset(PROJECT_ID)
34+
35+
# import some data to it
36+
import_data.import_data(dataset.name, 'TEXT',
37+
'gs://cloud-samples-data/datalabeling/text/text_dataset.csv')
38+
39+
yield dataset
40+
41+
# tear down
42+
manage_dataset.delete_dataset(dataset.name)
43+
44+
45+
@pytest.fixture(scope='function')
46+
def annotation_spec_set():
47+
# create a temporary annotation_spec_set
48+
annotation_spec_set = create_annotation_spec_set.create_annotation_spec_set(PROJECT_ID)
49+
50+
yield annotation_spec_set
51+
52+
# tear down
53+
client = datalabeling.DataLabelingServiceClient()
54+
client.delete_annotation_spec_set(annotation_spec_set.name)
55+
56+
57+
@pytest.fixture(scope='function')
58+
def instruction():
59+
# create a temporary instruction
60+
instruction = create_instruction.create_instruction(
61+
PROJECT_ID, 'TEXT',
62+
'gs://cloud-samples-data/datalabeling/instruction/test.pdf')
63+
64+
yield instruction
65+
66+
# tear down
67+
client = datalabeling.DataLabelingServiceClient()
68+
client.delete_instruction(instruction.name)
69+
70+
71+
# Passing in dataset as the last argument in test_label_text since it needs to be deleted before the annotation_spec_set can be deleted.
72+
@pytest.mark.slow
73+
def test_label_text(capsys, annotation_spec_set, instruction, dataset):
74+
75+
# Start labeling.
76+
response = label_text.label_text(dataset.name, instruction.name, annotation_spec_set.name)
77+
out, _ = capsys.readouterr()
78+
assert 'Label_text operation name: ' in out
79+
operation_name = response.operation.name
80+
81+
# Cancels the labeling operation.
82+
response.cancel()
83+
assert response.cancelled() == True
84+
85+
client = datalabeling.DataLabelingServiceClient()
86+
cancel_response = client.transport._operations_client.cancel_operation(
87+
operation_name)

0 commit comments

Comments
 (0)