Spaces:
Sleeping
Sleeping
| import argparse | |
| import sys | |
| import os | |
| from src.detectors.hybrid_detector import HybridDetector | |
| from src.redactors.mapper import StatefulMapper | |
| from src.redactors.docx_handler import DocxHandler | |
| from src.evaluation.metrics import Evaluator | |
| def main(): | |
| parser = argparse.ArgumentParser(description="PII Redaction Tool for Financial Documents") | |
| parser.add_argument( | |
| "--mode", | |
| type=str, | |
| required=True, | |
| choices=["redact", "evaluate"], | |
| help="Mode of operation: 'redact' a .docx file or 'evaluate' detection metrics." | |
| ) | |
| parser.add_argument( | |
| "--input", | |
| type=str, | |
| required=True, | |
| help="Path to the input file (.docx for redact, .json for evaluate)." | |
| ) | |
| parser.add_argument( | |
| "--output", | |
| type=str, | |
| help="Path to the output .docx file (required for redact mode)." | |
| ) | |
| parser.add_argument( | |
| "--whitelist", | |
| type=str, | |
| nargs='+', | |
| default=[ | |
| "Cap Price", "Floor Price", "Offer Price", "Book Building Process", | |
| "Cut-off Price", "Red Herring Prospectus", "Equity Shares", | |
| "Offer for Sale", "Fresh Issue", "SEBI", "BSE", "NSE", "Companies Act", | |
| "Selling Shareholder", "Selling Shareholders", "Promoter Selling Shareholder", | |
| "Company Secretary", "Compliance Officer", "Book Running Lead Manager", | |
| "Statutory Auditors", "Chartered Accountants", | |
| "REGISTERED OFFICE", "CORPORATE OFFICE", "CONTACT PERSON", | |
| "E-MAIL AND TELEPHONE", "WEBSITE", "DETAILS OF THE OFFER TO PUBLIC", | |
| "TYPE", "SIZE OF THE FRESH ISSUE", "TOTAL OFFER SIZE", | |
| "ELIGIBILITY AND SHARE RESERVATION", "Qualified Institutional Buyers", | |
| "Retail Individual Investors", "Stock Exchanges", "BSE Limited", | |
| "National Stock Exchange of India Limited", "ISSUER'S" | |
| ], | |
| help="List of terms to whitelist (ignore during PII detection)." | |
| ) | |
| args = parser.parse_args() | |
| if not os.path.exists(args.input): | |
| print(f"Error: Input file '{args.input}' not found.") | |
| sys.exit(1) | |
| # Initialize the core detector with whitelist | |
| detector = HybridDetector(whitelist=args.whitelist) | |
| if args.mode == "redact": | |
| if not args.output: | |
| print("Error: --output is required for 'redact' mode.") | |
| sys.exit(1) | |
| print(f"Starting redaction on {args.input}...") | |
| # Initialize mapper (seed can be specified if reproducible fakes are needed) | |
| mapper = StatefulMapper() | |
| # Initialize DOCX handler | |
| handler = DocxHandler(detector=detector, mapper=mapper) | |
| # Process and save | |
| handler.process_document(args.input, args.output) | |
| elif args.mode == "evaluate": | |
| print(f"Evaluating metrics using ground truth from {args.input}...") | |
| # Initialize evaluator | |
| evaluator = Evaluator(detector=detector) | |
| # Run evaluation | |
| evaluator.evaluate_from_file(args.input) | |
| if __name__ == "__main__": | |
| main() | |