From 63221f31aebf6a3cab61eae2181387e98e102700 Mon Sep 17 00:00:00 2001
From: Hironsan <light.tree.1.13@gmail.com>
Date: Fri, 25 Feb 2022 11:15:02 +0900
Subject: [PATCH] Add EntityAndRelationWriter

---
 backend/data_export/pipeline/writers.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/backend/data_export/pipeline/writers.py b/backend/data_export/pipeline/writers.py
index f1fbc284..aacf8904 100644
--- a/backend/data_export/pipeline/writers.py
+++ b/backend/data_export/pipeline/writers.py
@@ -156,3 +156,25 @@ class IntentAndSlotWriter(LineWriter):
                 {"id": record.id, "text": record.data, "cats": [], "entities": [], **record.metadata},
                 ensure_ascii=False,
             )
+
+
+class EntityAndRelationWriter(LineWriter):
+    extension = "jsonl"
+
+    def create_line(self, record):
+        if isinstance(record.label, dict):
+            return json.dumps(
+                {
+                    "id": record.id,
+                    "text": record.data,
+                    "relations": record.label.get("relations", []),
+                    "entities": record.label.get("entities", []),
+                    **record.metadata,
+                },
+                ensure_ascii=False,
+            )
+        else:
+            return json.dumps(
+                {"id": record.id, "text": record.data, "relations": [], "entities": [], **record.metadata},
+                ensure_ascii=False,
+            )