neo4j csv generation now includes structured fields into the csvs on …

…the Document node
nlp-tlp · Jan 18, 2024 · d988a37 · d988a37
1 parent cf79caa
commit d988a37
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 3 deletions.
diff --git a/dist/puggle-0.2.10-py3-none-any.whl b/dist/puggle-0.2.10-py3-none-any.whl
diff --git a/dist/puggle-0.2.10.tar.gz b/dist/puggle-0.2.10.tar.gz
diff --git a/puggle/Dataset.py b/puggle/Dataset.py
@@ -243,10 +243,15 @@ def create_neo4j_csvs(
         ent_idxs = {}
         rel_freqs = {}
 
+        document_fieldnames = []
+
         for i, d in enumerate(self.documents):
+            if i == 0:
+                print(d.fields)
+                document_fieldnames = d.fields.keys()
             ann = d.annotation
 
-            docs.add(tuple([i]))
+            docs.add(tuple([i, *d.fields.values()]))
 
             for mention in ann.mentions:
                 t = tuple([mention.label, " ".join(mention.tokens)])
@@ -281,7 +286,7 @@ def create_neo4j_csvs(
 
         with open(documents_path, "w", newline="", encoding="utf-8") as f:
             writer = csv.writer(f)
-            writer.writerow(["doc_idx"])
+            writer.writerow(["doc_idx", *document_fieldnames])
             for row in list(docs):
                 writer.writerow(row)
             logger.info(

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "puggle"
-version = "0.2.10"
+version = "0.2.11"
 description = ""
 authors = ["Michael Stewart <[email protected]>"]
 readme = "README.md"