data ingest

.gitignore

@@ -0,0 +1,4 @@
+.env
+*.json
+data
+.ipynb_checkpoints

main.py

@@ -0,0 +1,435 @@
+import os
+import re
+
+import dotenv
+import pandas as pd
+import requests
+import srsly
+from bs4 import BeautifulSoup
+from datasets import Dataset, load_dataset
+from huggingface_hub import create_repo, login, whoami
+from retry import retry
+from tqdm.auto import tqdm
+
+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-answer-ai-articles"
+
+
+########################################################
+### GET ZOTERO ITEMS
+########################################################
+
+
+@retry(tries=3, delay=8)
+def _fetch_one_zotero_batch(url, headers, params):
+    """
+    Fetch articles from Zotero API
+    """
+    response = requests.get(url, headers=headers, params=params)
+    response.raise_for_status()
+    return response.json()
+
+
+def get_zotero_items(debug=False):
+    """
+    fetch items from zotero library
+    """
+
+    GROUP_ID = os.getenv("GROUP_ID")
+    API_KEY = os.getenv("API_KEY")
+    BASE_URL = f"https://api.zotero.org/groups/{GROUP_ID}/items"
+    LIMIT = 100
+
+    headers = {"Zotero-API-Key": API_KEY, "Content-Type": "application/json"}
+
+    items = []
+    start = 0
+
+    i = 1
+    while True:
+        i += 1
+        params = {"limit": LIMIT, "start": start}
+        page_items = _fetch_one_zotero_batch(BASE_URL, headers, params)
+
+        if not page_items:
+            break
+
+        items.extend(page_items)
+        start += LIMIT
+        print(f"# items fetched {len(items)}")
+
+        if debug:
+            if len(items) > 300:
+                break
+
+    return items
+
+
+########################################################
+### EXTRACT ARXIV LINKS AND PDFs
+########################################################
+
+
+def get_arxiv_items(items):
+    visited = set()
+
+    arxiv_items = []
+    arxiv_pattern = re.compile(r"arxiv.org/abs/(\d+\.\d+)")
+
+    for item in items:
+        data = item.get("data", {})
+        attachments = item.get("links", {}).get("attachment", {})
+
+        arxiv_url = None
+        pdf_url = None
+
+        if "url" in data and "arxiv.org" in data["url"]:
+            arxiv_match = arxiv_pattern.search(data["url"])
+            if arxiv_match:
+                arxiv_url = data["url"]
+
+        if attachments:
+            pdf_url = attachments["href"]
+
+        if arxiv_url:
+            arxiv_id = arxiv_url.split("/")[-1]
+            if arxiv_id in visited:
+                continue
+
+            arxiv_items.append(
+                {
+                    "arxiv_id": arxiv_id,
+                    "arxiv_url": arxiv_url,
+                    "pdf_url": pdf_url,
+                    "added_by": item["meta"]["createdByUser"]["username"],
+                    "date_added": data.get("dateAdded", ""),
+                }
+            )
+
+            visited.add(arxiv_id)
+
+    return arxiv_items
+
+
+@retry(tries=3, delay=15, backoff=2)
+def fetch_arxiv_html(arxiv_id):
+    url = f"https://ar5iv.labs.arxiv.org/html/{arxiv_id.split('v')[0]}"
+    response = requests.get(url)
+    return response.text if response.status_code == 200 else None
+
+
+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
+        else:
+            print(f"failed to fetch html for {item['arxiv_id']}")
+            item["raw_html"] = "Error"
+
+    return arxiv_items
+
+
+########################################################
+### PARSE CONTENT FROM ARXIV HTML #
+########################################################
+
+
+def parse_html_content(html):
+    """
+    Parse content from arxiv html
+    """
+    arxiv_id_match = re.search(r"\[(\d+\.\d+(v\d+)?)\]", html)
+    arxiv_id = arxiv_id_match.group(1) if arxiv_id_match else None
+    soup = BeautifulSoup(html, "html.parser")
+    result = []
+
+    # Extract paper title
+    try:
+        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)
+
+    for math in soup.find_all("math"):
+        math.decompose()
+    for cite in soup.find_all("cite"):
+        cite.decompose()
+
+    # Extract abstract
+    abstract = soup.find("div", class_="ltx_abstract")
+    if abstract:
+        result.append(
+            {
+                "content": " ".join(p.get_text(strip=True) for p in abstract.find_all("p")).replace(")", ") "),
+                "title": "Abstract",
+                "paper_title": paper_title,
+                "content_type": "abstract",
+            }
+        )
+    # Extract sections
+    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_content = section.get_text(strip=True).replace(")", ") ")
+
+        content_type = "body"
+        if index == 0:
+            content_type = "introduction"
+        elif index == len(sections) - 1:
+            content_type = "conclusion"
+
+        result.append(
+            {
+                "content": section_content,
+                "title": section_title,
+                "paper_title": paper_title,
+                "content_type": content_type,
+            }
+        )
+
+    for c in result:
+        c["arxiv_id"] = arxiv_id
+
+    return result
+
+
+########################################################
+### GET TEXTS FROM PDF & PARSE
+########################################################
+
+
+def get_pdf_text(arxiv_id):
+    url = "http://147.189.194.113:80/extract"  # fix: currently down
+
+    try:
+        response = requests.get(url, params={"arxiv_id": arxiv_id})
+        response = response.json()
+        if "text" in response:
+            return response["text"]
+        return None
+    except Exception as e:
+        print(e)
+        return None
+
+
+def get_content_type(section_type, section_count):
+    """Determine the content type based on the section type and count"""
+    if section_type == "abstract":
+        return "abstract"
+    elif section_type == "introduction" or section_count == 1:
+        return "introduction"
+    elif section_type == "conclusion" or section_type == "references":
+        return section_type
+    else:
+        return "body"
+
+
+def get_section_type(title):
+    """Determine the section type based on the title"""
+    title_lower = title.lower()
+    if "abstract" in title_lower:
+        return "abstract"
+    elif "introduction" in title_lower:
+        return "introduction"
+    elif "conclusion" in title_lower:
+        return "conclusion"
+    elif "reference" in title_lower:
+        return "references"
+    else:
+        return "body"
+
+
+def parse_markdown_content(md_content, arxiv_id):
+    """
+    Parses markdown content to identify and extract sections based on headers.
+    """
+
+    lines = md_content.split("\n")
+    parsed = []
+    current_section = None
+    content = []
+    paper_title = None
+    current_title = None
+
+    # identify sections based on headers
+    for line in lines:
+        if line.startswith("#"):
+            if paper_title is None:
+                paper_title = line.lstrip("#").strip()
+                continue
+            if content:
+                if current_title:
+                    parsed.append(
+                        {
+                            "content": " ".join(content),
+                            "title": current_title,
+                            "paper_title": paper_title,
+                            "content_type": get_content_type(current_section, len(parsed)),
+                            "arxiv_id": arxiv_id,
+                        }
+                    )
+                content = []
+
+            current_title = line.lstrip("#").lstrip("#").lstrip()
+            if "bit" not in current_title:
+                current_title = (
+                    current_title.lstrip("123456789")
+                    .lstrip()
+                    .lstrip(".")
+                    .lstrip()
+                    .lstrip("123456789")
+                    .lstrip()
+                    .lstrip(".")
+                    .lstrip()
+                )
+            current_section = get_section_type(current_title)
+
+        else:
+            content.append(line)
+
+    # Add the last section
+    if content and current_title:
+        parsed.append(
+            {
+                "content": " ".join(content).replace(")", ") "),
+                "title": current_title,
+                "paper_title": paper_title,
+                "content_type": get_content_type(current_section, len(parsed)),
+                "arxiv_id": arxiv_id,
+            }
+        )
+
+    return parsed
+
+
+########################################################
+### HF UPLOAD
+########################################################
+
+
+def upload_to_hf(abstract_df, contents_df, processed_arxiv_ids):
+    repo_id = HF_REPO_ID
+    create_repo(
+        repo_id=repo_id,
+        token=os.environ.get("HF_TOKEN"),
+        private=True,
+        repo_type="dataset",
+        exist_ok=True,
+    )
+
+    # 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"))
+
+
+########################################################
+### MAIN
+########################################################
+
+
+def main():
+    items = get_zotero_items(debug=True)
+    print(f"# of items fetched from zotero: {len(items)}")
+    arxiv_items = get_arxiv_items(items)
+    print(f"# of arxiv papers: {len(arxiv_items)}")
+
+    # get already processed arxiv ids from HF
+    try:
+        existing_arxiv_ids = load_dataset(HF_REPO_ID, "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 = set(existing_arxiv_ids)
+    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 = fetch_arxiv_htmls(arxiv_items)
+    print(f"# of new arxiv items: {len(arxiv_items)}")
+
+    processed_arxiv_ids = set()
+    for item in arxiv_items:
+        try:
+            item["contents"] = parse_html_content(item["raw_html"])
+            processed_arxiv_ids.add(item["arxiv_id"])
+        except Exception as e:
+            print(f"Failed to parse html for {item['arxiv_id']}: {e}")
+            print("Extracting from pdf...")
+            md_content = get_pdf_text(item["arxiv_id"])
+            if md_content:
+                item["contents"] = parse_markdown_content(md_content, item["arxiv_id"])
+                processed_arxiv_ids.add(item["arxiv_id"])
+            else:
+                item["contents"] = []
+
+    # save contents ---
+    processed_arxiv_ids = list(processed_arxiv_ids)
+    print(f"# of processed arxiv ids: {len(processed_arxiv_ids)}")
+
+    # save abstracts ---
+    id_to_abstract = {}
+    for item in arxiv_items:
+        for entry in item["contents"]:
+            if entry["content_type"] == "abstract":
+                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"})
+    print(abstract_df.head())
+
+    # add to existing dataset
+    try:
+        old_abstract_df = load_dataset(HF_REPO_ID, "abstracts")["train"].to_pandas()
+    except Exception as e:
+        print(e)
+        old_abstract_df = pd.DataFrame(columns=abstract_df.columns)
+    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)
+
+    # contents
+    contents_df = pd.DataFrame(arxiv_items)
+    print(contents_df.head())
+    try:
+        old_contents_df = load_dataset(HF_REPO_ID, "articles")["train"].to_pandas()
+    except Exception as e:
+        print(e)
+        old_contents_df = pd.DataFrame(columns=contents_df.columns)
+    if len(old_contents_df) > 0:
+        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)
+
+    # upload to hf
+    processed_arxiv_ids = list(set(processed_arxiv_ids + list(processed_arxiv_ids)))
+    upload_to_hf(abstract_df, contents_df, processed_arxiv_ids)
+
+    # save as local copy
+    os.makedirs("data", exist_ok=True)
+    abstract_df.to_parquet("data/abstracts.parquet")
+    contents_df.to_parquet("data/contents.parquet")
+    srsly.write_json("data/processed_arxiv_ids.json", processed_arxiv_ids)
+
+
+if __name__ == "__main__":
+    main()

