Text-to-Document-Generation/PDF-Redaction-API

1"""2Example client for PDF Redaction API3"""4import requests5from pathlib import Path6import sys7 8 9def redact_pdf(api_url: str, pdf_path: str, output_path: str = "redacted.pdf", 10               dpi: int = 300, entity_types: str = None):11    """12    Redact a PDF file using the API13    14    Args:15        api_url: Base URL of the API16        pdf_path: Path to the PDF file to redact17        output_path: Path to save the redacted PDF18        dpi: DPI for OCR processing19        entity_types: Comma-separated list of entity types to redact20    """21    # Check if file exists22    if not Path(pdf_path).exists():23        print(f"Error: File {pdf_path} not found")24        return False25    26    print(f"Uploading {pdf_path}...")27    28    # Prepare request29    files = {"file": open(pdf_path, "rb")}30    params = {"dpi": dpi}31    32    if entity_types:33        params["entity_types"] = entity_types34    35    try:36        # Upload and redact37        response = requests.post(f"{api_url}/redact", files=files, params=params)38        response.raise_for_status()39        40        result = response.json()41        print(f"\nStatus: {result['status']}")42        print(f"Message: {result['message']}")43        44        # Display found entities45        if result.get('entities'):46            print("\nEntities redacted:")47            for i, entity in enumerate(result['entities'], 1):48                print(f"  {i}. {entity['entity_type']}: {entity['entity_text']} "49                      f"(Page {entity['page']}, {entity['word_count']} words)")50        51        # Download redacted file52        job_id = result['job_id']53        print(f"\nDownloading redacted PDF...")54        55        download_response = requests.get(f"{api_url}/download/{job_id}")56        download_response.raise_for_status()57        58        # Save file59        with open(output_path, "wb") as f:60            f.write(download_response.content)61        62        print(f"✓ Redacted PDF saved to: {output_path}")63        64        # Cleanup (optional)65        # requests.delete(f"{api_url}/cleanup/{job_id}")66        67        return True68    69    except requests.exceptions.RequestException as e:70        print(f"Error: {e}")71        return False72    finally:73        files["file"].close()74 75 76def check_health(api_url: str):77    """Check API health"""78    try:79        response = requests.get(f"{api_url}/health")80        response.raise_for_status()81        data = response.json()82        83        print(f"API Status: {data['status']}")84        print(f"Version: {data['version']}")85        print(f"Model Loaded: {data['model_loaded']}")86        87        return True88    except requests.exceptions.RequestException as e:89        print(f"Error checking health: {e}")90        return False91 92 93def get_stats(api_url: str):94    """Get API statistics"""95    try:96        response = requests.get(f"{api_url}/stats")97        response.raise_for_status()98        data = response.json()99        100        print("API Statistics:")101        print(f"  Pending uploads: {data['pending_uploads']}")102        print(f"  Processed files: {data['processed_files']}")103        print(f"  Model loaded: {data['model_loaded']}")104        105        return True106    except requests.exceptions.RequestException as e:107        print(f"Error getting stats: {e}")108        return False109 110 111if __name__ == "__main__":112    # Example usage113    114    # For local development115    API_URL = "http://localhost:7860"116    117    # For HuggingFace Spaces (replace with your space URL)118    # API_URL = "https://your-username-pdf-redaction-api.hf.space"119    120    if len(sys.argv) < 2:121        print("Usage:")122        print("  python client_example.py <pdf_file> [output_file] [dpi]")123        print("\nOr check health:")124        print("  python client_example.py --health")125        print("\nOr get stats:")126        print("  python client_example.py --stats")127        sys.exit(1)128    129    if sys.argv[1] == "--health":130        check_health(API_URL)131    elif sys.argv[1] == "--stats":132        get_stats(API_URL)133    else:134        pdf_path = sys.argv[1]135        output_path = sys.argv[2] if len(sys.argv) > 2 else "redacted.pdf"136        dpi = int(sys.argv[3]) if len(sys.argv) > 3 else 300137        138        # Optional: Filter specific entity types139        # entity_types = "PER,ORG"  # Only redact persons and organizations140        entity_types = None  # Redact all entity types141        142        redact_pdf(API_URL, pdf_path, output_path, dpi, entity_types)143