pszemraj/document-summarization

1# -*- coding: utf-8 -*-2"""3pdf2text.py - convert pdf files to text files using OCR4"""5import logging6import os7import re8import shutil9import time10from datetime import date11from os.path import join12from pathlib import Path13 14logging.basicConfig(15    level=logging.INFO,16    format="%(asctime)s %(levelname)s %(message)s",17    datefmt="%m/%d/%Y %I:%M:%S",18)19 20 21os.environ["USE_TORCH"] = "1"22 23from cleantext import clean24from doctr.io import DocumentFile25from doctr.models import ocr_predictor26from spellchecker import SpellChecker27 28 29def simple_rename(filepath, target_ext=".txt"):30    """simple_rename - get a new str to rename a file"""31    _fp = Path(filepath)32    basename = _fp.stem33    return f"OCR_{basename}_{target_ext}"34 35 36def rm_local_text_files(name_contains="RESULT_"):37    """38    rm_local_text_files - remove local text files39    """40    files = [41        f42        for f in Path.cwd().iterdir()43        if f.is_file() and f.suffix == ".txt" and name_contains in f.name44    ]45    logging.info(f"removing {len(files)} text files")46    for f in files:47        os.remove(f)48    logging.info("done")49 50 51def corr(52    s: str,53    add_space_when_numerics=False,54    exceptions=["e.g.", "i.e.", "etc.", "cf.", "vs.", "p."],55) -> str:56    """corrects spacing in a string57 58    Args:59        s (str): the string to correct60        add_space_when_numerics (bool, optional): [add a space when a period is between two numbers, example 5.73]. Defaults to False.61        exceptions (list, optional): [do not change these substrings]. Defaults to ['e.g.', 'i.e.', 'etc.', 'cf.', 'vs.', 'p.'].62 63    Returns:64        str: the corrected string65    """66    if add_space_when_numerics:67        s = re.sub(r"(\d)\.(\d)", r"\1. \2", s)68 69    s = re.sub(r"\s+", " ", s)70    s = re.sub(r'\s([?.!"](?:\s|$))', r"\1", s)71 72    # fix space before apostrophe73    s = re.sub(r"\s\'", r"'", s)74    # fix space after apostrophe75    s = re.sub(r"'\s", r"'", s)76    # fix space before comma77    s = re.sub(r"\s,", r",", s)78 79    for e in exceptions:80        expected_sub = re.sub(r"\s", "", e)81        s = s.replace(expected_sub, e)82 83    return s84 85 86def fix_punct_spaces(string: str) -> str:87    """88    fix_punct_spaces - fix spaces around punctuation89 90    :param str string: input string91    :return str: string with spaces fixed92    """93 94    fix_spaces = re.compile(r"\s*([?!.,]+(?:\s+[?!.,]+)*)\s*")95    string = fix_spaces.sub(lambda x: "{} ".format(x.group(1).replace(" ", "")), string)96    string = string.replace(" ' ", "'")97    string = string.replace(' " ', '"')98    return string.strip()99 100 101def clean_OCR(ugly_text: str) -> str:102    """103    clean_OCR - clean up the OCR text104 105    :param str ugly_text: input text to be cleaned106    :return str: cleaned text107    """108    # Remove all the newlines.109    cleaned_text = ugly_text.replace("\n", " ")110    # Remove all the tabs.111    cleaned_text = cleaned_text.replace("\t", " ")112    # Remove all the double spaces.113    cleaned_text = cleaned_text.replace("  ", " ")114    # Remove all the spaces at the beginning of the text.115    cleaned_text = cleaned_text.lstrip()116    # remove all instances of "- " and " - "117    cleaned_text = cleaned_text.replace("- ", "")118    cleaned_text = cleaned_text.replace(" -", "")119    return fix_punct_spaces(cleaned_text)120 121 122def move2completed(123    from_dir, filename, new_folder: str = "completed", verbose: bool = False124):125    """126    move2completed - move a file to a new folder127    """128    old_filepath = join(from_dir, filename)129 130    new_filedirectory = join(from_dir, new_folder)131 132    if not os.path.isdir(new_filedirectory):133        os.mkdir(new_filedirectory)134        if verbose:135            print("created new directory for files at: \n", new_filedirectory)136    new_filepath = join(new_filedirectory, filename)137 138    try:139        shutil.move(old_filepath, new_filepath)140        logging.info("successfully moved the file {} to */completed.".format(filename))141    except:142        logging.info(143            "ERROR! unable to move file to \n{}. Please investigate".format(144                new_filepath145            )146        )147 148 149custom_replace_list = {150    "t0": "to",151    "'$": "'s",152    ",,": ", ",153    "_ ": " ",154    " '": "'",155}156 157replace_corr_exceptions = {158    "i. e.": "i.e.",159    "e. g.": "e.g.",160    "e. g": "e.g.",161    " ,": ",",162}163 164 165spell = SpellChecker()166 167 168def check_word_spelling(word: str) -> bool:169    """170    check_word_spelling - check the spelling of a word171 172    Args:173        word (str): word to check174 175    Returns:176        bool: True if word is spelled correctly, False if not177    """178 179    misspelled = spell.unknown([word])180 181    return len(misspelled) == 0182 183 184def eval_and_replace(text: str, match_token: str = "- ") -> str:185    """186    eval_and_replace  - conditionally replace all instances of a substring in a string based on whether the eliminated substring results in a valid word187 188    Args:189        text (str): text to evaluate190        match_token (str, optional): token to replace. Defaults to "- ".191 192    Returns:193        str:  text with replaced tokens194    """195 196    if match_token not in text:197        return text198    else:199        while True:200            full_before_text = text.split(match_token, maxsplit=1)[0]201            before_text = [202                char for char in full_before_text.split()[-1] if char.isalpha()203            ]204            before_text = "".join(before_text)205            full_after_text = text.split(match_token, maxsplit=1)[-1]206            after_text = [char for char in full_after_text.split()[0] if char.isalpha()]207            after_text = "".join(after_text)208            full_text = before_text + after_text209            if check_word_spelling(full_text):210                text = full_before_text + full_after_text211            else:212                text = full_before_text + " " + full_after_text213            if match_token not in text:214                break215        return text216 217 218def cleantxt_ocr(ugly_text, lower=False, lang: str = "en") -> str:219    """220    cleantxt_ocr - clean text from OCR221 222        https://pypi.org/project/clean-text/223    Args:224        ugly_text (str): text to clean225        lower (bool, optional): lowercase text. Defaults to False.226        lang (str, optional): language of text. Defaults to "en".227 228    Returns:229        str: cleaned text230    """231 232    cleaned_text = clean(233        ugly_text,234        fix_unicode=True,  # fix various unicode errors235        to_ascii=True,  # transliterate to closest ASCII representation236        lower=lower,  # lowercase text237        no_line_breaks=True,  # fully strip line breaks as opposed to only normalizing them238        no_urls=True,  # replace all URLs with a special token239        no_emails=True,  # replace all email addresses with a special token240        no_phone_numbers=True,  # replace all phone numbers with a special token241        no_numbers=False,  # replace all numbers with a special token242        no_digits=False,  # replace all digits with a special token243        no_currency_symbols=False,  # replace all currency symbols with a special token244        no_punct=False,  # remove punctuations245        replace_with_punct="",  # instead of removing punctuations you may replace them246        replace_with_url="this url",247        replace_with_email="this email",248        replace_with_phone_number="this phone number",249        lang=lang,  # set to 'de' for German special handling250    )251 252    return cleaned_text253 254 255def format_ocr_out(OCR_data):256    """format OCR output to text"""257    if isinstance(OCR_data, list):258        text = " ".join(OCR_data)259    else:260        text = str(OCR_data)261    _clean = cleantxt_ocr(text)262    return corr(_clean)263 264 265def postprocess(text: str) -> str:266    """to be used after recombining the lines"""267 268    proc = corr(cleantxt_ocr(text))269 270    for k, v in custom_replace_list.items():271        proc = proc.replace(str(k), str(v))272 273    proc = corr(proc)274 275    for k, v in replace_corr_exceptions.items():276        proc = proc.replace(str(k), str(v))277 278    return eval_and_replace(proc)279 280 281def result2text(result, as_text=False) -> str or list:282    """Convert OCR result to text"""283 284    full_doc = []285    for i, page in enumerate(result.pages, start=1):286        text = ""287        for block in page.blocks:288            text += "\n\t"289            for line in block.lines:290                for word in line.words:291                    # print(dir(word))292                    text += word.value + " "293        full_doc.append(text)294 295    return "\n".join(full_doc) if as_text else full_doc296 297 298def convert_PDF_to_Text(299    PDF_file,300    ocr_model=None,301    max_pages: int = 20,302) -> str:303    """304    convert_PDF_to_Text - convert a PDF file to text305 306    :param str PDF_file: path to PDF file307    :param ocr_model: model to use for OCR, defaults to None (uses the default model)308    :param int max_pages: maximum number of pages to process, defaults to 20309    :return str: text from PDF310    """311    st = time.perf_counter()312    PDF_file = Path(PDF_file)313    ocr_model = ocr_predictor(pretrained=True) if ocr_model is None else ocr_model314    logging.info(f"starting OCR on {PDF_file.name}")315    doc = DocumentFile.from_pdf(PDF_file)316    truncated = False317    if len(doc) > max_pages:318        logging.warning(319            f"PDF has {len(doc)} pages, which is more than {max_pages}.. truncating"320        )321        doc = doc[:max_pages]322        truncated = True323 324    # Analyze325    logging.info(f"running OCR on {len(doc)} pages")326    result = ocr_model(doc)327    raw_text = result2text(result)328    proc_text = [format_ocr_out(r) for r in raw_text]329    fin_text = [postprocess(t) for t in proc_text]330 331    ocr_results = "\n\n".join(fin_text)332 333    fn_rt = time.perf_counter() - st334 335    logging.info("OCR complete")336 337    results_dict = {338        "num_pages": len(doc),339        "runtime": round(fn_rt, 2),340        "date": str(date.today()),341        "converted_text": ocr_results,342        "truncated": truncated,343        "length": len(ocr_results),344    }345 346    return results_dict347