datamule 0.381__py3-none-any.whl → 1.0.2__py3-none-any.whl

datamule/dataset_builder/dataset_builder.py

@@ -1,260 +0,0 @@
-import pandas as pd
-import json
-import os
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from tqdm import tqdm
-import google.generativeai as genai
-import time
-from datetime import datetime
-import psutil
-from threading import Lock
-
-class RateLimiter:
-    def __init__(self, max_rpm):
-        self.min_delay = 62.0 / max_rpm # 58 seconds to allow for some buffer WIP
-        self.last_request = time.time()
-        self.lock = Lock()
-        self.request_count = 0
-    
-    def acquire(self):
-        with self.lock:
-            now = time.time()
-            time_since_last = now - self.last_request
-            delay_needed = self.min_delay - time_since_last
-            self.last_request = now + max(0, delay_needed)  # Update based on expected completion
-            self.request_count += 1
-            count = self.request_count
-        
-        # Sleep outside the lock
-        if delay_needed > 0:
-            time.sleep(delay_needed)
-        
-        return count
-
-class DatasetBuilder:
-    def __init__(self):
-        self.base_prompt = None
-        self.response_schema = None
-        self.input_path = None
-        self.output_path = None
-        self.failed_path = None
-        self.max_rpm = 1450
-        self.max_workers = 30
-        self.save_frequency = 100
-        self.output_columns = None
-        self.buffer = []
-        self.buffer_lock = Lock()
-        self.failed_ids = set()
-        self.failed_lock = Lock()
-        self.model_name = "gemini-1.5-flash-8b"  # Default model
-        self.model_config = {}  # Additional model configuration
-        self.api_key = None
-
-    def set_api_key(self, api_key):
-        """Set the API key for Google's Generative AI."""
-        self.api_key = api_key
-        genai.configure(api_key=api_key)
-        return self
-
-    def set_paths(self, input_path, output_path, failed_path):
-        """Set input and output file paths."""
-        self.input_path = input_path
-        self.output_path = output_path
-        self.failed_path = failed_path
-        return self
-
-    def set_base_prompt(self, prompt):
-        """Set the base prompt for LLM processing."""
-        self.base_prompt = prompt
-        return self
-
-    def set_response_schema(self, schema):
-        """Set the response schema and derive output columns."""
-        self.response_schema = schema
-        # Derive output columns from schema
-        if schema and 'items' in schema and 'properties' in schema['items']:
-            properties = schema['items']['properties']
-            self.output_columns = ['accession_number'] + list(properties.keys())
-        return self
-
-    def set_rpm(self, max_rpm=1450):
-        """Set the maximum requests per minute."""
-        self.max_rpm = max_rpm
-        return self
-
-    def set_max_workers(self, max_workers=30):
-        """Set the maximum number of concurrent workers."""
-        self.max_workers = max_workers
-        return self
-
-    def set_save_frequency(self, frequency=100):
-        """Set how often to save progress."""
-        self.save_frequency = frequency
-        return self
-
-    def set_model(self, model_name="gemini-1.5-flash-8b", **model_config):
-        """Set the model name and configuration."""
-        self.model_name = model_name
-        self.model_config = model_config
-        return self
-
-    def validate_config(self):
-        """Validate that all required configurations are set."""
-        if not all([self.base_prompt, self.response_schema, self.input_path, 
-                   self.output_path, self.failed_path, self.api_key]):
-            raise ValueError("""Missing required configuration. Please ensure you have set:
-                           - API key
-                           - Paths (input_path, output_path, failed_path)
-                           - Base prompt
-                           - Response schema""")
-
-    def get_processed_ids(self):
-        """Get set of processed accession numbers from output file."""
-        if not os.path.exists(self.output_path):
-            return set()
-        
-        try:
-            # Read only the accession_number column for memory efficiency
-            df = pd.read_csv(self.output_path, usecols=['accession_number'])
-            return set(df['accession_number'])
-        except Exception as e:
-            print(f"Warning: Error reading processed IDs: {e}")
-            return set()
-
-    def save_data(self, df_new):
-        """Append new data to existing CSV."""
-        df_new.to_csv(self.output_path, mode='a', header=not os.path.exists(self.output_path), index=False)
-
-    def save_failed_ids(self):
-        """Save failed accession numbers to file."""
-        with open(self.failed_path, 'w') as f:
-            for acc in self.failed_ids:
-                f.write(f"{acc}\n")
-
-    def process_text(self, args):
-        """Process a single text entry through the model."""
-        model, text, accession_number, rate_limiter = args
-        
-        current_requests = rate_limiter.acquire()
-        
-        full_prompt = self.base_prompt + "\n\nINFORMATION:\n" + text
-
-        try:
-            generation_config = genai.GenerationConfig(
-                response_mime_type="application/json",
-                response_schema=self.response_schema,
-                **self.model_config
-            )
-            
-            response = model.generate_content(
-                full_prompt,
-                generation_config=generation_config
-            )
-            results = json.loads(response.text)
-            
-            for result in results:
-                result['accession_number'] = accession_number
-            
-            with self.buffer_lock:
-                self.buffer.extend(results)
-            
-            return True, current_requests
-        except Exception as e:
-            with self.failed_lock:
-                self.failed_ids.add(accession_number)
-            return False, f"Error processing {accession_number}: {str(e)}"
-
-    def build(self):
-        """Main processing method to build the dataset."""
-        self.validate_config()
-
-        # Initialize model and rate limiter
-        model = genai.GenerativeModel(self.model_name)
-        rate_limiter = RateLimiter(self.max_rpm)
-        
-        # Load data
-        print("Loading data...")
-        df_input = pd.read_csv(self.input_path)
-        processed_ids = self.get_processed_ids()
-        df_to_process = df_input[~df_input['accession_number'].isin(processed_ids)]
-        
-        total_in_dataset = len(df_input)
-        already_processed = len(processed_ids)
-        to_process = len(df_to_process)
-        
-        print(f"Total entries in dataset: {total_in_dataset}")
-        print(f"Already processed: {already_processed}")
-        print(f"New entries to process: {to_process}")
-        
-        if len(df_to_process) == 0:
-            print("All entries already processed!")
-            return
-
-        work_items = [
-            (model, row['text'], row['accession_number'], rate_limiter) 
-            for _, row in df_to_process.iterrows()
-        ]
-        
-        start_time = time.time()
-        last_save_time = time.time()
-        processed_count = 0
-        
-        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            futures = {executor.submit(self.process_text, item): item for item in work_items}
-            
-            with tqdm(total=total_in_dataset, initial=already_processed, desc="Processing entries") as pbar:
-                for future in as_completed(futures):
-                    success, result = future.result()
-                    
-                    if not success:
-                        print(f"\n{result}")
-                    
-                    processed_count += 1
-                    pbar.update(1)
-                    
-                    elapsed = time.time() - start_time
-                    rpm = processed_count / (elapsed / 60)
-                    memory_usage = psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024
-                    
-                    current_progress = already_processed + processed_count
-                    pbar.set_description(
-                        f"Processed {current_progress}/{total_in_dataset} | {rpm:.0f} RPM | Mem: {memory_usage:.0f}MB"
-                    )
-                    
-                    # Save periodically using append
-                    if len(self.buffer) >= self.save_frequency:
-                        with self.buffer_lock:
-                            df_new = pd.DataFrame(self.buffer)
-                            self.buffer = []
-                        
-                        if not df_new.empty:
-                            self.save_data(df_new)
-                            last_save_time = time.time()
-                    
-                    # Save failed IDs periodically
-                    if self.failed_ids and time.time() - last_save_time > 300:
-                        self.save_failed_ids()
-                        last_save_time = time.time()
-        
-        # Save any remaining results
-        if self.buffer:
-            with self.buffer_lock:
-                df_new = pd.DataFrame(self.buffer)
-                self.buffer = []
-            
-            if not df_new.empty:
-                self.save_data(df_new)
-        
-        if self.failed_ids:
-            self.save_failed_ids()
-        
-        # Print final statistics
-        elapsed = time.time() - start_time
-        final_rpm = processed_count / (elapsed / 60)
-        
-        print(f"\nProcessing complete:")
-        print(f"Total processed in this run: {processed_count}")
-        print(f"Average speed: {final_rpm:.0f} RPM")
-        print(f"Failed entries: {len(self.failed_ids)}")
-        if self.failed_ids:
-            print(f"Failed entries saved to: {self.failed_path}")

