From d8c642e57682da5a664bfe85581f03b1fd3faa82 Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Fri, 24 Mar 2023 18:56:52 -0400 Subject: [PATCH] modified dataset storage of metadata --- src/dataset.py | 61 ++++++++++++++++++++++++++++++++------------ src/text_splitter.py | 3 +-- 2 files changed, 46 insertions(+), 18 deletions(-) diff --git a/src/dataset.py b/src/dataset.py index 19a11db..7a2b48b 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -56,7 +56,8 @@ class Dataset: self.metadata: List[Tuple[str]] = [] # List of tuples, each containing the title of an article, its URL, and text. E.g.: [('title', 'url', 'text'), ...] self.embedding_strings: List[str] = [] # List of strings, each being a few paragraphs from a single article (not exceeding 1000 words). - + self.embeddings_metadata: List[int] # List of integers, each being the index of the article from which the embedding string was taken. + self.articles_count: DefaultDict[str, int] = defaultdict(int) # Number of articles per source. E.g.: {'source1': 10, 'source2': 20, 'total': 30} if self.custom_sources is not None: @@ -82,28 +83,28 @@ class Dataset: Returns: Tuple[str]: a tuple containing the title, author, date, URL, tags, and text of the article. """ - title = None - author = None - date_published = None - url = None - tags = None - text = None + title: str = "" + author: str = "" + date_published: str = None + url: str = None + tags: str = None + text: str = None # Get title if 'title' in article and 'book_title' in article and article['title']: title = article['title'] elif 'book_title' in article and 'title' not in article and article['book_title']: title = article['book_title'] - if title[-1] == '\n': title = title[:-1] elif 'title' in article and article['title']: title = article['title'] - if title[-1] == '\n': title = title[:-1] - else: title = None + title = title.strip('\n').replace('\n', ' ')[:100] # Get author if 'author' in article and 'authors' in article and article['author']: author = article['author'] elif 'authors' in article and article['authors']: author = article['authors'] elif 'author' in article and article['author']: author = article['author'] - else: author = None + if type(author) == str: author = get_authors_list(author) + if type(author) == list: author = ', '.join(author) + author = author.strip('\n').replace('\n', ' ')[:100] # Get date published if 'date_published' in article and article['date_published'] and len(article['date_published']) >= 10: date_published = article['date_published'][:10] @@ -156,6 +157,7 @@ class Dataset: # if we specified custom sources, only include articles from those sources if (self.custom_sources is not None) and (entry['source'] not in self.custom_sources): continue + self.articles_count[entry['source']] += 1 self.total_articles_count += 1 @@ -165,16 +167,20 @@ class Dataset: # Get signature signature = "" if title: signature += f"Title: {title}, " - if author: signature += f"Author: {author}, " - if date_published: signature += f"Date published: {date_published}, " - if url: signature += f"URL: {url}, " + else: signature += f"Title: None, " + if author: signature += f"Author: {author}" + else: signature += f"Author: None" + # if date_published: signature += f"Date published: {date_published}, " + # if url: signature += f"URL: {url}, " # if tags: signature += f"Tags: {tags}, " # Temporary decision to not include tags in the signature - if signature: signature = signature[:-2] + # if signature: signature = signature[:-2] + signature = signature.replace("\n", " ") # Add info to metadata and embedding strings - self.metadata.append((title, author, date_published, url, tags, text)) + self.metadata.append((title, author, date_published, url, tags)) blocks = text_splitter.split(text, signature) self.embedding_strings.extend(blocks) + self.embeddings_metadata.extend([self.total_articles_count] * len(blocks)) # Update counts self.total_char_count += len(text) @@ -222,6 +228,29 @@ class Dataset: with open(path, 'wb') as f: pickle.dump(self, f) + + + + +def get_authors_list(authors_string: str) -> List[str]: + """ + Given a string of authors, return a list of the authors, even if the string contains a single author. + """ + authors_string = authors_string.replace(" and ", ",") + authors_string = authors_string.replace('\n', ' ') + authors = [] + if authors_string is None: + return [] + if "," in authors_string: + authors = [author.strip() for author in authors_string.split(",")] + else: + authors = [authors_string.strip()] + return authors + + + + + if __name__ == "__main__": # List of possible sources: all_sources = ["https://aipulse.org", "ebook", "https://qualiacomputing.com", "alignment forum", "lesswrong", "manual", "arxiv", "https://deepmindsafetyresearch.medium.com", "waitbutwhy.com", "GitHub", "https://aiimpacts.org", "arbital.com", "carado.moe", "nonarxiv_papers", "https://vkrakovna.wordpress.com", "https://jsteinhardt.wordpress.com", "audio-transcripts", "https://intelligence.org", "youtube", "reports", "https://aisafety.camp", "curriculum", "https://www.yudkowsky.net", "distill", diff --git a/src/text_splitter.py b/src/text_splitter.py index f720ae0..5285607 100644 --- a/src/text_splitter.py +++ b/src/text_splitter.py @@ -98,13 +98,12 @@ class TokenSplitter: # Check all block elements are strings assert all([isinstance(block, str) for block in blocks]), "block elements are not strings" - output = [f"{block}\n - {signature}" for block in blocks] + output = [f'"{block}"\n- {signature}' for block in blocks] # Check all output elements are strings assert all([isinstance(block, str) for block in output]), "output elements are not strings" return output - if __name__ == "__main__": text = """This post has been recorded as part of the LessWrong Curated Podcast, and an be listened to on Spotify, Apple Podcasts, and Libsyn.