refresh

.gitignore

@@ -2,4 +2,5 @@
 *.json
 data
 .ipynb_checkpoints
-__pycache__
+__pycache__
+.sesskey

app.py

@@ -27,14 +27,13 @@ app, rt = fast_app(html_style=(style,))
 login(token=os.environ.get("HF_TOKEN"))
 
 hf_user = whoami(os.environ.get("HF_TOKEN"))["name"]
-HF_REPO_ID_TXT = f"{hf_user}/zotero-answer-ai-article-texts"
-HF_REPO_ID_IMG = f"{hf_user}/zotero-answer-ai-article-images"
+HF_REPO_ID_TXT = f"{hf_user}/zotero-answer-ai-texts"
+HF_REPO_ID_IMG = f"{hf_user}/zotero-answer-ai-images"
 
 abstract_ds = load_dataset(HF_REPO_ID_TXT, "abstracts", split="train")
 article_ds = load_dataset(HF_REPO_ID_TXT, "articles", split="train")
 
-image_ds = load_dataset(HF_REPO_ID_IMG, "images", split="train")
-image_ds = image_ds.filter(lambda x: x["page_number"] == 1)
+image_ds = load_dataset(HF_REPO_ID_IMG, "images_first_page", split="train")
 
 
 def parse_date(date_string):
@@ -56,11 +55,21 @@ for article in article_ds:
 
 weeks = sorted(week2articles.keys(), reverse=True)
 
+arxiv2article = {article["arxiv_id"]: article for article in article_ds}
+arxiv2abstract = {abstract["arxiv_id"]: abstract for abstract in abstract_ds}
+arxiv2image = {image["arxiv_id"]: image for image in image_ds}
+
+# def get_article_details(arxiv_id):
+#     article = article_ds.filter(lambda x: x["arxiv_id"] == arxiv_id)[0]
+#     abstract = abstract_ds.filter(lambda x: x["arxiv_id"] == arxiv_id)
+#     image = image_ds.filter(lambda x: x["arxiv_id"] == arxiv_id)
+#     return article, abstract, image
+
 
 def get_article_details(arxiv_id):
-    article = article_ds.filter(lambda x: x["arxiv_id"] == arxiv_id)[0]
-    abstract = abstract_ds.filter(lambda x: x["arxiv_id"] == arxiv_id)
-    image = image_ds.filter(lambda x: x["arxiv_id"] == arxiv_id)
+    article = arxiv2article.get(arxiv_id, {})
+    abstract = arxiv2abstract.get(arxiv_id, {})
+    image = arxiv2image.get(arxiv_id, {})
     return article, abstract, image
 
 
@@ -103,7 +112,7 @@ def generate_week_content(current_week):
         ]
 
         if image:
-            pil_image = image[0]["image"]
+            pil_image = image["image"]  # image[0]["image"]
             img_byte_arr = BytesIO()
             pil_image.save(img_byte_arr, format="JPEG")
             img_byte_arr = img_byte_arr.getvalue()

main.py

@@ -1,5 +1,6 @@
 import os
 import re
+import shutil
 import time
 
 import dotenv
@@ -19,8 +20,8 @@ dotenv.load_dotenv()
 login(token=os.environ.get("HF_TOKEN"))
 
 hf_user = whoami(os.environ.get("HF_TOKEN"))["name"]
-HF_REPO_ID_TXT = f"{hf_user}/zotero-answer-ai-article-texts"
-HF_REPO_ID_IMG = f"{hf_user}/zotero-answer-ai-article-images"
+HF_REPO_ID_TXT = f"{hf_user}/zotero-answer-ai-texts"
+HF_REPO_ID_IMG = f"{hf_user}/zotero-answer-ai-images"
 
 
 ########################################################
@@ -67,7 +68,7 @@ def get_zotero_items(debug=False):
         print(f"# items fetched {len(items)}")
 
         if debug:
-            if len(items) > 1500:
+            if len(items) > 1600:
                 break
 
     return items
@@ -334,7 +335,7 @@ def download_arxiv_pdf(arxiv_id):
         raise Exception(f"Failed to download PDF. Status code: {response.status_code}")
 
 
-def pdf_to_jpegs(pdf_content, output_folder):
+def pdf_to_jpegs(pdf_content, output_folder, max_pages=128):
     # Create output folder if it doesn't exist
     os.makedirs(output_folder, exist_ok=True)
 
@@ -353,6 +354,9 @@ def pdf_to_jpegs(pdf_content, output_folder):
         pix.save(image_path)
         # print(f"Saved {image_path}")
 
+        if page_num >= max_pages:
+            break
+
     doc.close()
 
 
@@ -444,6 +448,13 @@ def upload_to_hf(abstract_df, contents_df, processed_arxiv_ids):
     except Exception as e:
         print(e)
 
+    # upload first pages only
+    try:
+        img_ds = img_ds.filter(lambda x: x["page_number"] == 1)
+        img_ds.push_to_hub(HF_REPO_ID_IMG, "images_first_page", token=os.environ.get("HF_TOKEN"))
+    except Exception as e:
+        print(e)
+
     try:
         # push id_to_abstract
         abstract_ds = Dataset.from_pandas(abstract_df)
@@ -479,11 +490,8 @@ def main():
         existing_arxiv_ids = load_dataset(HF_REPO_ID_TXT, "processed_arxiv_ids")["train"]["arxiv_id"]
     except Exception as e:
         print(e)
-        try:
-            existing_arxiv_ids = srsly.read_json("data/processed_arxiv_ids.json")
-        except Exception as e:
-            print(e)
-            existing_arxiv_ids = []
+        existing_arxiv_ids = []
+
     existing_arxiv_ids = set(existing_arxiv_ids)
     print(f"# of existing arxiv ids: {len(existing_arxiv_ids)}")
 
@@ -492,9 +500,20 @@ def main():
     arxiv_items = fetch_arxiv_htmls(arxiv_items)
     print(f"# of new arxiv items: {len(arxiv_items)}")
 
+    if len(arxiv_items) == 0:
+        print("No new arxiv items to process")
+        return
+
     processed_arxiv_ids = set()
     pbar = tqdm(range(len(arxiv_items)))
 
+    # remove "data" directory if it exists
+    if os.path.exists("data"):
+        try:
+            shutil.rmtree("data")
+        except Exception as e:
+            print(e)
+
     for item in arxiv_items:
         # download images --
         save_arxiv_article_images(item["arxiv_id"])