""" utils.py - Utility functions for the project. """ import logging import re import subprocess from collections import defaultdict, deque from datetime import datetime, timedelta from itertools import combinations from pathlib import Path from typing import List logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", level=logging.INFO, ) import torch from natsort import natsorted from nltk.tokenize import WhitespaceTokenizer, sent_tokenize, word_tokenize from rapidfuzz import fuzz STOPWORDS = set( "a about above after again all also am an and any are aren't as at back be because been before being below between both but by can't cannot could couldn't did didn't do does doesn't doing don't down during each few for from further had hadn't has hasn't have haven't having he'd he'll he's hence her here here's hers herself him himself his how how's however i'd i'll i'm i've if in into is isn't it's its itself just let's me more moreover most mustn't my myself new nor now of off on once only or other ought our ours ourselves out over own really same shan't she'd she'll she's should shouldn't so some such than that that's the their theirs them themselves then there there's therefore these they they'd they'll they're they've this those through thus to too under until up use used using very was wasn't we we'd we'll we're we've were weren't what what's when when's where where's which while who who's whom why why's with won't would wouldn't you'd you'll you're you've your yours yourself yourselves".split() ) def contraction_aware_tokenize(text: str) -> List[str]: """contraction_aware_tokenize - merges words containing apostrophes as one token.""" # Tokenize the text using the WhitespaceTokenizer tokenizer = WhitespaceTokenizer() tokens = tokenizer.tokenize(text) merged_tokens = [] merged_token = "" for token in tokens: if re.search(r"\w+'\w+", token): # Token contains an apostrophe, merge with previous token merged_token += token else: # no apostrophe, add previous merged token (if any) and current if merged_token: merged_tokens.append(merged_token) merged_token = "" merged_tokens.append(token) # Add the last merged token (if any) if merged_token: merged_tokens.append(merged_token) return merged_tokens def remove_stopwords( text: str, stopwords: List[str] = STOPWORDS, contraction_tokenize: bool = True ) -> str: """ remove_stopwords - Remove stopwords from text. :param str text: input text :param List[str] stopwords: list of stopwords, defaults to STOPWORDS :param bool contraction_tokenize: use custom apostrophe tokenizer, defaults to True :return str: text with stopwords removed """ lines = text.split("\n") filtered_lines = [] def fix_commas(text: str) -> str: """fixes commas in text to have a space after them""" spaced_text = text.replace(",", ", ") return spaced_text.replace(" ", " ").strip() for line in lines: sentences = sent_tokenize(line) filtered_sentences = [] for sentence in sentences: # Add space around punctuations for the regex to work correctly, only if they are followed by a letter sentence_with_spaces = re.sub(r"([.,!?])(\w)", r"\1 \2", sentence[:-1]) words = ( contraction_aware_tokenize(sentence_with_spaces) if contraction_tokenize else word_tokenize(sentence_with_spaces) ) filtered_words = [] for word in words: if word.lower() not in stopwords: filtered_words.append(word) filtered_sentence = " ".join(filtered_words) # Restore original spaces around punctuation marks filtered_sentence = re.sub(r"([.,!?])\s*", r"\1", filtered_sentence) filtered_sentences.append(filtered_sentence + sentence[-1]) filtered_line = " ".join(filtered_sentences) # Replace multiple consecutive whitespaces with a single space filtered_line = re.sub(r"\s+", " ", filtered_line) filtered_line = fix_commas(filtered_line.strip()) filtered_lines.append(filtered_line) filtered_text = "\n".join(filtered_lines) return filtered_text def remove_stagnant_files( freq: str = "hourly", search_path: str = ".", substring="DocSumm", remove_suffix=".txt", ): """ remove_stagnant_files - Remove files that have not been modified in a certain amount of time. :param str freq: frequency of file removal, defaults to "hourly" :param str search_path: location to search for files, defaults to "." :param str substring: substring to search for in file names, defaults to "DocSumm" :param str remove_suffix: suffix of files to remove, defaults to ".txt" :raises ValueError: if freq is not one of "hourly", "daily", or "weekly" """ current_time = datetime.now() search_path = Path(search_path) if freq == "hourly": time_threshold = current_time - timedelta(hours=1) elif freq == "daily": time_threshold = current_time - timedelta(days=1) elif freq == "weekly": time_threshold = current_time - timedelta(weeks=1) else: raise ValueError( "Invalid frequency. Supported values are 'hourly', 'daily', and 'weekly'." ) files_to_remove = [] potential_files = [ f for f in search_path.iterdir() if f.is_file() and f.suffix == remove_suffix ] logging.info(f"Found {len(potential_files)} files.") for candidate in potential_files: if ( candidate.is_file() and substring in candidate.name and candidate.stat().st_mtime < time_threshold.timestamp() ): files_to_remove.append(candidate) logging.debug(f"File {candidate} last modified at {candidate.stat().st_mtime}") logging.info(f"Removing {len(files_to_remove)} files.") for file_path in files_to_remove: file_path.unlink() logging.debug(f"Removed files: {files_to_remove}") def compare_model_size(model_name: str, threshold: int = 500) -> bool: """ compare_model_size - compare string representations of model size to a threshold :param str model_name: the model name to compare :param int threshold: the threshold to compare against in millions, defaults to 500 :return: True if the model size is greater than the threshold, False or None otherwise """ pattern = r"(\d+)(M|G|k|b)?" # param regex matches = re.findall(pattern, model_name) if not matches: return None # Extract the parameter count and unit parameter_count, unit = matches[-1] parameter_count = int(parameter_count) # Convert to the standard form (M for million, G for billion, k for thousand) if unit == "G" or unit == "b": parameter_count *= 1000 elif unit == "M": pass elif unit == "k": parameter_count /= 1000 else: return None # Unknown return parameter_count > threshold def validate_pytorch2(torch_version: str = None) -> bool: """ validate_pytorch2 - validate that the PyTorch version is 2.0 or greater :param str torch_version: the PyTorch version to validate, defaults to None :return: True if the PyTorch version is 2.0 or greater, False otherwise """ torch_version = torch.__version__ if torch_version is None else torch_version pattern = r"^2\.\d+(\.\d+)*" return True if re.match(pattern, torch_version) else False def get_timestamp(detailed=False) -> str: """ get_timestamp - get a timestamp for the current time :param bool detailed: whether to include seconds and microseconds, defaults to False :return: str, the timestamp """ return ( datetime.now().strftime("%b%d%Y_%H%M%S%f") if detailed else datetime.now().strftime("%b%d%Y_%H") ) def truncate_word_count(text: str, max_words=1024) -> dict: """ truncate_word_count - truncate a text to a maximum number of words :param str text: the text to truncate :param int max_words: the maximum number of words to keep, defaults to 1024 :return: dict, the processed text """ words = contraction_aware_tokenize(str(text)) processed = {} if len(words) > max_words: processed["was_truncated"] = True processed["processed_text"] = " ".join(words[:max_words]) else: processed["was_truncated"] = False processed["processed_text"] = text return processed def load_examples(src, filetypes=[".txt", ".pdf"]): """ load_examples - a helper function for the gradio module to load examples :param str src: the path to the examples """ src = Path(src) src.mkdir(exist_ok=True) pdf_url = ( "https://www.dropbox.com/s/y92xy7o5qb88yij/all_you_need_is_attention.pdf?dl=1" ) subprocess.run(["wget", pdf_url, "-O", src / "all_you_need_is_attention.pdf"]) examples = [f for f in src.iterdir() if f.suffix in filetypes] examples = natsorted(examples) # load the examples into a list text_examples = [] for example in examples: with open(example, "r") as f: text = f.read() text_examples.append([text, "base", 2, 1024, 0.7, 3.5, 3]) return text_examples def load_example_filenames(example_path: str or Path): """ load_example_filenames - a helper function for the gradio module to load examples Returns: dict, the examples (filename:full path) """ example_path = Path(example_path) # load the examples into a list examples = {f.name: f for f in example_path.glob("*.txt")} return examples def textlist2html(text_batches: List[str]) -> str: """textlist2html - convert a list of text summaries into a single HTML string""" # Step 1: Generate each summary batch as a string of HTML formatted_batches = [ f"""

