Script to download IRS data in bulk and prep for notebooklm

README.md

IRS Bulk Download

pip install tqdm requests

irs-build-download-markitdown.py

import os
import shutil
import zipfile
import requests
from tqdm import tqdm
from markitdown import MarkItDown

# Base URL for GovInfo's JSON-formatted Bulk Data API
BULK_DATA_CFR_URL = "https://www.govinfo.gov/bulkdata/json/CFR"

def download_and_convert_cfr(output_dir="./latest_cfr", markdown_dir="./cfr_markdown", notebook_dir="./notebooklm_sources"):
    # Initialize the MarkItDown converter instance
    md_converter = MarkItDown()

    # Create output directories
    for path in [output_dir, markdown_dir, notebook_dir]:
        if not os.path.exists(path):
            os.makedirs(path)
        
    print("Fetching CFR directory list from GovInfo...")
    headers = {"Accept": "application/json"}
    
    try:
        response = requests.get(BULK_DATA_CFR_URL, headers=headers)
        response.raise_for_status()
        data = response.json()
    except requests.exceptions.RequestException as e:
        print(f"Error connecting to GovInfo API: {e}")
        return

    # Extract available years and find the latest edition
    folders = data.get("files", [])
    years = [int(f["name"]) for f in folders if f.get("name") and f["name"].isdigit()]
            
    if not years:
        print("Could not parse available CFR years from API response.")
        return
        
    latest_year = max(years)
    print(f"Latest available CFR publication year identified: {latest_year}\n")

    # Fetch the file list for the latest year directory
    year_url = f"{BULK_DATA_CFR_URL}/{latest_year}"
    try:
        response = requests.get(year_url, headers=headers)
        response.raise_for_status()
        year_data = response.json()
    except requests.exceptions.RequestException as e:
        print(f"Error fetching data for year {latest_year}: {e}")
        return

    # Filter out only the ZIP archives
    title_files = [
        item for item in year_data.get("files", []) 
        if item.get("name") and item["name"].endswith(".zip")
    ]
    
    total_files = len(title_files)
    if total_files == 0:
        print("No CFR ZIP files found for download.")
        return

    print(f"Found {total_files} files to process. Starting download and conversion stage...\n")

    # Master progress bar tracking overall file progression
    with tqdm(total=total_files, desc="Conversion Progress", unit="file", position=0) as overall_bar:
        for item in title_files:
            filename = item.get("name")
            download_url = item.get("link")
            local_zip_path = os.path.join(output_dir, filename)
            
            archive_name = os.path.splitext(filename)[0]
            target_md_folder = os.path.join(markdown_dir, archive_name)
            
            try:
                # 1. Download/Verify local copy
                if os.path.exists(local_zip_path):
                    tqdm.write(f"[INFO] Local copy of {filename} found. Skipping download.")
                else:
                    with requests.get(download_url, stream=True) as file_stream:
                        file_stream.raise_for_status()
                        total_size = int(file_stream.headers.get('content-length', 0))
                        
                        with tqdm(
                            total=total_size, 
                            desc=f" ↳ Downloading {filename}", 
                            unit='B', 
                            unit_scale=True, 
                            unit_divisor=1024,
                            position=1, 
                            leave=False
                        ) as file_bar:
                            
                            with open(local_zip_path, "wb") as f:
                                for chunk in file_stream.iter_content(chunk_size=8192):
                                    if chunk:
                                        f.write(chunk)
                                        file_bar.update(len(chunk))
                
                # 2. Reset Markdown subdirectories
                if os.path.exists(target_md_folder):
                    shutil.rmtree(target_md_folder)
                os.makedirs(target_md_folder)

                # 3. Process zip archive
                with zipfile.ZipFile(local_zip_path, 'r') as zip_ref:
                    xml_files = [f for f in zip_ref.namelist() if f.lower().endswith('.xml')]
                    
                    for xml_file in xml_files:
                        # Extract the XML file to a temporary file layout, as markitdown 
                        # natively relies on file paths or extensions to resolve conversion logic
                        temp_xml_path = os.path.join(target_md_folder, os.path.basename(xml_file))
                        with zip_ref.open(xml_file) as xml_content, open(temp_xml_path, "wb") as temp_out:
                            temp_out.write(xml_content.read())
                        
                        # Convert XML file via MarkItDown
                        # markitdown returns a result object containing text content
                        result = md_converter.convert(temp_xml_path)
                        markdown_text = result.text_content
                        
                        # Save out the generated markdown
                        md_filename = os.path.splitext(os.path.basename(xml_file))[0] + ".md"
                        md_file_path = os.path.join(target_md_folder, md_filename)
                        
                        with open(md_file_path, "w", encoding="utf-8") as md_file:
                            md_file.write(markdown_text)
                        
                        # Clean up the raw temp XML file to keep things tidy
                        os.remove(temp_xml_path)
                                    
                overall_bar.update(1)
                
            except Exception as e:
                tqdm.write(f"\n[ERROR] Processing error on {filename}: {e}")
                overall_bar.update(1)

    # 4. Post-Processing: Create Optimized NotebookLM Chunks
    print("\n--------------------------------------------------")
    print("Starting NotebookLM Source Assembly and Chunking...")
    print("--------------------------------------------------")
    
    if os.path.exists(notebook_dir):
        shutil.rmtree(notebook_dir)
    os.makedirs(notebook_dir)

    current_chunk_words = []
    current_chunk_sources = set()
    chunk_index = 1
    max_words = 50000

    for root, _, files in os.walk(markdown_dir):
        for file in sorted(files):
            if file.endswith(".md"):
                file_path = os.path.join(root, file)
                parent_folder = os.path.basename(root)
                source_identifier = f"{parent_folder}/{file}"
                
                with open(file_path, "r", encoding="utf-8") as f:
                    content = f.read()
                
                file_words = content.split()
                
                if len(current_chunk_words) + len(file_words) > max_words and current_chunk_words:
                    save_notebooklm_chunk(notebook_dir, chunk_index, current_chunk_words, current_chunk_sources, latest_year)
                    chunk_index += 1
                    current_chunk_words = []
                    current_chunk_sources = set()
                
                current_chunk_words.extend(file_words)
                current_chunk_sources.add(source_identifier)
                
    if current_chunk_words:
        save_notebooklm_chunk(notebook_dir, chunk_index, current_chunk_words, current_chunk_sources, latest_year)

    print(f"\nSuccessfully generated {chunk_index} notebook-ready text files.")
    print(f"Target Destination: '{notebook_dir}'")

def save_notebooklm_chunk(directory, index, word_list, source_set, year):
    chunk_filename = f"cfr_edition_{year}_chunk_{index:03d}.md"
    chunk_path = os.path.join(directory, chunk_filename)
    
    header = (
        "================================================================================\n"
        f"NOTEBOOKLM SOURCE BLOCK: CHUNK {index:03d}\n"
        f"REGULATORY EDITION YEAR: {year}\n"
        "COMPILING METADATA FROM SOURCE DOCUMENTS INCLUDED IN THIS FILE:\n"
    )
    for src in sorted(source_set):
        header += f" - Source Path Reference: {src}\n"
    header += "================================================================================\n\n"
    
    body_text = " ".join(word_list)
    
    with open(chunk_path, "w", encoding="utf-8") as out_f:
        out_f.write(header + body_text)

if __name__ == "__main__":
    download_and_convert_cfr()