learningcircuit--local-deep-research/tests/download_stuff_for_local.py

import os

import arxiv
from datasets import load_dataset

max_results = 50
# Search for papers and limit to 400
search = arxiv.Search(query="machine learning", max_results=50)
client = arxiv.Client()

# Download papers
os.makedirs("./../local_search_files/research_papers", exist_ok=True)
for i, result in enumerate(client.results(search)):
    try:
        filename = f"./../local_search_files/research_papers/paper_{i}.pdf"
        result.download_pdf(filename=filename)
        print(f"Downloaded {i + 1} / {max_results}: {result.title}")
    except Exception as e:
        print(f"Error downloading {result.title}: {e}")


wiki = load_dataset(
    "wikipedia",
    "20220301.en",
    split=f"train[:{max_results}]",
    trust_remote_code=True,
)

os.makedirs("./../local_search_files/wiki_sample", exist_ok=True)
for i, article in enumerate(wiki):
    with open(f"./../local_search_files/wiki_sample/article_{i}.txt", "w") as f:
        f.write(article["title"] + "\n\n" + article["text"])