refresh

main.py

@@ -66,7 +66,7 @@ def get_zotero_items(debug=False):
         print(f"# items fetched {len(items)}")
 
         if debug:
-            if len(items) > 200:
+            if len(items) > 500:
                 break
 
     return items
@@ -153,9 +153,7 @@ def parse_html_content(html):
 
     # Extract paper title
     try:
-        paper_title = soup.find("h1", class_="ltx_title ltx_title_document").get_text(
-            strip=True
-        )
+        paper_title = soup.find("h1", class_="ltx_title ltx_title_document").get_text(strip=True)
     except Exception:
         paper_title = soup.find("title").get_text(strip=True)
         paper_title = re.sub(r"^\[\d+\.\d+(v\d+)?\]\s*", "", paper_title)
@@ -170,9 +168,7 @@ def parse_html_content(html):
     if abstract:
         result.append(
             {
-                "content": " ".join(
-                    p.get_text(strip=True) for p in abstract.find_all("p")
-                ).replace(")", ") "),
+                "content": " ".join(p.get_text(strip=True) for p in abstract.find_all("p")).replace(")", ") "),
                 "title": "Abstract",
                 "paper_title": paper_title,
                 "content_type": "abstract",
@@ -182,11 +178,7 @@ def parse_html_content(html):
     sections = soup.find_all("section", class_="ltx_section")
     for index, section in enumerate(sections):
         section_title = section.find("h2", class_="ltx_title ltx_title_section")
-        section_title = (
-            section_title.get_text(strip=True)
-            if section_title
-            else f"Section {index + 1}"
-        )
+        section_title = section_title.get_text(strip=True) if section_title else f"Section {index + 1}"
         section_content = section.get_text(strip=True).replace(")", ") ")
 
         content_type = "body"
@@ -281,9 +273,7 @@ def parse_markdown_content(md_content, arxiv_id):
                             "content": " ".join(content),
                             "title": current_title,
                             "paper_title": paper_title,
-                            "content_type": get_content_type(
-                                current_section, len(parsed)
-                            ),
+                            "content_type": get_content_type(current_section, len(parsed)),
                             "arxiv_id": arxiv_id,
                         }
                     )
@@ -393,13 +383,7 @@ def create_hf_image_dataset(base_dir):
 
                 # Add the data
                 data.append(
-                    {
-                        "image": image_path,
-                        "arxiv_id": arxiv_id,
-                        "page_number": page_number,
-                        "width": width,
-                        "height": height,
-                    }
+                    {"image": image_path, "arxiv_id": arxiv_id, "page_number": page_number, "width": width, "height": height}
                 )
 
     # Create the dataset
@@ -435,23 +419,24 @@ def upload_to_hf(abstract_df, contents_df, processed_arxiv_ids):
     )
 
     # upload image dataset
-    img_ds = create_hf_image_dataset("data/arxiv_images")
-    img_ds.push_to_hub(repo_id, "images", token=os.environ.get("HF_TOKEN"))
-
-    # push id_to_abstract
-    abstract_ds = Dataset.from_pandas(abstract_df)
-    abstract_ds.push_to_hub(repo_id, "abstracts", token=os.environ.get("HF_TOKEN"))
-
-    # push arxiv_items
-    arxiv_ds = Dataset.from_pandas(contents_df)
-    arxiv_ds.push_to_hub(repo_id, "articles", token=os.environ.get("HF_TOKEN"))
-
-    # push processed_arxiv_ids
-    processed_arxiv_ids = [{"arxiv_id": arxiv_id} for arxiv_id in processed_arxiv_ids]
-    processed_arxiv_ids_ds = Dataset.from_list(processed_arxiv_ids)
-    processed_arxiv_ids_ds.push_to_hub(
-        repo_id, "processed_arxiv_ids", token=os.environ.get("HF_TOKEN")
-    )
+    try:
+        img_ds = create_hf_image_dataset("data/arxiv_images")
+        img_ds.push_to_hub(repo_id, "images", token=os.environ.get("HF_TOKEN"))
+
+        # push id_to_abstract
+        abstract_ds = Dataset.from_pandas(abstract_df)
+        abstract_ds.push_to_hub(repo_id, "abstracts", token=os.environ.get("HF_TOKEN"))
+
+        # push arxiv_items
+        arxiv_ds = Dataset.from_pandas(contents_df)
+        arxiv_ds.push_to_hub(repo_id, "articles", token=os.environ.get("HF_TOKEN"))
+
+        # push processed_arxiv_ids
+        processed_arxiv_ids = [{"arxiv_id": arxiv_id} for arxiv_id in processed_arxiv_ids]
+        processed_arxiv_ids_ds = Dataset.from_list(processed_arxiv_ids)
+        processed_arxiv_ids_ds.push_to_hub(repo_id, "processed_arxiv_ids", token=os.environ.get("HF_TOKEN"))
+    except Exception as e:
+        print(e)
 
 
 ########################################################
@@ -467,9 +452,7 @@ def main():
 
     # get already processed arxiv ids from HF
     try:
-        existing_arxiv_ids = load_dataset(HF_REPO_ID, "processed_arxiv_ids")["train"][
-            "arxiv_id"
-        ]
+        existing_arxiv_ids = load_dataset(HF_REPO_ID, "processed_arxiv_ids")["train"]["arxiv_id"]
     except Exception as e:
         print(e)
         try:
@@ -481,9 +464,7 @@ def main():
     print(f"# of existing arxiv ids: {len(existing_arxiv_ids)}")
 
     # new arxiv items
-    arxiv_items = [
-        item for item in arxiv_items if item["arxiv_id"] not in existing_arxiv_ids
-    ]
+    arxiv_items = [item for item in arxiv_items if item["arxiv_id"] not in existing_arxiv_ids]
     arxiv_items = fetch_arxiv_htmls(arxiv_items)
     print(f"# of new arxiv items: {len(arxiv_items)}")
 
@@ -521,11 +502,7 @@ def main():
                 id_to_abstract[item["arxiv_id"]] = entry["content"]
                 break
     print(f"# of abstracts: {len(id_to_abstract)}")
-    abstract_df = (
-        pd.Series(id_to_abstract)
-        .reset_index()
-        .rename(columns={"index": "arxiv_id", 0: "abstract"})
-    )
+    abstract_df = pd.Series(id_to_abstract).reset_index().rename(columns={"index": "arxiv_id", 0: "abstract"})
     print(abstract_df.head())
 
     # add to existing dataset
@@ -537,9 +514,7 @@ def main():
     print(old_abstract_df.head())
 
     abstract_df = pd.concat([old_abstract_df, abstract_df]).reset_index(drop=True)
-    abstract_df = abstract_df.drop_duplicates(
-        subset=["arxiv_id"], keep="last"
-    ).reset_index(drop=True)
+    abstract_df = abstract_df.drop_duplicates(subset=["arxiv_id"], keep="last").reset_index(drop=True)
 
     # contents
     contents_df = pd.DataFrame(arxiv_items)
@@ -553,9 +528,7 @@ def main():
         print(old_contents_df.sample().T)
 
     contents_df = pd.concat([old_contents_df, contents_df]).reset_index(drop=True)
-    contents_df = contents_df.drop_duplicates(
-        subset=["arxiv_id"], keep="last"
-    ).reset_index(drop=True)
+    contents_df = contents_df.drop_duplicates(subset=["arxiv_id"], keep="last").reset_index(drop=True)
 
     # upload to hf
     processed_arxiv_ids = list(set(processed_arxiv_ids + list(processed_arxiv_ids)))