ruff.toml

@@ -0,0 +1,3 @@
+line-length = 128
+target-version = "py311"
+ignore = ["E402"]

zotero_refactor.ipynb

@@ -0,0 +1,746 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "bfd057e4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%ai reset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "1c3aa2dc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%ai -c\n",
+    "I'm refactoring the following code to ingest and process items from a zotero lib.\n",
+    "\n",
+    "```python\n",
+    "import re\n",
+    "\n",
+    "import requests\n",
+    "import srsly\n",
+    "from bs4 import BeautifulSoup\n",
+    "from retry import retry\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "\n",
+    "\n",
+    "def get_pdf_text(arxiv_id):\n",
+    "    url = \"http://147.189.194.113:80/extract\"\n",
+    "\n",
+    "    try:\n",
+    "        response = requests.get(url, params={\"arxiv_id\": arxiv_id})\n",
+    "        response = response.json()\n",
+    "        if \"text\" in response:\n",
+    "            return response['text']\n",
+    "        return None\n",
+    "    except Exception as e:\n",
+    "        print(e)\n",
+    "        return None\n",
+    "    \n",
+    "\n",
+    "\n",
+    "def get_zotero_items():\n",
+    "    \"\"\"fetch all items from zotero library\"\"\"\n",
+    "\n",
+    "    GROUP_ID = \"5004697\"\n",
+    "    API_KEY = \"HFfTT8uuL7z9tN0L5VCHe7Ow\"\n",
+    "    BASE_URL = f\"https://api.zotero.org/groups/{GROUP_ID}/items\"\n",
+    "    LIMIT = 100\n",
+    "\n",
+    "    headers = {\"Zotero-API-Key\": API_KEY, \"Content-Type\": \"application/json\"}\n",
+    "\n",
+    "    items = []\n",
+    "    start = 0\n",
+    "\n",
+    "    ITEM_TYPES = [\n",
+    "        \"journalArticle\",\n",
+    "        \"report\",\n",
+    "        \"conferencePaper\",\n",
+    "        \"manuscript\",\n",
+    "        \"book\",\n",
+    "        \"thesis\",\n",
+    "    ]  # Replace or extend with desired item types\n",
+    "\n",
+    "    i = 1\n",
+    "    while True:\n",
+    "        i += 1\n",
+    "        params = {\"limit\": LIMIT, \"start\": start}\n",
+    "        response = requests.get(BASE_URL, headers=headers, params=params)\n",
+    "        page_items = response.json()\n",
+    "        if ITEM_TYPES:\n",
+    "            params[\"itemType\"] = \",\".join(ITEM_TYPES)\n",
+    "\n",
+    "        if not page_items:\n",
+    "            break\n",
+    "\n",
+    "        items.extend(page_items)\n",
+    "        start += LIMIT\n",
+    "        print(f\"# items fetched {len(items)}\")\n",
+    "        # break # debug\n",
+    "\n",
+    "    return items\n",
+    "\n",
+    "\n",
+    "def extract_arxiv_links_and_pdfs(items):\n",
+    "    \"\"\"Extract mapping between arxiv & pdf urls for a list of zotero items\n",
+    "\n",
+    "    Args:\n",
+    "        items (list): list of zotero items\n",
+    "\n",
+    "    Returns:\n",
+    "        arxiv_info (dict): dictionary with arxiv links as keys and pdf links as values\n",
+    "    \"\"\"\n",
+    "    arxiv_info = {}\n",
+    "    arxiv_pattern = re.compile(r\"arxiv.org/abs/(\\d+\\.\\d+)\")\n",
+    "\n",
+    "    for item in items:\n",
+    "        data = item.get(\"data\", {})\n",
+    "        attachments = item.get(\"links\", {}).get(\"attachment\", {})\n",
+    "        arxiv_url = None\n",
+    "        pdf_url = None\n",
+    "\n",
+    "        if \"url\" in data and \"arxiv.org\" in data[\"url\"]:\n",
+    "            arxiv_match = arxiv_pattern.search(data[\"url\"])\n",
+    "            if arxiv_match:\n",
+    "                arxiv_url = data[\"url\"]\n",
+    "\n",
+    "        if attachments:\n",
+    "            pdf_url = attachments[\"href\"]\n",
+    "\n",
+    "        if arxiv_url:\n",
+    "            arxiv_info[arxiv_url] = pdf_url\n",
+    "\n",
+    "    return arxiv_info\n",
+    "\n",
+    "\n",
+    "@retry(tries=3, delay=15, backoff=2)\n",
+    "def get_arxiv_with_retry(arxiv_id):\n",
+    "    return requests.get(f\"https://ar5iv.labs.arxiv.org/html/{arxiv_id}\")\n",
+    "\n",
+    "\n",
+    "def get_arxiv(arxiv_ids):\n",
+    "    \"\"\"\n",
+    "    Fetch html version of arxiv papers\n",
+    "    \"\"\"\n",
+    "    html5ivs = {}\n",
+    "    for i, arxiv_id in tqdm(enumerate(arxiv_ids)):\n",
+    "        if \"v\" in arxiv_id:\n",
+    "            arxiv_id = arxiv_id.split(\"v\")[0]\n",
+    "        response = get_arxiv_with_retry(arxiv_id)\n",
+    "        html5ivs[arxiv_id] = response.text\n",
+    "        if response.status_code != 200:\n",
+    "            html5ivs[arxiv_id] = \"ERROR\"\n",
+    "    return html5ivs\n",
+    "\n",
+    "\n",
+    "def extract_html_content(html):\n",
+    "    \"\"\"\n",
+    "    Extract content from arxiv html\n",
+    "    \"\"\"\n",
+    "    arxiv_id_match = re.search(r\"\\[(\\d+\\.\\d+(v\\d+)?)\\]\", html)\n",
+    "    arxiv_id = arxiv_id_match.group(1) if arxiv_id_match else None\n",
+    "    soup = BeautifulSoup(html, \"html.parser\")\n",
+    "    result = []\n",
+    "\n",
+    "    # Extract paper title\n",
+    "    paper_title = soup.find(\"h1\", class_=\"ltx_title ltx_title_document\").get_text(strip=True)\n",
+    "\n",
+    "    for math in soup.find_all(\"math\"):\n",
+    "        math.decompose()\n",
+    "    for cite in soup.find_all(\"cite\"):\n",
+    "        cite.decompose()\n",
+    "\n",
+    "    # Extract abstract\n",
+    "    abstract = soup.find(\"div\", class_=\"ltx_abstract\")\n",
+    "    if abstract:\n",
+    "        result.append(\n",
+    "            {\n",
+    "                \"content\": \" \".join(p.get_text(strip=True) for p in abstract.find_all(\"p\")).replace(\")\", \") \"),\n",
+    "                \"title\": \"Abstract\",\n",
+    "                \"paper_title\": paper_title,\n",
+    "                \"content_type\": \"abstract\",\n",
+    "            }\n",
+    "        )\n",
+    "\n",
+    "    # Extract introduction\n",
+    "    intro = soup.find(\"section\", id=\"S1\")\n",
+    "    if intro:\n",
+    "        section_title = \"Introduction\"\n",
+    "        paragraphs = intro.find_all(\"div\", class_=\"ltx_para\")\n",
+    "        for para in paragraphs:\n",
+    "            result.append(\n",
+    "                {\n",
+    "                    \"content\": para.get_text(strip=True).replace(\")\", \") \"),\n",
+    "                    \"title\": section_title,\n",
+    "                    \"paper_title\": paper_title,\n",
+    "                    \"content_type\": \"introduction\",\n",
+    "                }\n",
+    "            )\n",
+    "\n",
+    "    # Extract other sections (excluding introduction and conclusion)\n",
+    "    sections = soup.find_all(\"section\", class_=\"ltx_section\")\n",
+    "    for section in sections[1:-1]:\n",
+    "        section_title = section.find(\"h2\", class_=\"ltx_title ltx_title_section\").get_text(strip=True)\n",
+    "        paragraphs = section.find_all(\"div\", class_=\"ltx_para\")\n",
+    "        for para in paragraphs:\n",
+    "            result.append(\n",
+    "                {\n",
+    "                    \"content\": para.get_text(strip=True).replace(\")\", \") \"),\n",
+    "                    \"title\": section_title,\n",
+    "                    \"paper_title\": paper_title,\n",
+    "                    \"content_type\": \"body\",\n",
+    "                }\n",
+    "            )\n",
+    "\n",
+    "    # Extract conclusion\n",
+    "    conclusion = soup.find(\"section\", id=\"S4\")\n",
+    "    if conclusion:\n",
+    "        result.append(\n",
+    "            {\n",
+    "                \"content\": \" \".join(\n",
+    "                    p.get_text(strip=True) for p in conclusion.find_all(\"div\", class_=\"ltx_para\")\n",
+    "                ).replace(\")\", \") \"),\n",
+    "                \"title\": conclusion.find(\"h2\", class_=\"ltx_title ltx_title_section\").get_text(strip=True),\n",
+    "                \"paper_title\": paper_title,\n",
+    "                \"content_type\": \"conclusion\",\n",
+    "            }\n",
+    "        )\n",
+    "\n",
+    "    for c in result:\n",
+    "        c[\"arxiv_id\"] = arxiv_id\n",
+    "\n",
+    "    return result\n",
+    "\n",
+    "\n",
+    "def pool_html_content(parsed, pool_size=3, step_size=2):\n",
+    "    \"\"\"\n",
+    "    Pool together individual bits of contents from \"body\" and \"introduction\" sections where the \"title\" is the same.\n",
+    "    It pools them 'pool_size' by 'pool_size' as a sliding window with a step size of 'step_size'.\n",
+    "    Keeps all other attributes for each element and includes abstract and conclusion without pooling.\n",
+    "\n",
+    "    Args:\n",
+    "    parsed (list): List of dictionaries containing parsed content.\n",
+    "    pool_size (int): Size of the pool for the sliding window.\n",
+    "    step_size (int): Step size for the sliding window.\n",
+    "\n",
+    "    Returns:\n",
+    "    list: A list of pooled contents including non-pooled abstract and conclusion.\n",
+    "    \"\"\"\n",
+    "    pooled_contents = []\n",
+    "    i = 0\n",
+    "    while i < len(parsed) - pool_size + 1:\n",
+    "        current_window = parsed[i : i + pool_size]\n",
+    "        # Check if all entries in the current window have the same title and are from 'body' or 'introduction'\n",
+    "        if all(\n",
+    "            entry[\"content_type\"] in [\"body\", \"introduction\"] and entry[\"title\"] == current_window[0][\"title\"]\n",
+    "            for entry in current_window\n",
+    "        ):\n",
+    "            pooled_content = \"\\n\".join(entry[\"content\"] for entry in current_window)\n",
+    "            if \"bit\" not in current_window[0][\"title\"][:5]:\n",
+    "                title = current_window[0][\"title\"].lstrip(\"0123456789\")\n",
+    "            pooled_content = f\"{current_window[0]['paper_title']}\\n{title}\\n\\n\" + pooled_content\n",
+    "            pooled_entry = {\n",
+    "                \"title\": title,\n",
+    "                \"text\": pooled_content,\n",
+    "                \"paper_title\": current_window[0][\"paper_title\"],\n",
+    "                \"content_type\": current_window[0][\"content_type\"],\n",
+    "                \"arxiv_id\": current_window[0][\"arxiv_id\"],\n",
+    "            }\n",
+    "            pooled_contents.append(pooled_entry)\n",
+    "        i += step_size\n",
+    "\n",
+    "    # Include abstract and conclusion without pooling\n",
+    "    for entry in parsed:\n",
+    "        if entry[\"content_type\"] in [\"abstract\", \"conclusion\"]:\n",
+    "            if \"content\" in entry:\n",
+    "                entry[\"text\"] = entry[\"content\"]\n",
+    "                del entry[\"content\"]\n",
+    "            if \"bit\" not in entry[\"title\"][:5]:\n",
+    "                entry[\"title\"] = entry[\"title\"].lstrip(\"0123456789\")\n",
+    "            entry[\"text\"] = f\"{entry['paper_title']}\\n{entry['title']}\\n\\n\" + entry[\"text\"]\n",
+    "            pooled_contents.append(entry)\n",
+    "\n",
+    "    return pooled_contents\n",
+    "\n",
+    "\n",
+    "def parse_markdown_content(md_content, arxiv_id):\n",
+    "    \"\"\"\n",
+    "    Parses markdown content to identify and extract sections based on headers.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    lines = md_content.split(\"\\n\")\n",
+    "    parsed = []\n",
+    "    current_section = None\n",
+    "    content = []\n",
+    "    paper_title = None\n",
+    "    current_title = None\n",
+    "\n",
+    "    # Identify sections based on headers\n",
+    "    for line in lines:\n",
+    "        if line.startswith(\"#\"):\n",
+    "            if paper_title is None:\n",
+    "                paper_title = line.lstrip(\"#\").lstrip(\"#\").strip()\n",
+    "                continue\n",
+    "            if content:\n",
+    "                if not current_title:\n",
+    "                    print(\"skipping pre-abstract!\")\n",
+    "                parsed.append(\n",
+    "                    {\n",
+    "                        \"title\": current_title,\n",
+    "                        \"text\": content,\n",
+    "                        \"paper_title\": paper_title,\n",
+    "                        \"content_type\": \"body\"\n",
+    "                        if current_section not in [\"Abstract\", \"Introduction\", \"Conclusion\"]\n",
+    "                        else current_section.lower(),\n",
+    "                        \"arxiv_id\": arxiv_id,\n",
+    "                    }\n",
+    "                )\n",
+    "                content = []\n",
+    "\n",
+    "            if \"Introduction\" in line:\n",
+    "                current_section = \"Introduction\"\n",
+    "            elif \"Abstract\" in line:\n",
+    "                current_section = \"Abstract\"\n",
+    "            elif \"Conclusion\" in line:\n",
+    "                current_section = \"Conclusion\"\n",
+    "            elif \"Reference\" in line:\n",
+    "                break\n",
+    "            else:\n",
+    "                current_section = \"body\"\n",
+    "\n",
+    "            current_title = line.lstrip(\"#\").lstrip(\"#\").lstrip()\n",
+    "\n",
+    "            if \"bit\" not in current_title:\n",
+    "                current_title = (\n",
+    "                    current_title.lstrip(\"123456789\")\n",
+    "                    .lstrip()\n",
+    "                    .lstrip(\".\")\n",
+    "                    .lstrip()\n",
+    "                    .lstrip(\"123456789\")\n",
+    "                    .lstrip()\n",
+    "                    .lstrip(\".\")\n",
+    "                    .lstrip()\n",
+    "                )\n",
+    "        else:\n",
+    "            content.append(line)\n",
+    "\n",
+    "    if content:\n",
+    "        parsed.append(\n",
+    "            {\n",
+    "                \"title\": current_section,\n",
+    "                \"text\": content,\n",
+    "                \"paper_title\": paper_title,\n",
+    "                \"content_type\": \"body\"\n",
+    "                if current_section not in [\"Abstract\", \"Introduction\", \"Conclusion\"]\n",
+    "                else current_section.lower(),\n",
+    "                \"arxiv_id\": arxiv_id,\n",
+    "            }\n",
+    "        )\n",
+    "\n",
+    "    return parsed\n",
+    "\n",
+    "\n",
+    "def pool_markdown(parsed, pool_size=3, step_size=2):\n",
+    "    \"\"\"\n",
+    "    pool together markdown contents with sliding window of size 'pool_size' and step size of 'step_size'.\n",
+    "\n",
+    "    Args:\n",
+    "        parsed (list): List of dictionaries containing parsed content.\n",
+    "        pool_size (int): Size of the pool for the sliding window.\n",
+    "        step_size (int): Step size for the sliding window.\n",
+    "\n",
+    "    Returns:\n",
+    "        list: A list of pooled contents including non-pooled abstract and conclusion.\n",
+    "    \"\"\"\n",
+    "    cleaned_contents = []\n",
+    "    for section in parsed:\n",
+    "        cleaned_lines = [line for line in section[\"text\"] if line and \"|\" not in line and len(line) >= 40]\n",
+    "        section[\"text\"] = cleaned_lines\n",
+    "        cleaned_contents.append(section)\n",
+    "\n",
+    "    pooled_contents = []\n",
+    "    for section in cleaned_contents:\n",
+    "        if section[\"content_type\"] in [\"body\", \"introduction\"]:\n",
+    "            i = 0\n",
+    "            while i <= len(section[\"text\"]) - pool_size:\n",
+    "                pooled_content = \"\\n\".join(section[\"text\"][i : i + pool_size])\n",
+    "                pooled_content = f\"{section['paper_title']}\\n{section['title']}\\n\\n\" + pooled_content\n",
+    "                pooled_entry = {\n",
+    "                    \"title\": section[\"title\"],\n",
+    "                    \"text\": pooled_content,\n",
+    "                    \"paper_title\": section[\"paper_title\"],\n",
+    "                    \"content_type\": section[\"content_type\"],\n",
+    "                    \"arxiv_id\": section[\"arxiv_id\"],\n",
+    "                }\n",
+    "                pooled_contents.append(pooled_entry)\n",
+    "                i += step_size\n",
+    "        else:\n",
+    "            section[\"text\"] = f\"{section['paper_title']}\\n\" + \"\\n\".join(section[\"text\"])\n",
+    "            pooled_contents.append(section)\n",
+    "    return pooled_contents\n",
+    "\n",
+    "\n",
+    "def clean_and_pool_markdown_content(md, arxiv_id):\n",
+    "    md = parse_markdown_content(md, arxiv_id)\n",
+    "    return pool_markdown(md)\n",
+    "\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    items = get_zotero_items()\n",
+    "    print(f\"# of items fetched from zotero: {len(items)}\")\n",
+    "    arxiv_info = extract_arxiv_links_and_pdfs(items)\n",
+    "\n",
+    "    # extract arxiv information ---\n",
+    "    arxiv_ids = set()\n",
+    "    for arxiv_url, pdf_url in arxiv_info.items():\n",
+    "        arxiv_ids.add(arxiv_url.split(\"/\")[-1])\n",
+    "    print(f\"# of arxiv papers: {len(arxiv_ids)}\")\n",
+    "\n",
+    "    html5ivs = get_arxiv(arxiv_ids)\n",
+    "\n",
+    "    broken = []\n",
+    "    extracted = {}\n",
+    "    extracted_from_pdf = {}\n",
+    "    broken_tex = {}\n",
+    "\n",
+    "    pbar = tqdm(range(len(html5ivs)))\n",
+    "    for arxiv_id, x in html5ivs.items():\n",
+    "        try:\n",
+    "            extracted[arxiv_id] = extract_html_content(x)\n",
+    "        except Exception as e:\n",
+    "            # print(e)\n",
+    "            print(f'getting pdf text for {arxiv_id}...')\n",
+    "            pdf_text = get_pdf_text(arxiv_id) \n",
+    "            if pdf_text is not None:\n",
+    "                extracted_from_pdf[arxiv_id] = pdf_text # get_pdf_text(arxiv_id) # <- original\n",
+    "            else:\n",
+    "                broken.append(arxiv_id)\n",
+    "        pbar.update(1)\n",
+    "    pbar.close()\n",
+    "\n",
+    "    print(f\"# of arxiv papers extracted: {len(extracted)}\")\n",
+    "    print(f\"# of arxiv papers extracted from pdf: {len(extracted_from_pdf)}\")\n",
+    "    print(f\"# of broken links: {len(broken)}\")\n",
+    "\n",
+    "    pooled = [pool_html_content(x) for _, x in extracted.items()]\n",
+    "    pooled += [\n",
+    "        clean_and_pool_markdown_content(x, arxiv_id) for arxiv_id, x in extracted_from_pdf.items()\n",
+    "    ]  # No pdf content yet, no op for now\n",
+    "    flat_pool = [item for sublist in pooled for item in sublist]\n",
+    "\n",
+    "    id_to_abstract = {}\n",
+    "    for doc in pooled:\n",
+    "        for entry in doc:\n",
+    "            if entry[\"content_type\"] == \"abstract\":\n",
+    "                id_to_abstract[entry[\"arxiv_id\"]] = entry[\"text\"]\n",
+    "                break\n",
+    "\n",
+    "    srsly.write_json(\"id_to_abstract.json\", id_to_abstract)\n",
+    "    srsly.write_json(\"flatpool.json\", flat_pool)\n",
+    "```\n",
+    "go through it and reply OK."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "24e7ef23",
+   "metadata": {},
+   "source": [
+    "OK. I've reviewed the code you provided for ingesting and processing items from a Zotero library."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "84628bb2",
+   "metadata": {},
+   "source": [
+    "OK. I've reviewed the code you provided for ingesting and processing items from a Zotero library."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "bada6931",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%ai\n",
+    "Let's work step by step. I currently have\n",
+    "\n",
+    "```\n",
+    "import os\n",
+    "import re\n",
+    "\n",
+    "import dotenv\n",
+    "import requests\n",
+    "from bs4 import BeautifulSoup\n",
+    "from retry import retry\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "dotenv.load_dotenv()\n",
+    "\n",
+    "\n",
+    "########################################################\n",
+    "### GET ZOTERO ITEMS\n",
+    "########################################################\n",
+    "@retry(tries=3, delay=8)\n",
+    "def _fetch_one_zotero_batch(url, headers, params):\n",
+    "    \"\"\"\n",
+    "    Fetch articles from Zotero API\n",
+    "    \"\"\"\n",
+    "    response = requests.get(url, headers=headers, params=params)\n",
+    "    response.raise_for_status()\n",
+    "    return response.json()\n",
+    "\n",
+    "\n",
+    "def get_zotero_items(debug=False):\n",
+    "    \"\"\"\n",
+    "    fetch items from zotero library\n",
+    "    \"\"\"\n",
+    "\n",
+    "    GROUP_ID = os.getenv(\"GROUP_ID\")\n",
+    "    API_KEY = os.getenv(\"API_KEY\")\n",
+    "    BASE_URL = f\"https://api.zotero.org/groups/{GROUP_ID}/items\"\n",
+    "    LIMIT = 100\n",
+    "\n",
+    "    headers = {\"Zotero-API-Key\": API_KEY, \"Content-Type\": \"application/json\"}\n",
+    "\n",
+    "    items = []\n",
+    "    start = 0\n",
+    "\n",
+    "    i = 1\n",
+    "    while True:\n",
+    "        i += 1\n",
+    "        params = {\"limit\": LIMIT, \"start\": start}\n",
+    "        page_items = _fetch_one_zotero_batch(BASE_URL, headers, params)\n",
+    "\n",
+    "        if not page_items:\n",
+    "            break\n",
+    "\n",
+    "        items.extend(page_items)\n",
+    "        start += LIMIT\n",
+    "        print(f\"# items fetched {len(items)}\")\n",
+    "\n",
+    "        if debug:\n",
+    "            break\n",
+    "\n",
+    "    return items\n",
+    "\n",
+    "\n",
+    "########################################################\n",
+    "### EXTRACT ARXIV LINKS AND PDFs\n",
+    "########################################################\n",
+    "def get_arxiv_items(items):\n",
+    "    arxiv_items = []\n",
+    "    arxiv_pattern = re.compile(r\"arxiv.org/abs/(\\d+\\.\\d+)\")\n",
+    "\n",
+    "    for item in items:\n",
+    "        data = item.get(\"data\", {})\n",
+    "        attachments = item.get(\"links\", {}).get(\"attachment\", {})\n",
+    "\n",
+    "        arxiv_url = None\n",
+    "        pdf_url = None\n",
+    "\n",
+    "        if \"url\" in data and \"arxiv.org\" in data[\"url\"]:\n",
+    "            arxiv_match = arxiv_pattern.search(data[\"url\"])\n",
+    "            if arxiv_match:\n",
+    "                arxiv_url = data[\"url\"]\n",
+    "\n",
+    "        if attachments:\n",
+    "            pdf_url = attachments[\"href\"]\n",
+    "\n",
+    "        if arxiv_url:\n",
+    "            arxiv_items.append(\n",
+    "                {\n",
+    "                    \"arxiv_id\": arxiv_url.split(\"/\")[-1],\n",
+    "                    \"arxiv_url\": arxiv_url,\n",
+    "                    \"pdf_url\": pdf_url,\n",
+    "                    \"added_by\": item[\"meta\"][\"createdByUser\"][\"username\"],\n",
+    "                    \"date_added\": data.get(\"dateAdded\", \"\"),\n",
+    "                }\n",
+    "            )\n",
+    "\n",
+    "    return arxiv_items\n",
+    "\n",
+    "\n",
+    "@retry(tries=3, delay=15, backoff=2)\n",
+    "def fetch_arxiv_html(arxiv_id):\n",
+    "    url = f\"https://ar5iv.labs.arxiv.org/html/{arxiv_id.split('v')[0]}\"\n",
+    "    response = requests.get(url)\n",
+    "    return response.text if response.status_code == 200 else None\n",
+    "\n",
+    "\n",
+    "def get_arxiv_html(arxiv_items):\n",
+    "    for item in tqdm(arxiv_items):\n",
+    "        html = fetch_arxiv_html(item[\"arxiv_id\"])\n",
+    "        if html:\n",
+    "            item[\"raw_html\"] = html\n",
+    "        else:\n",
+    "            print(f\"failed to fetch html for {item['arxiv_id']}\")\n",
+    "\n",
+    "    return arxiv_items\n",
+    "\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    items = get_zotero_items(debug=True)\n",
+    "    print(f\"# of items fetched from zotero: {len(items)}\")\n",
+    "    arxiv_items = get_arxiv_items(items)\n",
+    "    print(f\"# of arxiv papers: {len(arxiv_items)}\")\n",
+    "    arxiv_htmls = get_arxiv_html(arxiv_items)  # raw html files\n",
+    "--\n",
+    "review for any mistake until now. What shall we do next."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "60ab1dfc",
+   "metadata": {},
+   "source": [
+    "The code looks good so far. You've implemented functions to fetch items from Zotero, extract ArXiv-related items, and fetch the HTML content for ArXiv papers. \n",
+    "\n",
+    "Next steps could include:\n",
+    "\n",
+    "1. Parse the HTML content to extract structured information from the papers.\n",
+    "2. Implement error handling for failed HTML fetches.\n",
+    "3. Add a function to extract text from PDFs for papers where HTML fetch fails.\n",
+    "4. Store the extracted data in a suitable format (e.g., JSON, database).\n",
+    "5. Implement content pooling as in your original code.\n",
+    "\n",
+    "Which of these would you like to focus on first?"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "696318f6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%aip\n",
+    "Let's work on step 1. Let's follow the original code as much as possible.\n",
+    "Note: these are arxiv htmls. Let's also keep full body text."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "e7e4a5a6",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'arxiv_htmls' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[6], line 18\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess_arxiv_items\u001b[39m(arxiv_items):\n\u001b[1;32m     16\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m [\u001b[38;5;28mdict\u001b[39m(item, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mextract_html_content(item[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mraw_html\u001b[39m\u001b[38;5;124m\"\u001b[39m])) \u001b[38;5;28;01mfor\u001b[39;00m item \u001b[38;5;129;01min\u001b[39;00m arxiv_items \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mraw_html\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m item]\n\u001b[0;32m---> 18\u001b[0m processed_items \u001b[38;5;241m=\u001b[39m process_arxiv_items(\u001b[43marxiv_htmls\u001b[49m)\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'arxiv_htmls' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "def extract_html_content(html):\n",
+    "    soup = BeautifulSoup(html, \"html.parser\")\n",
+    "    arxiv_id = re.search(r\"\\[(\\d+\\.\\d+(v\\d+)?)\\]\", html).group(1) if re.search(r\"\\[(\\d+\\.\\d+(v\\d+)?)\\]\", html) else None\n",
+    "    paper_title = soup.find(\"h1\", class_=\"ltx_title ltx_title_document\").get_text(strip=True)\n",
+    "    [m.decompose() for m in soup.find_all([\"math\", \"cite\"])]\n",
+    "    abstract = \" \".join(p.get_text(strip=True) for p in soup.find(\"div\", class_=\"ltx_abstract\").find_all(\"p\")).replace(\")\", \") \") if soup.find(\"div\", class_=\"ltx_abstract\") else \"\"\n",
+    "    sections = soup.find_all(\"section\", class_=\"ltx_section\")\n",
+    "    body = []\n",
+    "    for section in sections:\n",
+    "        title = section.find(\"h2\", class_=\"ltx_title ltx_title_section\").get_text(strip=True)\n",
+    "        content = \" \".join(p.get_text(strip=True).replace(\")\", \") \") for p in section.find_all(\"div\", class_=\"ltx_para\"))\n",
+    "        body.append({\"title\": title, \"content\": content})\n",
+    "    return {\"arxiv_id\": arxiv_id, \"paper_title\": paper_title, \"abstract\": abstract, \"body\": body, \"full_text\": soup.get_text(strip=True)}\n",
+    "\n",
+    "def process_arxiv_items(arxiv_items):\n",
+    "    return [dict(item, **extract_html_content(item[\"raw_html\"])) for item in arxiv_items if \"raw_html\" in item]\n",
+    "\n",
+    "processed_items = process_arxiv_items(arxiv_htmls)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "70a3caaa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%aip\n",
+    "the above doesn't work and missed may details of the original implementation. Do not focus on pooling now\n",
+    "just do extract_html_content better\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4040257d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def extract_html_content(html):\n",
+    "    soup, result = BeautifulSoup(html, \"html.parser\"), []\n",
+    "    arxiv_id = re.search(r\"\\[(\\d+\\.\\d+(v\\d+)?)\\]\", html)\n",
+    "    arxiv_id = arxiv_id.group(1) if arxiv_id else None\n",
+    "    paper_title = soup.find(\"h1\", class_=\"ltx_title ltx_title_document\").get_text(strip=True)\n",
+    "    [m.decompose() for m in soup.find_all([\"math\", \"cite\"])]\n",
+    "    abstract = soup.find(\"div\", class_=\"ltx_abstract\")\n",
+    "    if abstract:\n",
+    "        result.append({\"content\": \" \".join(p.get_text(strip=True) for p in abstract.find_all(\"p\")).replace(\")\", \") \"), \"title\": \"Abstract\", \"paper_title\": paper_title, \"content_type\": \"abstract\"})\n",
+    "    intro = soup.find(\"section\", id=\"S1\")\n",
+    "    if intro:\n",
+    "        for para in intro.find_all(\"div\", class_=\"ltx_para\"):\n",
+    "            result.append({\"content\": para.get_text(strip=True).replace(\")\", \") \"), \"title\": \"Introduction\", \"paper_title\": paper_title, \"content_type\": \"introduction\"})\n",
+    "    for section in soup.find_all(\"section\", class_=\"ltx_section\")[1:-1]:\n",
+    "        section_title = section.find(\"h2\", class_=\"ltx_title ltx_title_section\").get_text(strip=True)\n",
+    "        for para in section.find_all(\"div\", class_=\"ltx_para\"):\n",
+    "            result.append({\"content\": para.get_text(strip=True).replace(\")\", \") \"), \"title\": section_title, \"paper_title\": paper_title, \"content_type\": \"body\"})\n",
+    "    conclusion = soup.find(\"section\", id=\"S4\")\n",
+    "    if conclusion:\n",
+    "        result.append({\"content\": \" \".join(p.get_text(strip=True) for p in conclusion.find_all(\"div\", class_=\"ltx_para\")).replace(\")\", \") \"), \"title\": conclusion.find(\"h2\", class_=\"ltx_title ltx_title_section\").get_text(strip=True), \"paper_title\": paper_title, \"content_type\": \"conclusion\"})\n",
+    "    for c in result:\n",
+    "        c[\"arxiv_id\"] = arxiv_id\n",
+    "    return result\n",
+    "\n",
+    "processed_items = [extract_html_content(item[\"raw_html\"]) for item in arxiv_items if \"raw_html\" in item]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0242f021",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}