pszemraj's picture
⚰️ 🎨 clean up and rm verbose testing code
c1cba4f
raw
history blame
16.3 kB
"""
utils.py - Utility functions for the project.
"""
import logging
import os
import re
import string
import subprocess
from collections import defaultdict, deque
from datetime import datetime, timedelta
from itertools import combinations, islice
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"""
<div style="
margin-bottom: 20px;
font-size: 18px;
line-height: 1.5em;
color: #333;
">
<h2 style="font-size: 22px; color: #555;">Batch {i}:</h2>
<p style="white-space: pre-line;">{s}</p>
</div>
"""
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"""
<div style="
border: 1px solid #ddd;
border-radius: 5px;
padding: 20px;
">
{joined_batches}
</div>
"""
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'<h2 style="font-size: 22px; color: #555;">(.*?)</h2>\s*<p style="white-space: pre-line;">(.*?)</p>'
# 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(f"## 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 outpath