Skip to content

Commit 2ce8be6

Browse files
author
Val Brodsky
committed
Add integration test
1 parent 15f8181 commit 2ce8be6

File tree

10 files changed

+344
-7
lines changed

10 files changed

+344
-7
lines changed

‎labelbox/data/annotation_types/annotation.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import abc
2-
from typing import Any, Dict, List, Optional, Union, Type
2+
from typing import Any, Dict, List, Optional, Union
33

44
from labelbox.data.mixins import ConfidenceNotSupportedMixin, ConfidenceMixin
55

‎labelbox/data/annotation_types/ner/conversation_entity.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,3 @@
1-
from typing import Dict, Any
2-
3-
from pydantic import BaseModel, root_validator
4-
51
from labelbox.data.annotation_types.ner.text_entity import TextEntity
62
from labelbox.utils import _CamelCaseMixin
73

‎labelbox/data/serialization/ndjson/label.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ def _infer_media_type(
8282

8383
types = {type(annotation) for annotation in annotations}
8484
data = ImageData
85-
if TextEntity in types or ConversationEntity:
85+
if (TextEntity in types) or (ConversationEntity in types):
8686
data = TextData
8787
elif VideoClassificationAnnotation in types or VideoObjectAnnotation in types:
8888
data = VideoData

‎labelbox/data/serialization/ndjson/objects.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from ast import Bytes
22
from io import BytesIO
3-
from typing import Any, Dict, List, Tuple, Type, Union, Optional
3+
from typing import Any, Dict, List, Tuple, Union, Optional
44
import base64
55
from labelbox.data.annotation_types.ner.conversation_entity import ConversationEntity
66
from labelbox.data.mixins import ConfidenceMixin

‎tests/assets/conversation-1.json‎

Lines changed: 235 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,235 @@
1+
{
2+
"type":
3+
"application/vnd.labelbox.conversational",
4+
"version":
5+
1,
6+
"messages": [{
7+
"user": {
8+
"userId": "0",
9+
"name": "User"
10+
},
11+
"canLabel":
12+
true,
13+
"content":
14+
"I'd really like to take my client out to a nice restaurant that serves indian food.",
15+
"messageId":
16+
"0",
17+
"align":
18+
"left"
19+
}, {
20+
"user": {
21+
"userId": "1",
22+
"name": "Bot"
23+
},
24+
"canLabel":
25+
false,
26+
"content":
27+
"I show many restaurants that serve Indian food in that price range. What area would you like to travel to?",
28+
"messageId":
29+
"1",
30+
"align":
31+
"right"
32+
}, {
33+
"user": {
34+
"userId": "0",
35+
"name": "User"
36+
},
37+
"canLabel":
38+
true,
39+
"content":
40+
"I am looking for an expensive indian restaurant in the area of centre.",
41+
"messageId":
42+
"2",
43+
"align":
44+
"left"
45+
}, {
46+
"user": {
47+
"userId": "1",
48+
"name": "Bot"
49+
},
50+
"canLabel":
51+
false,
52+
"content":
53+
"Might I recommend Saffron Brasserie? That is an expensive Indian restaurant in the center of town. I can book a table for you, if you like.",
54+
"messageId":
55+
"3",
56+
"align":
57+
"right"
58+
}, {
59+
"user": {
60+
"userId": "0",
61+
"name": "User"
62+
},
63+
"canLabel": true,
64+
"content": "Sure thing, please book for 6 people at 19:30 on Saturday.",
65+
"messageId": "4",
66+
"align": "left"
67+
}, {
68+
"user": {
69+
"userId": "1",
70+
"name": "Bot"
71+
},
72+
"canLabel":
73+
false,
74+
"content":
75+
"Booking was successful. The table will be reserved for 15 minutes. Your Reference number is : RF00JUFQ.",
76+
"messageId":
77+
"5",
78+
"align":
79+
"right"
80+
}, {
81+
"user": {
82+
"userId": "0",
83+
"name": "User"
84+
},
85+
"canLabel":
86+
true,
87+
"content":
88+
"Okay great! Thank you so much. Could you also help me find a 3 star hotel in the area. I don't need wifi either.",
89+
"messageId":
90+
"6",
91+
"align":
92+
"left"
93+
}, {
94+
"user": {
95+
"userId": "1",
96+
"name": "Bot"
97+
},
98+
"canLabel":
99+
false,
100+
"content":
101+
"The alpha-milton guest house is in the north and is moderately priced. It has 3 stars and no internet or parking. Would you like to book a room?",
102+
"messageId":
103+
"7",
104+
"align":
105+
"right"
106+
}, {
107+
"user": {
108+
"userId": "0",
109+
"name": "User"
110+
},
111+
"canLabel": true,
112+
"content": "That sounds great. Please book that now.",
113+
"messageId": "8",
114+
"align": "left"
115+
}, {
116+
"user": {
117+
"userId": "1",
118+
"name": "Bot"
119+
},
120+
"canLabel": false,
121+
"content": "May I ask how many people are in your group?",
122+
"messageId": "9",
123+
"align": "right"
124+
}, {
125+
"user": {
126+
"userId": "0",
127+
"name": "User"
128+
},
129+
"canLabel": true,
130+
"content": "I have 6 people in my group. ",
131+
"messageId": "10",
132+
"align": "left"
133+
}, {
134+
"user": {
135+
"userId": "1",
136+
"name": "Bot"
137+
},
138+
"canLabel": false,
139+
"content": "How many days would you like to stay?",
140+
"messageId": "11",
141+
"align": "right"
142+
}, {
143+
"user": {
144+
"userId": "0",
145+
"name": "User"
146+
},
147+
"canLabel": true,
148+
"content": "2 nights, starting the same day as the reservation.",
149+
"messageId": "12",
150+
"align": "left"
151+
}, {
152+
"user": {
153+
"userId": "1",
154+
"name": "Bot"
155+
},
156+
"canLabel":
157+
false,
158+
"content":
159+
"I'm sorry. It looks like they're full. Would you like me to look for something else?",
160+
"messageId":
161+
"13",
162+
"align":
163+
"right"
164+
}, {
165+
"user": {
166+
"userId": "0",
167+
"name": "User"
168+
},
169+
"canLabel":
170+
true,
171+
"content":
172+
"Yes please. Is there something else available in that area? ",
173+
"messageId":
174+
"14",
175+
"align":
176+
"left"
177+
}, {
178+
"user": {
179+
"userId": "1",
180+
"name": "Bot"
181+
},
182+
"canLabel":
183+
false,
184+
"content":
185+
"I'm sorry, it looks like that is the only 3 star hotel in that area, would you like me to look somewhere else?",
186+
"messageId":
187+
"15",
188+
"align":
189+
"right"
190+
}, {
191+
"user": {
192+
"userId": "0",
193+
"name": "User"
194+
},
195+
"canLabel":
196+
true,
197+
"content":
198+
"Can we try for 1 night instead of 2? I'll need the reference number please.",
199+
"messageId":
200+
"16",
201+
"align":
202+
"left"
203+
}, {
204+
"user": {
205+
"userId": "1",
206+
"name": "Bot"
207+
},
208+
"canLabel":
209+
false,
210+
"content":
211+
"I was able to get one night, the reference number is 9XVT8M5T.",
212+
"messageId":
213+
"17",
214+
"align":
215+
"right"
216+
}, {
217+
"user": {
218+
"userId": "0",
219+
"name": "User"
220+
},
221+
"canLabel": true,
222+
"content": "Thank you so much!",
223+
"messageId": "18",
224+
"align": "left"
225+
}, {
226+
"user": {
227+
"userId": "1",
228+
"name": "Bot"
229+
},
230+
"canLabel": false,
231+
"content": "I'm glad to help, you're welcome!",
232+
"messageId": "19",
233+
"align": "right"
234+
}]
235+
}

‎tests/data/assets/ndjson/conversation_entity_import.json‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
[{
2+
"location": {
3+
"start": 67,
4+
"end": 128
5+
},
6+
"messageId": "some-message-id",
7+
"uuid": "5ad9c52f-058d-49c8-a749-3f20b84f8cd4",
8+
"dataRow": {
9+
"id": "cl6xnv9h61fv0085yhtoq06ht"
10+
},
11+
"name": "some-text-entity",
12+
"schemaId": "cl6xnuwt95lqq07330tbb3mfd",
13+
"classifications": [],
14+
"confidence": 0.53
15+
}]

‎tests/data/assets/ndjson/conversation_entity_without_confidence_import.json‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
[{
2+
"location": {
3+
"start": 67,
4+
"end": 128
5+
},
6+
"messageId": "some-message-id",
7+
"uuid": "5ad9c52f-058d-49c8-a749-3f20b84f8cd4",
8+
"dataRow": {
9+
"id": "cl6xnv9h61fv0085yhtoq06ht"
10+
},
11+
"name": "some-text-entity",
12+
"schemaId": "cl6xnuwt95lqq07330tbb3mfd",
13+
"classifications": []
14+
}]

‎tests/integration/annotation_import/conftest.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,16 @@ def dataset_pdf_entity(client, rand_gen, document_data_row):
338338
dataset.delete()
339339

340340

341+
@pytest.fixture
342+
def dataset_conversation_entity(client, rand_gen, conversation_entity_data_row):
343+
dataset = client.create_dataset(name=rand_gen(str))
344+
data_row_ids = []
345+
data_row = dataset.create_data_row(conversation_entity_data_row)
346+
data_row_ids.append(data_row.uid)
347+
yield dataset, data_row_ids
348+
dataset.delete()
349+
350+
341351
@pytest.fixture
342352
def configured_project_without_data_rows(client, configured_project, rand_gen):
343353
project = client.create_project(name=rand_gen(str),

‎tests/integration/annotation_import/test_conversation_import.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
import uuid
2+
from labelbox.data.annotation_types.annotation import ObjectAnnotation
3+
from labelbox.data.annotation_types.label import Label
4+
from labelbox.data.annotation_types.data.text import TextData
5+
from labelbox.data.annotation_types.ner import ConversationEntity
6+
7+
from labelbox.schema.annotation_import import MALPredictionImport
8+
9+
10+
def test_conversation_entity(client, configured_project_without_data_rows,
11+
dataset_conversation_entity, rand_gen):
12+
13+
conversation_entity_annotation = ConversationEntity(name="named-entity",
14+
start=0,
15+
end=8,
16+
message_id="4")
17+
18+
entities_annotation = ObjectAnnotation(name="named-entity",
19+
value=conversation_entity_annotation)
20+
21+
labels = []
22+
_, data_row_uids = dataset_conversation_entity
23+
configured_project_without_data_rows.create_batch(
24+
rand_gen(str),
25+
data_row_uids, # sample of data row objects
26+
5 # priority between 1(Highest) - 5(lowest)
27+
)
28+
29+
for data_row_uid in data_row_uids:
30+
labels.append(
31+
Label(data=TextData(uid=data_row_uid),
32+
annotations=[
33+
entities_annotation,
34+
]))
35+
36+
import_annotations = MALPredictionImport.create_from_objects(
37+
client=client,
38+
project_id=configured_project_without_data_rows.uid,
39+
name=f"import {str(uuid.uuid4())}",
40+
predictions=labels)
41+
import_annotations.wait_until_done()
42+
43+
assert import_annotations.errors == []

‎tests/integration/conftest.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,30 @@ def pdf_url(client):
177177
return {"row_data": {"pdf_url": pdf_url,}, "global_key": str(uuid.uuid4())}
178178

179179

180+
@pytest.fixture(scope="session")
181+
def pdf_entity_data_row(client):
182+
pdf_url = client.upload_file(
183+
'tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483.pdf')
184+
text_layer_url = client.upload_file(
185+
'tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483-lb-textlayer.json'
186+
)
187+
188+
return {
189+
"row_data": {
190+
"pdf_url": pdf_url,
191+
"text_layer_url": text_layer_url
192+
},
193+
"global_key": str(uuid.uuid4())
194+
}
195+
196+
197+
@pytest.fixture(scope="session")
198+
def conversation_entity_data_row(client):
199+
conversation_url = client.upload_file('tests/assets/conversation-1.json')
200+
201+
return {"row_data": conversation_url, "global_key": str(uuid.uuid1())}
202+
203+
180204
@pytest.fixture
181205
def project(client, rand_gen):
182206
project = client.create_project(name=rand_gen(str),

0 commit comments

Comments
 (0)