Spaces:
Running
Running
| """ | |
| utils.py - Utility functions for the project. | |
| """ | |
| import re | |
| from pathlib import Path | |
| from datetime import datetime | |
| from natsort import natsorted | |
| import subprocess | |
| def get_timestamp()->str: | |
| """ | |
| get_timestamp - get a timestamp for the current time | |
| Returns: | |
| str, the timestamp | |
| """ | |
| return datetime.now().strftime("%Y%m%d_%H%M%S") | |
| def truncate_word_count(text, max_words=512): | |
| """ | |
| truncate_word_count - a helper function for the gradio module | |
| Parameters | |
| ---------- | |
| text : str, required, the text to be processed | |
| max_words : int, optional, the maximum number of words, default=512 | |
| Returns | |
| ------- | |
| dict, the text and whether it was truncated | |
| """ | |
| # split on whitespace with regex | |
| words = re.split(r"\s+", text) | |
| processed = {} | |
| if len(words) > max_words: | |
| processed["was_truncated"] = True | |
| processed["truncated_text"] = " ".join(words[:max_words]) | |
| else: | |
| processed["was_truncated"] = False | |
| processed["truncated_text"] = text | |
| return processed | |
| def load_examples(src, filetypes=[".txt", ".pdf"]): | |
| """ | |
| load_examples - a helper function for the gradio module to load examples | |
| Returns: | |
| list of str, 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 saves_summary(summarize_output, outpath:str or Path=None, add_signature=True): | |
| """ | |
| saves_summary - save the summary generated from summarize_via_tokenbatches() to a text file | |
| _summaries = summarize_via_tokenbatches( | |
| text, | |
| batch_length=token_batch_length, | |
| batch_stride=batch_stride, | |
| **settings, | |
| ) | |
| """ | |
| outpath = Path.cwd() / f"document_summary_{get_timestamp()}.txt" if outpath is None else Path(outpath) | |
| sum_text = [s["summary"][0] 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\t".join(sum_text) | |
| with open( | |
| outpath, | |
| "w", | |
| ) as fo: | |
| if add_signature: | |
| fo.write( | |
| "Generated with the Document Summarization space :) https://hf.co/spaces/pszemraj/document-summarization\n\n" | |
| ) | |
| fo.writelines(full_summary) | |
| with open( | |
| outpath, | |
| "a", | |
| ) as fo: | |
| fo.write("\n" * 3) | |
| fo.write(f"\n\nSection Scores:\n") | |
| fo.writelines(scores_text) | |
| fo.write("\n\n---\n") | |
| return outpath | |