From e159e9cd4fcd6d4270eaabdce875138cfdee87ca Mon Sep 17 00:00:00 2001
From: adamjanovsky
Date: Thu, 21 Sep 2023 14:43:10 +0200
Subject: clean labels when loading dataframes

---
 src/sec_certs/model/references/segment_extractor.py | 2 ++
 1 file changed, 2 insertions(+)

(limited to 'src')

diff --git a/src/sec_certs/model/references/segment_extractor.py b/src/sec_certs/model/references/segment_extractor.py
index bbb526b1..ffa576f2 100644
--- a/src/sec_certs/model/references/segment_extractor.py
+++ b/src/sec_certs/model/references/segment_extractor.py
@@ -276,6 +276,7 @@ class ReferenceSegmentExtractor:
             return (
                 pd.read_csv(
                     pth,
+                    na_values=["None"],
                     dtype={
                         "dgst": str,
                         "canonical_reference_keyword": str,
@@ -348,6 +349,7 @@ class ReferenceSegmentExtractor:
                 label=lambda df_: [
                     annotations_dict.get(x) for x in zip(df_["dgst"], df_["canonical_reference_keyword"])
                 ],
+                actual_reference_keywords=lambda df_: df_.actual_reference_keywords.map(list),
                 cert_name=lambda df_: df_.dgst.map(dgst_to_cert_name),
                 referenced_cert_name=lambda df_: df_.canonical_reference_keyword.map(cert_id_to_cert_name),
                 cert_versions=lambda df_: df_.dgst.map(dgst_to_extracted_versions),
-- 
cgit v1.3.1