Fix dict key mismatch

app.py

@@ -33,14 +33,6 @@ def get_pipeline():
 @spaces.GPU
 def classify_abstracts(abstracts:Dict[str, str],batch_size=64, progress=gr.Progress()) -> None:
     pipe = get_pipeline()
-    # pmids = list(abstracts.keys())
-    # batch_size = 64
-    # classification = []
-    # abstracts_list = list(abstracts.values())
-    # for i in range(0, len(abstracts), batch_size):
-    #     classification.extend(pipe(abstracts_list[i:i+batch_size]))
-
-    
         
     # return classification
     results = []
@@ -67,7 +59,7 @@ def classify_abstracts(abstracts:Dict[str, str],batch_size=64, progress=gr.Progr
             for pmid, classification in zip(batch_pmids, classifications):
                 results.append({
                     'pmid': pmid,
-                    'classification': classification['label'],
+                    'classification': label_lookup[classification['label']],
                     'score': classification['score']
                 })
             
@@ -78,9 +70,7 @@ def classify_abstracts(abstracts:Dict[str, str],batch_size=64, progress=gr.Progr
         except Exception as e:
             print(f"Error classifying batch starting at index {i}: {str(e)}")
             continue
-    for pmid, abs in zip(pmids, results):
-        abs['label'] = label_lookup[abs['label']]
-        abs['pmid'] = pmid
+
     progress(1.0, desc="Classification complete!")
     return results