up

app.py

@@ -74,6 +74,7 @@ def match_code(lines, codewords):
 
 def main(filename, codewords_mapping):
     lines = retrieve_lines(filename)
+    files = []
 
     for label, codewords in codewords_mapping.items():
         match = match_code(lines, codewords)
@@ -98,28 +99,32 @@ def main(filename, codewords_mapping):
             with open(result_file, "a") as f:
                 f.write(out)
 
+        files.append(result_file)
+
+    return files
+
+
 def convert(*keywords):
-    # cached_folder = snapshot_download("claudiag/atlas", token=os.environ.get("HF_TOKEN"))
+    cached_folder = snapshot_download("claudiag/atlas", token=os.environ.get("HF_TOKEN"))
     codewords_mapping = {k: v for k,v in zip(CATEGORIES, keywords)}
 
     num_files = 0
 
-    return "_".join(codewords_mapping.values())
-
-    for folder in tqdm.tqdm(glob.glob("./*")):
-        shutil.rmtree(RESULTS_FOLDER, ignore_errors=True)
-        os.makedirs(RESULTS_FOLDER)
+    shutil.rmtree(RESULTS_FOLDER, ignore_errors=True)
+    os.makedirs(RESULTS_FOLDER)
 
+    result_files = []
+    for folder in tqdm.tqdm(glob.glob(os.path.join(cached_folder, "/*"))):
         all_files = tqdm.tqdm(glob.glob(f"./{folder}/*"))
         num_files += len(all_files)
 
         for filename in all_files:
             try:
-                main(filename)
+                result_files += main(filename)
             except Exception as e:
                 print(f"{filename} not working because \n {e}")
 
-        return f"Retrieved from {num_files}"
+            return f"Retrieved from {num_files}"
 
 inputs = [gr.Textbox(label=f"Enter your keywords for {k}", max_lines=2, placeholder=CAT_TO_CODEWORDS[k]) for k in CATEGORIES]