datamule/downloader/dropbox_downloader.py

@@ -1,225 +0,0 @@
-import asyncio
-import aiohttp
-from aiolimiter import AsyncLimiter
-import os
-from tqdm import tqdm
-from tqdm.asyncio import tqdm as atqdm
-from datetime import datetime, timedelta
-from urllib.parse import urlparse, parse_qs, urlencode
-import math
-import re
-import aiofiles
-import json
-import csv
-from pkg_resources import resource_filename
-import zipfile
-import shutil
-
-class DropboxDownloader:
-    """
-    Asynchronous downloader for handling multiple file downloads with rate limiting and progress tracking.
-    
-    This class provides functionality to download multiple files concurrently with rate limiting,
-    progress bars, and automatic handling of zip archives (including multi-part archives).
-
-    Parameters
-    ----------
-    concurrent_downloads : int, optional
-        Maximum number of concurrent downloads allowed (default is 5)
-    rate_limit : int, optional
-        Maximum number of requests per second (default is 10)
-
-    Attributes
-    ----------
-    semaphore : asyncio.Semaphore
-        Controls the number of concurrent downloads
-    rate_limiter : AsyncLimiter
-        Handles rate limiting of requests
-    session : aiohttp.ClientSession
-        HTTP session for making requests
-    progress_bars : dict
-        Dictionary of progress bars for active downloads
-    """
-
-    def __init__(self, concurrent_downloads=5, rate_limit=10):
-        """Initialize the DropboxDownloader with specified concurrency and rate limits."""
-        self.semaphore = asyncio.Semaphore(concurrent_downloads)
-        self.rate_limiter = AsyncLimiter(rate_limit, 1)  # rate_limit requests per second
-        self.session = None
-        self.progress_bars = {}
-
-    async def create_session(self):
-        """
-        Create an aiohttp client session.
-
-        This method should be called before starting any downloads.
-        """
-        self.session = aiohttp.ClientSession()
-
-    async def close_session(self):
-        """
-        Close the aiohttp client session.
-
-        This method should be called after all downloads are complete.
-        """
-        if self.session:
-            await self.session.close()
-
-    async def download_file(self, url, dest_folder):
-        """
-        Download a single file with progress tracking and automatic unzipping.
-
-        Parameters
-        ----------
-        url : str
-            URL of the file to download
-        dest_folder : str
-            Destination folder where the file will be saved
-
-        Notes
-        -----
-        - Uses a progress bar to show download progress
-        - Automatically handles zip files and multi-part archives
-        - Rate limited based on the instance settings
-        """
-        async with self.semaphore:
-            await self.rate_limiter.acquire()
-            try:
-                file_name = os.path.basename(urlparse(url).path)
-                file_path = os.path.join(dest_folder, file_name)
-
-                async with self.session.get(url) as response:
-                    if response.status != 200:
-                        print(f"Failed to download {url}: HTTP {response.status}")
-                        return
-
-                    file_size = int(response.headers.get('Content-Length', 0))
-                    
-                    if file_name not in self.progress_bars:
-                        self.progress_bars[file_name] = tqdm(
-                            total=file_size,
-                            unit='iB',
-                            unit_scale=True,
-                            desc=file_name
-                        )
-
-                    async with aiofiles.open(file_path, 'wb') as f:
-                        chunk_size = 8192
-                        downloaded = 0
-                        async for chunk in response.content.iter_chunked(chunk_size):
-                            await f.write(chunk)
-                            downloaded += len(chunk)
-                            self.progress_bars[file_name].update(len(chunk))
-
-                    self.progress_bars[file_name].close()
-                    del self.progress_bars[file_name]
-                    
-                print(f"Downloaded {file_name}")
-                
-                if re.match(r'.*\.zip(\.001)?$', file_name):
-                    await self.unzip_file(file_path, dest_folder)
-
-            except Exception as e:
-                print(f"Error downloading {url}: {str(e)}")
-
-    async def unzip_file(self, file_path, dest_folder):
-        """
-        Extract contents of a zip file and clean up archive files.
-
-        Handles both single zip files and multi-part archives.
-
-        Parameters
-        ----------
-        file_path : str
-            Path to the zip file
-        dest_folder : str
-            Destination folder for extracted contents
-
-        Notes
-        -----
-        - Automatically combines multi-part archives
-        - Deletes archive files after successful extraction
-        - Handles both .zip and .zip.001 format files
-        """
-        try:
-            base_name = os.path.splitext(file_path)[0]
-            if file_path.endswith('.001'):
-                base_name = os.path.splitext(base_name)[0]
-            
-            combined_zip = f"{base_name}.zip"
-            
-            # Combine parts if necessary
-            if not os.path.exists(combined_zip):
-                with open(combined_zip, 'wb') as outfile:
-                    part_num = 1
-                    while True:
-                        part_path = f"{base_name}.zip.{part_num:03d}"
-                        if not os.path.exists(part_path):
-                            break
-                        with open(part_path, 'rb') as infile:
-                            shutil.copyfileobj(infile, outfile)
-                        part_num += 1
-
-            # Unzip the combined file
-            with zipfile.ZipFile(combined_zip, 'r') as zip_ref:
-                zip_ref.extractall(dest_folder)
-            print(f"Unzipped {os.path.basename(base_name)}")
-
-            # Remove archive files
-            if os.path.exists(combined_zip):
-                os.remove(combined_zip)
-            
-            part_num = 1
-            while True:
-                part_path = f"{base_name}.zip.{part_num:03d}"
-                if not os.path.exists(part_path):
-                    break
-                os.remove(part_path)
-                part_num += 1
-            
-            standalone_zip = f"{base_name}.zip"
-            if os.path.exists(standalone_zip):
-                os.remove(standalone_zip)
-                
-            print(f"Cleaned up zip files for {os.path.basename(base_name)}")
-        except Exception as e:
-            print(f"Error unzipping {file_path}: {str(e)}")
-
-    async def _download_urls(self, urls, dest_folder):
-        """
-        Internal method to handle multiple URL downloads.
-
-        Parameters
-        ----------
-        urls : list of str
-            List of URLs to download
-        dest_folder : str
-            Destination folder for downloaded files
-        """
-        os.makedirs(dest_folder, exist_ok=True)
-        await self.create_session()
-        tasks = [self.download_file(url, dest_folder) for url in urls]
-        await asyncio.gather(*tasks)
-        await self.close_session()
-
-    def download(self, urls, output_dir):
-        """
-        Download multiple URLs to a specified directory.
-
-        This is the main method to use for downloading files. It handles the creation
-        and cleanup of the async event loop.
-
-        Parameters
-        ----------
-        urls : list of str
-            List of URLs to download
-        output_dir : str
-            Directory where files will be saved
-
-        Examples
-        --------
-        >>> downloader = DropboxDownloader(concurrent_downloads=3, rate_limit=5)
-        >>> urls = ['http://example.com/file1.zip', 'http://example.com/file2.zip']
-        >>> downloader.download(urls, '/path/to/output')
-        """
-        return asyncio.run(self._download_urls(urls, output_dir))