Batch {i}:

{s}

""" for i, s in enumerate(text_batches, start=1) ] # Step 2: Join all the summary batches together into one string joined_batches = "".join(formatted_batches) # Step 3: Wrap the summary string in a larger div with background color, border, and padding text_html_block = f"""
{joined_batches}
""" return text_html_block def extract_batches(html_string: str, pattern=None, flags=None) -> list: """ Extract batches of text from an HTML string. Args: html_string (str): The HTML string to extract batches from. pattern (str, optional): The regular expression pattern to use. Defaults to a pattern that matches batches in the format provided. flags (int, optional): The flags to use with the regular expression. Defaults to re.DOTALL. Returns: list: A list of dictionaries where each dictionary represents a batch and has 'title' and 'content' keys. """ # Set default pattern if none provided if pattern is None: pattern = r'

(.*?)

\s*

(.*?)

' # Set default flags if none provided if flags is None: flags = re.DOTALL try: # Find all matches in the string matches = re.findall(pattern, html_string, flags) # Convert matches to a list of dictionaries batches = [ {"title": title.strip(), "content": content.strip()} for title, content in matches ] return batches except re.error as e: logging.error(f"An error occurred while trying to extract batches: {e}") return [] def extract_keywords( text: str, num_keywords: int = 3, window_size: int = 5, kw_max_len: int = 20 ) -> List[str]: """ Extracts keywords from a text using a simplified TextRank algorithm. Args: text: The text to extract keywords from. num_keywords: The number of keywords to extract. Default: 3 window_size: The number of words considered for co-occurrence. Default: 5 kw_max_len: The maximum length of a keyword (truncate longer keywords to max). Default: 20 Returns: A list of strings, where each string is a keyword extracted from the input text. """ logger = logging.getLogger(__name__) # Remove stopwords and tokenize the text into words words = [ word for word in re.findall(r"\b\w{3,}\b", text.lower()) if word not in STOPWORDS ] # Create a graph of word co-occurrences within a moving window of words cooccur = defaultdict(lambda: defaultdict(int)) deque_words = deque(maxlen=window_size) for word in words: for w1, w2 in combinations(deque_words, 2): cooccur[w1][w2] += 1 cooccur[w2][w1] += 1 deque_words.append(word) # Assign scores to words using a simplified TextRank algorithm scores = defaultdict(float) for _ in range(10): new_scores = defaultdict(float) for word, co_words in cooccur.items(): new_scores[word] = 0.15 + 0.85 * sum( cooccur[word][other] / sum(cooccur[other].values()) * scores[other] for other in co_words ) scores = new_scores # Sort the words by score and return the top num_keywords keywords keywords = sorted(scores, key=scores.get, reverse=True)[:num_keywords] logger.debug(f"All keywords: {keywords}") # Use fuzzy matching to remove similar keywords final_keywords = [] for keyword in keywords: if not any(fuzz.ratio(keyword, other) > 70 for other in final_keywords): final_keywords.append(keyword[:kw_max_len]) logger.debug(f"Keywords (max len. {kw_max_len}):\t{final_keywords}") return final_keywords def saves_summary( summarize_output, outpath: str or Path = None, add_signature=True, **kwargs ) -> Path: """ saves_summary - save the summary generated from summarize_via_tokenbatches() to a text file summarize_output: output from summarize_via_tokenbatches() outpath: path to the output file add_signature: whether to add a signature to the output file kwargs: additional keyword arguments to include in the output file """ logger = logging.getLogger(__name__) sum_text = [f"{s['summary'][0]}\n" for s in summarize_output] sum_scores = [f"\n - {round(s['summary_score'],4)}" for s in summarize_output] scores_text = "\n".join(sum_scores) full_summary = "\n".join(sum_text) keywords = "_".join(extract_keywords(full_summary, kw_max_len=4)) logger.debug(f"kw:\t{keywords}") outpath = ( Path.cwd() / f"DocSumm_{keywords}_{get_timestamp()}.txt" if outpath is None else Path(outpath) ) logger.info(f"Saving summary to:\t{outpath.name}") with open( outpath, "w", encoding="utf-8", ) as fo: fo.writelines(full_summary) fo.write("\n\n") if add_signature: fo.write("\n\n---\n\n") fo.write("Generated with the Document Summarization space :)\n\n") fo.write("https://hf.co/spaces/pszemraj/document-summarization\n\n") with open( outpath, "a", encoding="utf-8", ) as fo: fo.write("\n") fo.write("## Section Scores:\n\n") fo.writelines(scores_text) fo.write("\n\n") fo.write(f"Date: {get_timestamp()}\n\n") if kwargs: fo.write("---\n\n") fo.write("## Parameters:\n\n") for key, value in kwargs.items(): fo.write(f"{key}: {value}\n") return str(outpath.resolve())