diff --git a/nifi/user-scripts/cogstack_cohort_generate_data.py b/nifi/user-scripts/cogstack_cohort_generate_data.py
index 51ac3e0b..005a8f0f 100644
--- a/nifi/user-scripts/cogstack_cohort_generate_data.py
+++ b/nifi/user-scripts/cogstack_cohort_generate_data.py
@@ -173,7 +173,7 @@ def _process_annotation_records(annotation_records: list):
                 annotation_entity = annotation_record["_source"]
             docid = annotation_entity[ANNOTATION_DOCUMENT_ID_FIELD_NAME]
 
-            if str(docid) in unique_doc_ids.keys():
+            if str(docid) in unique_doc_ids:
                 patient_id = global_doc2ptt[str(docid)]
                 cui = annotation_entity["nlp.cui"]