fixes loading label studio matches into CCDataset

1 files changed, 7 insertions, 6 deletions

diff --git a/sec_certs/dataset/common_criteria.py b/sec_certs/dataset/common_criteria.py
index c1c869e2..2f180ed6 100644
--- a/sec_certs/dataset/common_criteria.py
+++ b/sec_certs/dataset/common_criteria.py
@@ -12,6 +12,7 @@ import json
 
 import pandas as pd
 from bs4 import Tag, BeautifulSoup
+from tqdm import tqdm
 
 from sec_certs import helpers as helpers, parallel_processing as cert_processing, constants as constants
 from sec_certs.dataset.cpe import CPEDataset
@@ -742,12 +743,12 @@ class CCDataset(Dataset, ComplexSerializableType):
 
         cpe_dset = self.prepare_cpe_dataset()
 
-        for annotation in [x for x in data if 'verified_cpe_match' in x]:
-            match_keys = annotation['verified_cpe_match']['choices']
-            if isinstance(match_keys, str):
-                match_keys = [match_keys]
-            match_keys = [x[1:] for x in match_keys]
-            cpes = itertools.chain.from_iterable([cpe_dset.get_cpes_from_title(annotation[x]) for x in match_keys])
+        logger.info('Translating label studio matches into their CPE representations and assigning to certificates.')
+        for annotation in tqdm([x for x in data if 'verified_cpe_match' in x], desc='Translating label studio matches'):
+            match_keys = annotation['verified_cpe_match']
+            match_keys = [match_keys] if isinstance(match_keys, str) else match_keys['choices']
+            match_keys = [x.lstrip('$') for x in match_keys]
+            cpes = list(itertools.chain.from_iterable([cpe_dset.get_cpes_from_title(annotation[x]) for x in match_keys]))
             certs = self.get_certs_from_name(annotation['text'])
 
             for c in certs: