app

main.py

@@ -9,7 +9,7 @@ import requests
 import schedule
 import srsly
 from bs4 import BeautifulSoup
-from datasets import Dataset, Image, load_dataset
+from datasets import Dataset, Image, concatenate_datasets, load_dataset
 from huggingface_hub import create_repo, login, whoami
 from PIL import Image as PILImage
 from retry import retry
@@ -19,7 +19,8 @@ dotenv.load_dotenv()
 login(token=os.environ.get("HF_TOKEN"))
 
 hf_user = whoami(os.environ.get("HF_TOKEN"))["name"]
-HF_REPO_ID = f"{hf_user}/zotero-articles"
+HF_REPO_ID_TXT = f"{hf_user}/zotero-answer-ai-article-texts"
+HF_REPO_ID_IMG = f"{hf_user}/zotero-answer-ai-article-images"
 
 
 ########################################################
@@ -66,7 +67,7 @@ def get_zotero_items(debug=False):
         print(f"# items fetched {len(items)}")
 
         if debug:
-            if len(items) > 500:
+            if len(items) > 1500:
                 break
 
     return items
@@ -103,11 +104,18 @@ def get_arxiv_items(items):
             if arxiv_id in visited:
                 continue
 
+            authors = []
+            for author in data.get("creators", []):
+                authors.append(f"{author.get('firstName', '')} {author.get('lastName', '')}")
+
             arxiv_items.append(
                 {
                     "arxiv_id": arxiv_id,
                     "arxiv_url": arxiv_url,
+                    "title": data.get("title", ""),
+                    "authors": authors,
                     "pdf_url": pdf_url,
+                    "date_published": data.get("date", ""),
                     "added_by": item["meta"]["createdByUser"]["username"],
                     "date_added": data.get("dateAdded", ""),
                 }
@@ -129,10 +137,10 @@ def fetch_arxiv_htmls(arxiv_items):
     for item in tqdm(arxiv_items):
         html = fetch_arxiv_html(item["arxiv_id"])
         if html:
-            item["raw_html"] = html
+            item["raw_content"] = html
         else:
             print(f"failed to fetch html for {item['arxiv_id']}")
-            item["raw_html"] = "Error"
+            item["raw_content"] = "Error"
 
     return arxiv_items
 
@@ -392,8 +400,6 @@ def create_hf_image_dataset(base_dir):
             "image": [d["image"] for d in data],
             "arxiv_id": [d["arxiv_id"] for d in data],
             "page_number": [d["page_number"] for d in data],
-            "width": [d["width"] for d in data],
-            "height": [d["height"] for d in data],
         }
     )
 
@@ -409,9 +415,17 @@ def create_hf_image_dataset(base_dir):
 
 
 def upload_to_hf(abstract_df, contents_df, processed_arxiv_ids):
-    repo_id = HF_REPO_ID
+    # repo_id = HF_REPO_ID
+    create_repo(
+        repo_id=HF_REPO_ID_TXT,
+        token=os.environ.get("HF_TOKEN"),
+        private=True,
+        repo_type="dataset",
+        exist_ok=True,
+    )
+
     create_repo(
-        repo_id=repo_id,
+        repo_id=HF_REPO_ID_IMG,
         token=os.environ.get("HF_TOKEN"),
         private=True,
         repo_type="dataset",
@@ -421,20 +435,28 @@ def upload_to_hf(abstract_df, contents_df, processed_arxiv_ids):
     # upload image dataset
     try:
         img_ds = create_hf_image_dataset("data/arxiv_images")
-        img_ds.push_to_hub(repo_id, "images", token=os.environ.get("HF_TOKEN"))
+        try:
+            old_img_ds = load_dataset(HF_REPO_ID_IMG, "images")["train"]
+            img_ds = concatenate_datasets([old_img_ds, img_ds])
+        except Exception as e:
+            print(e)
+        img_ds.push_to_hub(HF_REPO_ID_IMG, "images", token=os.environ.get("HF_TOKEN"))
+    except Exception as e:
+        print(e)
 
+    try:
         # 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"))
+        abstract_ds.push_to_hub(HF_REPO_ID_TXT, "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"))
+        arxiv_ds.push_to_hub(HF_REPO_ID_TXT, "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"))
+        processed_arxiv_ids_ds.push_to_hub(HF_REPO_ID_TXT, "processed_arxiv_ids", token=os.environ.get("HF_TOKEN"))
     except Exception as e:
         print(e)
 
@@ -454,7 +476,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_TXT, "processed_arxiv_ids")["train"]["arxiv_id"]
     except Exception as e:
         print(e)
         try:
@@ -471,14 +493,15 @@ def main():
     print(f"# of new arxiv items: {len(arxiv_items)}")
 
     processed_arxiv_ids = set()
+    pbar = tqdm(range(len(arxiv_items)))
+
     for item in arxiv_items:
         # download images --
         save_arxiv_article_images(item["arxiv_id"])
 
         # parse html
         try:
-            item["contents"] = parse_html_content(item["raw_html"])
-            processed_arxiv_ids.add(item["arxiv_id"])
+            item["contents"] = parse_html_content(item["raw_content"])
         except Exception as e:
             print(f"Failed to parse html for {item['arxiv_id']}: {e}")
             item["contents"] = []
@@ -486,12 +509,21 @@ def main():
         if len(item["contents"]) == 0:
             print("Extracting from pdf...")
             md_content = get_pdf_text(item["arxiv_id"])  # fix this
+            item["raw_content"] = md_content
+
             if md_content:
                 item["contents"] = parse_markdown_content(md_content, item["arxiv_id"])
-                processed_arxiv_ids.add(item["arxiv_id"])
             else:
                 item["contents"] = []
 
+        if len(item["contents"]) > 0:
+            processed_arxiv_ids.add(item["arxiv_id"])
+            if len(item["authors"]) == 0:
+                item["authors"] = []  # ["unknown"]
+                item["title"] = item["contents"][0]["paper_title"]
+        pbar.update(1)
+    pbar.close()
+
     # save contents ---
     processed_arxiv_ids = list(processed_arxiv_ids)
     print(f"# of processed arxiv ids: {len(processed_arxiv_ids)}")
@@ -509,7 +541,7 @@ def main():
 
     # add to existing dataset
     try:
-        old_abstract_df = load_dataset(HF_REPO_ID, "abstracts")["train"].to_pandas()
+        old_abstract_df = load_dataset(HF_REPO_ID_TXT, "abstracts")["train"].to_pandas()
     except Exception as e:
         print(e)
         old_abstract_df = pd.DataFrame(columns=abstract_df.columns)
@@ -522,7 +554,7 @@ def main():
     contents_df = pd.DataFrame(arxiv_items)
     print(contents_df.head())
     try:
-        old_contents_df = load_dataset(HF_REPO_ID, "articles")["train"].to_pandas()
+        old_contents_df = load_dataset(HF_REPO_ID_TXT, "articles")["train"].to_pandas()
     except Exception as e:
         print(e)
         old_contents_df = pd.DataFrame(columns=contents_df.columns)

requirements.txt

@@ -12,4 +12,5 @@ retry
 pandas
 datasets
 PyMuPDF
-pillow
+pillow
+tqdm