datamule/downloader/ftd.py

@@ -1,216 +0,0 @@
-from datetime import datetime, timedelta
-import pkg_resources
-import io
-import re
-import os
-import zipfile
-import csv
-from concurrent.futures import ThreadPoolExecutor
-from tqdm import tqdm
-
-def process_ftd_zip(zip_path):
-    """
-    Process a single FTD (Fails-to-Deliver) ZIP file by converting its contents to CSV format.
-
-    This function extracts the contents of a ZIP file containing tab-delimited data,
-    converts it to CSV format, and removes the original ZIP file.
-
-    Parameters
-    ----------
-    zip_path : str
-        Path to the ZIP file to process
-
-    Notes
-    -----
-    - Assumes each ZIP file contains exactly one data file
-    - Uses '|' as the delimiter for input data
-    - Removes the original ZIP file after processing
-    - Handles UTF-8 encoding with replacement for invalid characters
-
-    Examples
-    --------
-    >>> process_ftd_zip('path/to/cnsfails202301a.zip')
-    """
-    base_name = os.path.splitext(zip_path)[0]
-    csv_path = f"{base_name}.csv"
-    
-    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-        file_name = zip_ref.namelist()[0]  # Assuming only one file per zip
-        with zip_ref.open(file_name) as file:
-            content = io.TextIOWrapper(file, encoding='utf-8', errors='replace').read()
-    
-    # Convert tab-delimited content to CSV
-    with open(csv_path, 'w', newline='', encoding='utf-8') as csvfile:
-        writer = csv.writer(csvfile)
-        for line in content.split('\n'):
-            writer.writerow(line.split('|'))
-    
-    # Remove the original zip file
-    os.remove(zip_path)
-
-def process_all_ftd_zips(output_dir):
-    """
-    Process all FTD ZIP files in a directory concurrently.
-
-    Uses ThreadPoolExecutor for parallel processing with progress tracking.
-
-    Parameters
-    ----------
-    output_dir : str
-        Directory containing the ZIP files to process
-
-    Notes
-    -----
-    - Processes all files ending with '.zip' in the directory
-    - Shows progress bar during processing
-    - Uses system default number of worker threads
-
-    Examples
-    --------
-    >>> process_all_ftd_zips('/path/to/ftd/files')
-    """
-    zip_files = [f for f in os.listdir(output_dir) if f.endswith('.zip')]
-    
-    # Use ThreadPoolExecutor for parallel processing with tqdm
-    with ThreadPoolExecutor() as executor:
-        list(tqdm(executor.map(lambda f: process_ftd_zip(os.path.join(output_dir, f)), zip_files),
-                  total=len(zip_files),
-                  desc="Processing ZIP files",
-                  unit="file"))
-
-def load_csv_data():
-    """
-    Load existing FTD locations data from the package's data directory.
-
-    Returns
-    -------
-    list of dict
-        List of dictionaries containing the CSV data, where each dictionary
-        represents a row with keys from the CSV header
-
-    Notes
-    -----
-    - Reads from 'ftd_locations.csv' in the package's data directory
-    - Uses UTF-8 encoding
-    - Assumes CSV has a header row
-
-    Examples
-    --------
-    >>> data = load_csv_data()
-    >>> print(data[0]['url'])  # Print first URL in data
-    """
-    csv_content = pkg_resources.resource_string('datamule', 'data/ftd_locations.csv')
-    csv_data = []
-    csv_file = io.StringIO(csv_content.decode('utf-8'))
-    csv_reader = csv.DictReader(csv_file)
-    for row in csv_reader:
-        csv_data.append(row)
-    return csv_data
-
-def extract_date_from_url(url):
-    """
-    Extract the date from an FTD file URL.
-
-    Parameters
-    ----------
-    url : str
-        URL of the FTD file
-
-    Returns
-    -------
-    datetime or None
-        Datetime object representing the month and year from the URL,
-        or None if no date is found
-
-    Notes
-    -----
-    - Expects URLs in format containing 'cnsfails{YYYYMM}[a|b].zip'
-    - Returns None if the URL doesn't match the expected pattern
-
-    Examples
-    --------
-    >>> url = "https://www.sec.gov/files/data/fails-deliver-data/cnsfails202301a.zip"
-    >>> date = extract_date_from_url(url)
-    >>> print(date)
-    2023-01-01 00:00:00
-    """
-    match = re.search(r'cnsfails(\d{6})[ab]\.zip', url)
-    if match:
-        date_str = match.group(1)
-        return datetime.strptime(date_str, '%Y%m')
-    return None
-
-def generate_urls(start_date, end_date):
-    """
-    Generate FTD file URLs for a date range.
-
-    Parameters
-    ----------
-    start_date : datetime
-        Start date for URL generation
-    end_date : datetime
-        End date for URL generation
-
-    Returns
-    -------
-    list of str
-        List of URLs for FTD files in the date range
-
-    Notes
-    -----
-    - Generates two URLs per month ('a' and 'b' files)
-    - Uses 15-day intervals for half-month periods
-    - URLs follow SEC's file naming convention
-
-    Examples
-    --------
-    >>> start = datetime(2023, 1, 1)
-    >>> end = datetime(2023, 12, 31)
-    >>> urls = generate_urls(start, end)
-    """
-    urls = []
-    current_date = start_date
-    while current_date <= end_date:
-        for half in ['a', 'b']:
-            url = f"https://www.sec.gov/files/data/fails-deliver-data/cnsfails{current_date.strftime('%Y%m')}{half}.zip"
-            urls.append(url)
-        current_date += timedelta(days=15)  # Move to the next half-month
-    return urls
-
-def get_all_ftd_urls():
-    """
-    Get a complete list of FTD URLs, including both existing and new ones.
-
-    Returns
-    -------
-    list of str
-        Combined list of existing and newly generated URLs
-
-    Notes
-    -----
-    - Loads existing URLs from the package's data file
-    - Finds the latest date in existing URLs
-    - Generates new URLs from the month after the latest date up to current date
-    - Combines and returns all URLs
-
-    Examples
-    --------
-    >>> urls = get_all_ftd_urls()
-    >>> print(len(urls))  # Number of URLs
-    >>> print(urls[-1])   # Most recent URL
-    """
-    # Load existing URLs
-    csv_data = load_csv_data()
-    existing_urls = [row['url'] for row in csv_data]
-
-    # Find the last date in the existing URLs
-    last_date = max(extract_date_from_url(url) for url in existing_urls if extract_date_from_url(url))
-
-    # Generate new URLs starting from the month after the last date
-    start_date = (last_date.replace(day=1) + timedelta(days=32)).replace(day=1)
-    end_date = datetime.now()
-    
-    new_urls = generate_urls(start_date, end_date)
-
-    # Combine and return all URLs
-    return existing_urls + new_urls