From a25f217530623ee7b8d0d67b696e73a0d4d8c7bd Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Wed, 21 Jun 2023 00:30:30 -0400 Subject: [PATCH 01/49] simplified dataset code + uses huggingface dataset --- .gitignore | 2 + src/dataset/create_dataset.py | 416 ++++++++++------------------------ src/dataset/settings.py | 3 +- src/main.py | 98 ++------ 4 files changed, 142 insertions(+), 377 deletions(-) diff --git a/.gitignore b/.gitignore index 293af7e..3687207 100644 --- a/.gitignore +++ b/.gitignore @@ -129,6 +129,8 @@ dmypy.json .pyre/ # Other +*test.py +*test.ipynb *alignment_texts.jsonl *config.py *.DS_Store diff --git a/src/dataset/create_dataset.py b/src/dataset/create_dataset.py index 2456e44..8834c07 100644 --- a/src/dataset/create_dataset.py +++ b/src/dataset/create_dataset.py @@ -1,4 +1,3 @@ -import jsonlines import numpy as np from typing import List, Dict, Tuple, DefaultDict, Any from collections import defaultdict @@ -11,6 +10,9 @@ from pathlib import Path from tqdm.auto import tqdm from dateutil.parser import parse, ParserError import openai +from datasets import load_dataset +from langchain.document_loaders import HuggingFaceDatasetLoader + try: import config @@ -19,18 +21,18 @@ except ImportError: openai.api_key = os.environ.get('OPENAI_API_KEY') -from .settings import PATH_TO_RAW_DATA, PATH_TO_DATASET_PKL, PATH_TO_DATASET_DICT_PKL, EMBEDDING_MODEL, LEN_EMBEDDINGS +from .settings import PATH_TO_DATASET_PKL, PATH_TO_DATASET_DICT_PKL, EMBEDDING_MODEL, LEN_EMBEDDINGS from .text_splitter import TokenSplitter, split_into_sentences error_count_dict = { + "Entry has no id.": 0, "Entry has no source.": 0, "Entry has no title.": 0, "Entry has no text.": 0, - "Entry has no URL.": 0, - "Entry has wrong citation level.": 0 + "Entry has non-string required key.": 0 } @@ -38,221 +40,129 @@ class MissingDataException(Exception): pass -class Dataset: +class ChunkedARD: def __init__(self, - jsonl_data_path: str = PATH_TO_RAW_DATA, # Path to the dataset .jsonl file. - custom_sources: List[str] = None, # List of sources to include, like "alignment forum", "lesswrong", "arxiv",etc. - rate_limit_per_minute: int = 3_500, # Rate limit for the OpenAI API. min_tokens_per_block: int = 300, # Minimum number of tokens per block. max_tokens_per_block: int = 400, # Maximum number of tokens per block. - fraction_of_articles_to_use: float = 1.0, # Fraction of articles to use. If 1.0, use all articles. - ): - self.jsonl_data_path = jsonl_data_path + custom_sources: List[str] = None, # List of sources to include, like "alignmentforum", "lesswrong", "arxiv",etc. + rate_limit_per_minute: int = 3_500, # Rate limit for the OpenAI API. + ): + self.min_tokens_per_block = min_tokens_per_block # for the text splitter + self.max_tokens_per_block = max_tokens_per_block # for the text splitter + self.custom_sources = custom_sources self.rate_limit_per_minute = rate_limit_per_minute self.delay_in_seconds = 60.0 / self.rate_limit_per_minute - self.fraction_of_articles_to_use = fraction_of_articles_to_use - self.min_tokens_per_block = min_tokens_per_block # for the text splitter - self.max_tokens_per_block = max_tokens_per_block # for the text splitter - - self.metadata: List[Tuple[str]] = [] # List of tuples, each containing the title, author, date, URL, and tags of an article. - self.embedding_strings: List[str] = [] # List of strings, each being a few paragraphs from a single article (not exceeding max_tokens_per_block tokens). - self.embeddings_metadata_index: List[int] = [] # List of integers, each being the index of the article from which the embedding string was taken. + self.metadata: List[Dict[str, Any]] = [] # List of dicts, each containing: id, entry_id, source, title, text, url, date_published, authors. - self.articles_count: DefaultDict[str, int] = defaultdict(int) # Number of articles per source. E.g.: {'source1': 10, 'source2': 20, 'total': 30} + self.entries_per_source_count: DefaultDict[str, int] = defaultdict(int) # Number of entries per source. E.g.: {'source1': 10, 'source2': 20, 'total': 30} + self.total_counts = { + 'chars': 0, + 'words': 0, + 'sentences': 0, + 'chunks': 0, + 'entries': 0, + } if self.custom_sources is not None: for source in self.custom_sources: - self.articles_count[source] = 0 - self.total_articles_count = 0 + self.entries_per_source_count[source] = 0 - self.total_char_count = 0 - self.total_word_count = 0 - self.total_sentence_count = 0 - self.total_block_count = 0 - - self.sources_so_far: List[str] = [] - self.info_types: Dict[str, List[str]] = {} - - def extract_info_from_article(self, article: Dict[str, Any]) -> Tuple[str]: - """ - This function extracts the title, author, date, URL, tags, and text from an article. - - Args: - article (Dict[str, Any]): a dictionary containing the article's text and metadata. + def contains_required_metadata(self, entry: Dict[str, Any]): + metadata_types = { + 'id': str, + 'source': str, + 'title': str, + 'url': str, + 'date_published': str, + 'authors': list, # It appears to be a list, but actually it's a list-looking string. Like "['apple', 'orange', 'tomato']" is a string. +# 'summary': str # see previous comment + } + required_metadata_keys = ['id', 'source', 'title', 'text'] - Returns: - Tuple[str]: a tuple containing the title, author, date, URL, tags, and text of the article. - """ - 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'] - elif 'title' in article and article['title']: - title = article['title'] - title = title.strip('\n').replace('\n', ' ')[:100] + # Check that the 8 primary metadata keys all have the correct type + for key, key_type in metadata_types.items(): + if type(entry[key]) != key_type: + raise MissingDataException(f"Entry {entry['id']} has key {key} of type {type(entry[key])} when it should be {key_type}.") - # 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'] - if type(author) == str: author = get_authors_list(author) - if type(author) == list: author = ', '.join(author) - author = author.strip('\n').replace('\n', ' ')[:100] + # Check that the 4 required metadata keys are non-empty + for key in required_metadata_keys: + if not entry[key]: + raise MissingDataException(f"Entry {entry['id']} has no {key}.") - # Get date published - if 'date_published' in article and article['date_published'] and len(article['date_published']) >= 10: date_published = article['date_published'][:10] - elif 'published' in article and article['published'] and len(article['published']) >= 16: date_published = article['published'][:16] - else: date_published = None - if date_published is not None: - date_published = standardize_date(date_published) - - # Get URL - if 'link' in article and article['link']: url = article['link'] - elif 'url' in article and article['url']: url = article['url'] - elif 'doi' in article and article['doi']: url = article['doi'] - else: url = None - - # Get tags - if 'tags' in article and article['tags']: - if type(article['tags']) == list: tags = ', '.join([val['term'] for val in article['tags']]) - elif type(article['tags']) == str: tags = article['tags'] - else: tags = None - - # Get text - if 'text' in article and article['text']: text = article['text'] - else: - raise MissingDataException(f"Entry has no text.") - - return (title, author, date_published, url, tags, text) def get_alignment_texts(self): text_splitter = TokenSplitter(self.min_tokens_per_block, self.max_tokens_per_block) - with jsonlines.open(self.jsonl_data_path, "r") as reader: - for entry in tqdm(reader): - try: - if 'source' not in entry: - if 'url' in entry and entry['url'] == "https://www.cold-takes.com/": - entry["source"] = "Cold Takes" - elif 'question' in entry and 'answer' in entry: - entry["source"] = "printouts" - continue # for now, skip printouts - elif 'article_url' in entry and entry['article_url'] == "https://www.gwern.net": - entry["source"] = "gwern.net" - elif 'url' in entry and entry['url'] == "https://generative.ink/posts/": - entry["source"] = "generative.ink" - elif 'url' in entry and entry['url'][:24] == "https://greaterwrong.com": - entry["source"] = "greaterwrong.com" - else: - raise MissingDataException("Entry has no source.") - - # 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 - - if entry["source"] == 'alignment forum': - if int(entry['score'].replace('−', '-')) < 70: continue - elif entry["source"] == 'lesswrong': - if int(entry['score'].replace('−', '-')) < 150: continue - elif entry["source"] == 'arxiv': - if 'citation_level' != '0': continue + # Load the dataset. streaming allows you to load one entry at a time, + # so entries can be processed before the entire dataset has been saved. + iterable_data = load_dataset('StampyAI/alignment-research-dataset', 'aisafety.info', split='train', streaming=True) + + for entry in tqdm(iterable_data): + """Checks""" - # Dict describing the proportion of each source we want: - # E.g.: {'arxiv': 0.5, 'youtube': 0.5, 'lesswrong': 1.0} - desired_source_proportions = { - "https://aipulse.org": 1, - "ebook": 0, - "https://qualiacomputing.com": 0.02, - "alignment forum": .7, - "lesswrong": .5, - "manual": 1, - "arxiv": 0.1, - "https://deepmindsafetyresearch.medium.com/": 1, - "waitbutwhy.com": 1, - "GitHub": 1, - "https://aiimpacts.org": 0.2, - "arbital.com": 0.2, - "carado.moe": 0.3, - "nonarxiv_papers": 0.1, - "https://vkrakovna.wordpress.com": .5, - "https://jsteinhardt.wordpress.com": .5, - "audio-transcripts": 0.2, - "https://intelligence.org": .1, - "youtube": 0.07, - "reports": 0.4, - "https://aisafety.camp": 1, - "curriculum": 1, - "https://www.yudkowsky.net": 0.2, - "distill": 1, - "Cold Takes": 0.5, - "printouts": 1, - "gwern.net": 1, - "generative.ink": 1, - "greaterwrong.com": 0.2 - } - - random_number = random.random() - if random_number > desired_source_proportions[entry['source']]: - continue - - # if we specified a fraction of articles to use, only use that fraction from the remaining articles - random_number = random.random() - if random_number > self.fraction_of_articles_to_use: - continue - - # Get title, author, date, URL, tags, and text - title, author, date_published, url, tags, text = self.extract_info_from_article(entry) - - # If there's less than 2 of 'title', 'author' and 'url', ignore this text - if (((title or '').strip() == '') + ((author or '').strip() == '') + ((url or '').strip() == '')) > 1: - print(f'{entry["source"]}') - continue - - #if the text is too short, ignore this text - if len(text) < 500: - continue + # if we specified custom sources, only include entries from those sources + if (self.custom_sources is not None) and (entry['source'] not in self.custom_sources): + continue - #we're keeping the text so we inc the aticle count - self.articles_count[entry['source']] += 1 - self.total_articles_count += 1 - - # Get signature - signature = "" - if title: signature += f"Title: {title}, " - 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] - signature = signature.replace("\n", " ") - - # Add info to metadata and embedding strings - self.metadata.append((title, author, date_published, url, tags)) - blocks = text_splitter.split(text, signature) - self.embedding_strings.extend(blocks) - self.embeddings_metadata_index.extend([self.total_articles_count-1] * len(blocks)) - - # Update counts - self.total_char_count += len(text) - self.total_word_count += len(text.split()) - self.total_sentence_count += len(split_into_sentences(text)) - self.total_block_count += len(blocks) - - except MissingDataException as e: - if str(e) not in error_count_dict: - error_count_dict[str(e)] = 0 - error_count_dict[str(e)] += 1 + # raise error if the entry does not contain the required metadata + self.contains_required_metadata(entry) + + #if the text is too short, ignore this text + if len(entry['text']) < 500: + continue + + + """Checks are done, so we construct the metadata.""" + + # Get id, source, title, text, url, date_published, authors, and summary + entry_id: str = entry['id'] + source: str = entry['source'] + title: str = entry['title'] + text: str = entry['text'] + url: str = entry['url'] + date_published: str = entry['date_published'] + authors: list = entry['authors'] + + # Get signature + if authors: + signature = f"Title: {title}, Authors: {get_authors_str(authors)}" + else: + signature = f"Title: {title}" + + # We use the text_splitter to get the chunks from the entry, + # and we add a metadata element for each new chunk we add to the dataset. + chunks = text_splitter.split(text, signature) + num_chunks = len(chunks) + + for i in range(num_chunks): + self.metadata.append({ + 'id': f"{entry_id}_{str(i+1).zfill(6)}", + 'entry_id': entry_id, + 'source': source, + 'title': title, + 'text': chunks[i], + 'url': url, + 'date_published': date_published, + 'authors': authors + }) + + # Update counts + self.entries_per_source_count[entry['source']] += 1 + self.total_counts['entries'] += 1 + self.total_counts['chars'] += len(text) + self.total_counts['words'] += len(text.split()) + self.total_counts['sentences'] += len(split_into_sentences(text)) + self.total_counts['chunks'] += len(chunks) + + def show_stats(self): + print(f'Number of entries by source: {self.entries_per_source_count}') + print(f'Total entries count: {self.total_counts["entries"]}') + print(f'Total character count: {self.total_counts["chars"]}') + print(f'Total word count: {self.total_counts["words"]}') + print(f'Total sentence count: {self.total_counts["sentences"]}') + print(f'Total chunk count: {self.total_counts["chunks"]}') def get_embeddings(self): def get_embeddings_at_index(texts: str, batch_idx: int, batch_size: int = 200): # int, np.ndarray @@ -269,15 +179,16 @@ class Dataset: rate_limit = 3500 / 60 # Maximum embeddings per second start = time.time() - self.embeddings = np.zeros((len(self.embedding_strings), LEN_EMBEDDINGS)) + embedding_strings = [chunk_data['text'] for chunk_data in self.metadata] + self.embeddings = np.zeros((len(embedding_strings), LEN_EMBEDDINGS)) with concurrent.futures.ThreadPoolExecutor() as executor: futures = [executor.submit( get_embeddings_at_index, - self.embedding_strings[batch_idx:batch_idx+batch_size], + embedding_strings[batch_idx:batch_idx+batch_size], batch_idx, - len(self.embedding_strings[batch_idx:batch_idx+batch_size]) - ) for batch_idx in range(0, len(self.embedding_strings), batch_size)] + len(embedding_strings[batch_idx:batch_idx+batch_size]) + ) for batch_idx in range(0, len(embedding_strings), batch_size)] num_completed = 0 for future in concurrent.futures.as_completed(futures): batch_idx, embeddings = future.result() @@ -289,113 +200,32 @@ class Dataset: sleep_time = max(expected_time - elapsed_time, 0) time.sleep(sleep_time) - print(f"Completed {num_completed}/{len(self.embedding_strings)} embeddings in {elapsed_time:.2f} seconds.") + print(f"Completed {num_completed}/{len(embedding_strings)} embeddings in {elapsed_time:.2f} seconds.") - def save_embeddings(self, path: str): - np.save(path, self.embeddings) - - def load_embeddings(self, path: str): - self.embeddings = np.load(path) - - def save_class(self, path: str = PATH_TO_DATASET_PKL): - # Save the class to a pickle file - print(f"Saving class to {path}...") - with open(path, 'wb') as f: - pickle.dump(self, f) - def save_data(self, path: str = PATH_TO_DATASET_DICT_PKL): # Save the data to a pickle file print(f"Saving data to {path}...") data = { "metadata": self.metadata, - "embedding_strings": self.embedding_strings, - "embeddings_metadata_index": self.embeddings_metadata_index, "embeddings": self.embeddings.astype(np.float32), - "articles_count": self.articles_count, - "total_articles_count": self.total_articles_count, - "total_char_count": self.total_char_count, - "total_word_count": self.total_word_count, - "total_sentence_count": self.total_sentence_count, - "total_block_count": self.total_block_count + "total_counts": self.total_counts, + "entries_per_source_count": self.entries_per_source_count } with open(path, 'wb') as f: pickle.dump(data, 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(",")] +def get_authors_str(authors_lst: List[str]) -> str: + if authors_lst == []: return 'n/a' + if len(authors_lst) == 1: return authors_lst[0] else: - authors = [authors_string.strip()] - return authors + authors_lst = authors_lst[:3] + authors_str = ", ".join(authors_lst[:-1]) + " and " + authors_lst[-1] + return authors_str def standardize_date(date_string, default_date='n/a'): try: dt = parse(date_string) return dt.strftime('%Y-%m-%d') except (ParserError, ValueError): - return default_date - - - -""" -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", "Cold Takes", "printouts", "gwern.net", "generative.ink", "greaterwrong.com"] # These sources do not have a source field in the .jsonl file - - # List of sources we are using for the test run: - custom_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", - # "Cold Takes", - # "printouts", - # "gwern.net", - # "generative.ink", - # "greaterwrong.com" - ] - - dataset = Dataset( - jsonl_data_path=PATH_TO_RAW_DATA.resolve(), - custom_sources=custom_sources, - rate_limit_per_minute=3500, - min_tokens_per_block=200, max_tokens_per_block=300, - # fraction_of_articles_to_use=1/2000 - ) - dataset.get_alignment_texts() - dataset.get_embeddings() - # dataset.save_embeddings("data/embeddings.npy") - - dataset.save_class(PATH_TO_DATASET.resolve()) - # # dataset = pickle.load(open("dataset.pkl", "rb")) - """ - \ No newline at end of file + return default_date \ No newline at end of file diff --git a/src/dataset/settings.py b/src/dataset/settings.py index 28cb3da..abcb211 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -4,9 +4,8 @@ EMBEDDING_MODEL = "text-embedding-ada-002" COMPLETIONS_MODEL = "gpt-3.5-turbo" LEN_EMBEDDINGS = 1536 -MAX_LEN_PROMPT = 4095 # This may be 8191, unsure. +MAX_LEN_PROMPT = 8190 # This may be 8191, unsure. current_file_path = Path(__file__).resolve() -PATH_TO_RAW_DATA = str(current_file_path.parent / 'data' / 'alignment_texts.jsonl') PATH_TO_DATASET_PKL = str(current_file_path.parent / 'data' / 'dataset.pkl') PATH_TO_DATASET_DICT_PKL = str(current_file_path.parent / 'data' / 'dataset_dict.pkl') \ No newline at end of file diff --git a/src/main.py b/src/main.py index 900a14f..333795a 100644 --- a/src/main.py +++ b/src/main.py @@ -27,20 +27,14 @@ if str(src_path) not in sys.path: from dataset import create_dataset #from assistant import semantic_search -from settings import EMBEDDING_MODEL, PATH_TO_DATASET_DICT_PKL +from settings import EMBEDDING_MODEL, PATH_TO_DATASET_DICT_PKL, PATH_TO_DATASET_PKL import numpy as np import matplotlib.pyplot as plt def load_rawdata_into_pkl(): - """with open(PATH_TO_DATASET, 'rb') as f: - dataset = pickle.load(f) - AS = AlignmentSearch(dataset=dataset) - prompt = "What would be an idea to solve the Alignment Problem? Name the Lesswrong post by Quintin Pope that discusses this idea." - answer = AS.search_and_answer(prompt, 3, HyDE=False) - print(answer) - """ + """ # 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", "Cold Takes", "printouts", "gwern.net", "generative.ink", "greaterwrong.com"] # These last do not have a source field in the .jsonl file @@ -77,87 +71,27 @@ def load_rawdata_into_pkl(): "generative.ink", "greaterwrong.com" ] - - dataset = create_dataset.Dataset( - custom_sources=custom_sources, - rate_limit_per_minute=3500, - min_tokens_per_block=200, max_tokens_per_block=300, - # fraction_of_articles_to_use=1/150, + """ + + dataset = create_dataset.ChunkedARD( + min_tokens_per_block=200, max_tokens_per_block=300 ) dataset.get_alignment_texts() - - print(len(dataset.embedding_strings)) - print(dataset.total_word_count) - print(dataset.total_block_count) - print(dataset.articles_count) - dataset.get_embeddings() dataset.save_data() - + dataset.show_stats() + +def load_pkl_and_display_stuff(): + with open(PATH_TO_DATASET_DICT_PKL, 'rb') as f: + dataset_dict = pickle.load(f) + print(dataset_dict) + @retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(4)) def get_embedding(text: str) -> np.ndarray: result = openai.Embedding.create(model=EMBEDDING_MODEL, input=text) return np.array(result["data"][0]["embedding"]) -def print_out_dataset_stuff(): - with open(PATH_TO_DATASET_PKL, 'rb') as f: - dataset = pickle.load(f) - - embeddings_len = len(dataset.embedding_strings) - i1 = random.randint(0,embeddings_len-1) - #i2 = random.randint(0,embeddings_len-1) - #print(len(dataset.embeddings)) - #print(len(dataset.embedding_strings)) - #embedding_test = get_embedding(dataset.embedding_strings[i]) - #print(np.dot(embedding_test,dataset.embeddings[i])) - - metadata_i1 = dataset.embeddings_metadata_index[i1] - print("metadata:",dataset.metadata[metadata_i1]) - print("embedding_string:",dataset.embedding_strings[i1]) - #print("embedding_vector:",dataset.embeddings[i1]) - embedding_of_string1 = get_embedding(dataset.embedding_strings[i1]) - - #metadata_i2 = dataset.embeddings_metadata_index[i2] - #print("metadata:",dataset.metadata[metadata_i2]) - #print("embedding_string:",dataset.embedding_strings[i2]) - #print("embedding_vector:",dataset.embeddings[i1]) - #embedding_of_string2 = get_embedding(dataset.embedding_strings[i2]) - #embedding_of_string2 = get_embedding("000000000000000000000000000000000000000000000000000000000000000000000000000000000") - - #print(len(embedding_of_string1)) - vector = dataset.embeddings[i1] - plot_likelihood(vector) - #plot_likelihood(get_embedding("tst")) - print(max(vector), min(vector)) - print(sum([x**2 for x in vector])) - - - - - #print(np.dot(embedding_of_string1, embedding_of_string2)) - -def plot_likelihood(embeddings, num_buckets=200): - # Calculate the histogram - histogram, bin_edges = np.histogram(embeddings.flatten(), bins=num_buckets, range=(embeddings.min(), embeddings.max())) - - # Normalize the histogram to get likelihoods - likelihoods = histogram / embeddings.flatten().size - - # Plot the likelihoods - plt.bar(bin_edges[:-1], likelihoods, width=(bin_edges[1] - bin_edges[0]), edgecolor="k", alpha=0.7) - plt.xlabel("Value") - plt.ylabel("Likelihood") - plt.title("Likelihood of Floats in the Vector Embedding") - plt.savefig("bla.png") - - - - - - if __name__ == "__main__": - # load_rawdata_into_pkl() - # print_out_dataset_stuff() - - with open(PATH_TO_DATASET_DICT_PKL, 'rb') as f: - dataset = pickle.load(f) \ No newline at end of file + load_rawdata_into_pkl() + #load_pkl_and_display_stuff() + \ No newline at end of file From e9ebb25f88f31ab5cac60b0d53e64e82f65fce3f Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Wed, 21 Jun 2023 00:34:14 -0400 Subject: [PATCH 02/49] fixed minor errors from testing --- src/dataset/create_dataset.py | 4 +++- src/settings.py | 2 -- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/dataset/create_dataset.py b/src/dataset/create_dataset.py index 8834c07..516e4cb 100644 --- a/src/dataset/create_dataset.py +++ b/src/dataset/create_dataset.py @@ -97,7 +97,9 @@ class ChunkedARD: # Load the dataset. streaming allows you to load one entry at a time, # so entries can be processed before the entire dataset has been saved. - iterable_data = load_dataset('StampyAI/alignment-research-dataset', 'aisafety.info', split='train', streaming=True) + # iterable_data = load_dataset('StampyAI/alignment-research-dataset', 'aisafety.info', split='train', streaming=True) + + iterable_data = load_dataset('StampyAI/alignment-research-dataset', 'all', split='train', streaming=True) for entry in tqdm(iterable_data): """Checks""" diff --git a/src/settings.py b/src/settings.py index 2362cc7..94abeb5 100644 --- a/src/settings.py +++ b/src/settings.py @@ -7,6 +7,4 @@ LEN_EMBEDDINGS = 1536 MAX_LEN_PROMPT = 4095 # This may be 8191, unsure. current_file_path = Path(__file__).resolve() -PATH_TO_RAW_DATA = str(current_file_path.parent / 'dataset' / 'data' / 'alignment_texts.jsonl') -PATH_TO_DATASET_PKL = str(current_file_path.parent / 'dataset' / 'data' / 'dataset.pkl') PATH_TO_DATASET_DICT_PKL = str(current_file_path.parent / 'dataset' / 'data' / 'dataset_dict.pkl') \ No newline at end of file From a755ecc4801f47379b7a389bc4eaae96c0725527 Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Wed, 21 Jun 2023 00:36:14 -0400 Subject: [PATCH 03/49] removed PATH_TO_DATASET_PKL, stored in dict now --- src/dataset/create_dataset.py | 2 +- src/main.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/dataset/create_dataset.py b/src/dataset/create_dataset.py index 516e4cb..bdf5182 100644 --- a/src/dataset/create_dataset.py +++ b/src/dataset/create_dataset.py @@ -21,7 +21,7 @@ except ImportError: openai.api_key = os.environ.get('OPENAI_API_KEY') -from .settings import PATH_TO_DATASET_PKL, PATH_TO_DATASET_DICT_PKL, EMBEDDING_MODEL, LEN_EMBEDDINGS +from .settings import PATH_TO_DATASET_DICT_PKL, EMBEDDING_MODEL, LEN_EMBEDDINGS from .text_splitter import TokenSplitter, split_into_sentences diff --git a/src/main.py b/src/main.py index 333795a..029df5a 100644 --- a/src/main.py +++ b/src/main.py @@ -27,7 +27,7 @@ if str(src_path) not in sys.path: from dataset import create_dataset #from assistant import semantic_search -from settings import EMBEDDING_MODEL, PATH_TO_DATASET_DICT_PKL, PATH_TO_DATASET_PKL +from settings import EMBEDDING_MODEL, PATH_TO_DATASET_DICT_PKL import numpy as np import matplotlib.pyplot as plt From f2e160f8e33c99798113ee1aa453a216a68db2c2 Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Wed, 21 Jun 2023 00:37:16 -0400 Subject: [PATCH 04/49] removed commented-out code --- src/main.py | 39 --------------------------------------- 1 file changed, 39 deletions(-) diff --git a/src/main.py b/src/main.py index 029df5a..87a68e6 100644 --- a/src/main.py +++ b/src/main.py @@ -34,45 +34,6 @@ import matplotlib.pyplot as plt def load_rawdata_into_pkl(): - """ - # 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", - "Cold Takes", "printouts", "gwern.net", "generative.ink", "greaterwrong.com"] # These last do not have a source field in the .jsonl file - - # List of sources we are using for the test run: - custom_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", - "Cold Takes", - "printouts", - "gwern.net", - "generative.ink", - "greaterwrong.com" - ] - """ - dataset = create_dataset.ChunkedARD( min_tokens_per_block=200, max_tokens_per_block=300 ) From 5479ed7e93ffe1a2ad5ee39e8f9a60880c0cfe03 Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Wed, 21 Jun 2023 00:38:43 -0400 Subject: [PATCH 05/49] minor change to MAX_LEN_PROMPT --- src/dataset/settings.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/dataset/settings.py b/src/dataset/settings.py index abcb211..cf24c94 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -4,8 +4,7 @@ EMBEDDING_MODEL = "text-embedding-ada-002" COMPLETIONS_MODEL = "gpt-3.5-turbo" LEN_EMBEDDINGS = 1536 -MAX_LEN_PROMPT = 8190 # This may be 8191, unsure. +MAX_LEN_PROMPT = 4095 # This may be 8191, unsure. current_file_path = Path(__file__).resolve() -PATH_TO_DATASET_PKL = str(current_file_path.parent / 'data' / 'dataset.pkl') PATH_TO_DATASET_DICT_PKL = str(current_file_path.parent / 'data' / 'dataset_dict.pkl') \ No newline at end of file From 12c9db9cd0908613abed0da1c5fdff42ceccec90 Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Wed, 21 Jun 2023 00:40:26 -0400 Subject: [PATCH 06/49] load_dataset simpler than langchain's wrapper --- src/dataset/create_dataset.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/dataset/create_dataset.py b/src/dataset/create_dataset.py index bdf5182..0ec5891 100644 --- a/src/dataset/create_dataset.py +++ b/src/dataset/create_dataset.py @@ -11,7 +11,6 @@ from tqdm.auto import tqdm from dateutil.parser import parse, ParserError import openai from datasets import load_dataset -from langchain.document_loaders import HuggingFaceDatasetLoader try: From 9b8c470438808a29e9501e7ca67e7f68ecd6c125 Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Wed, 21 Jun 2023 00:41:54 -0400 Subject: [PATCH 07/49] removed error_count_dict --- src/dataset/create_dataset.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/dataset/create_dataset.py b/src/dataset/create_dataset.py index 0ec5891..6fe1e3f 100644 --- a/src/dataset/create_dataset.py +++ b/src/dataset/create_dataset.py @@ -25,16 +25,6 @@ from .settings import PATH_TO_DATASET_DICT_PKL, EMBEDDING_MODEL, LEN_EMBEDDINGS from .text_splitter import TokenSplitter, split_into_sentences - -error_count_dict = { - "Entry has no id.": 0, - "Entry has no source.": 0, - "Entry has no title.": 0, - "Entry has no text.": 0, - "Entry has non-string required key.": 0 -} - - class MissingDataException(Exception): pass From 8044111a1b11cd68fbd1a4a9f76754fdc3ebb500 Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Wed, 21 Jun 2023 00:45:40 -0400 Subject: [PATCH 08/49] commented change in authors' type from str to list --- src/dataset/create_dataset.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/dataset/create_dataset.py b/src/dataset/create_dataset.py index 6fe1e3f..ec20044 100644 --- a/src/dataset/create_dataset.py +++ b/src/dataset/create_dataset.py @@ -65,8 +65,8 @@ class ChunkedARD: 'title': str, 'url': str, 'date_published': str, - 'authors': list, # It appears to be a list, but actually it's a list-looking string. Like "['apple', 'orange', 'tomato']" is a string. -# 'summary': str # see previous comment + 'authors': list, # unsure as of yet if this is a string or a list. TODO it is subject to change +# 'summary': list # see previous comment } required_metadata_keys = ['id', 'source', 'title', 'text'] @@ -115,6 +115,7 @@ class ChunkedARD: url: str = entry['url'] date_published: str = entry['date_published'] authors: list = entry['authors'] + # summary is ignored for now, since most sources lack one. TODO: add summary. see self.metadata code as well. # Get signature if authors: From 3b667ce16a6002bfad4f5ff018b1477b885d3a29 Mon Sep 17 00:00:00 2001 From: Thomas Lemoine Date: Fri, 23 Jun 2023 15:00:50 -0400 Subject: [PATCH 09/49] simplified validation --- src/dataset/create_dataset.py | 48 ++++++++++++++++++++--------------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/src/dataset/create_dataset.py b/src/dataset/create_dataset.py index ec20044..ff3447a 100644 --- a/src/dataset/create_dataset.py +++ b/src/dataset/create_dataset.py @@ -58,39 +58,41 @@ class ChunkedARD: for source in self.custom_sources: self.entries_per_source_count[source] = 0 - def contains_required_metadata(self, entry: Dict[str, Any]): + @staticmethod + def _validate_required_metadata(entry: Dict[str, Any]): metadata_types = { 'id': str, 'source': str, 'title': str, 'url': str, 'date_published': str, - 'authors': list, # unsure as of yet if this is a string or a list. TODO it is subject to change -# 'summary': list # see previous comment + 'authors': list, + 'summary': list } - required_metadata_keys = ['id', 'source', 'title', 'text'] - # Check that the 8 primary metadata keys all have the correct type + # Check that the 8 primary metadata keys all have the correct type and the key exists for key, key_type in metadata_types.items(): - if type(entry[key]) != key_type: - raise MissingDataException(f"Entry {entry['id']} has key {key} of type {type(entry[key])} when it should be {key_type}.") - - # Check that the 4 required metadata keys are non-empty - for key in required_metadata_keys: if not entry[key]: raise MissingDataException(f"Entry {entry['id']} has no {key}.") + + if not isinstance(entry[key], key_type): + raise MissingDataException(f"Entry {entry['id']} has key {key} of type {type(entry[key])} when it should be {key_type}.") + + def iterable_data(self): + if not self.custom_sources: + return tqdm(load_dataset('StampyAI/alignment-research-dataset', 'all', split='train', streaming=True)) + + data = (entry for source in self.custom_sources for entry in load_dataset('StampyAI/alignment-research-dataset', source, split='train', streaming=True)) + return tqdm(data) def get_alignment_texts(self): text_splitter = TokenSplitter(self.min_tokens_per_block, self.max_tokens_per_block) # Load the dataset. streaming allows you to load one entry at a time, - # so entries can be processed before the entire dataset has been saved. - # iterable_data = load_dataset('StampyAI/alignment-research-dataset', 'aisafety.info', split='train', streaming=True) + iterable_data = self.iterable_data(self) - iterable_data = load_dataset('StampyAI/alignment-research-dataset', 'all', split='train', streaming=True) - - for entry in tqdm(iterable_data): + for entry in iterable_data: """Checks""" # if we specified custom sources, only include entries from those sources @@ -98,7 +100,14 @@ class ChunkedARD: continue # raise error if the entry does not contain the required metadata - self.contains_required_metadata(entry) + try: + self._validate_required_metadata(entry) + except MissingDataException as mde: + #logging.error(str(mde)) # Log the error message + print(str(mde)) + except Exception as e: + raise e + #if the text is too short, ignore this text if len(entry['text']) < 500: @@ -126,15 +135,14 @@ class ChunkedARD: # We use the text_splitter to get the chunks from the entry, # and we add a metadata element for each new chunk we add to the dataset. chunks = text_splitter.split(text, signature) - num_chunks = len(chunks) - for i in range(num_chunks): + for i, chunk in enumerate(chunks): self.metadata.append({ 'id': f"{entry_id}_{str(i+1).zfill(6)}", 'entry_id': entry_id, 'source': source, 'title': title, - 'text': chunks[i], + 'text': chunk, 'url': url, 'date_published': date_published, 'authors': authors @@ -218,6 +226,6 @@ def get_authors_str(authors_lst: List[str]) -> str: def standardize_date(date_string, default_date='n/a'): try: dt = parse(date_string) - return dt.strftime('%Y-%m-%d') + return dt.date().isoformat() except (ParserError, ValueError): return default_date \ No newline at end of file From 27a56b601e27e1ebc022bb9a45b16a84228a93c0 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:22:18 -0400 Subject: [PATCH 10/49] Updated gitignore --- .gitignore | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.gitignore b/.gitignore index 3687207..c98677c 100644 --- a/.gitignore +++ b/.gitignore @@ -147,3 +147,7 @@ api/dataset_big.pkl api/dataset_300.pkl api/.env.backup + +src/dataset_tests.ipynb +src/dataset/data/* +src/dataset/logs/* \ No newline at end of file From 72c28ad1334118bcc15c269b88ddc24b7adac555 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:22:33 -0400 Subject: [PATCH 11/49] added example .env in src/ --- src/.env.example | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 src/.env.example diff --git a/src/.env.example b/src/.env.example new file mode 100644 index 0000000..4e40cd4 --- /dev/null +++ b/src/.env.example @@ -0,0 +1,3 @@ +OPENAI_API_KEY="sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" +PINECONE_API_KEY="" # leave blank to use our online API instead +LOGGING_URL="" # leave blank if you're not testing logging specifically From 1514e9ec278715081b009f34827623e7b1c75119 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:22:56 -0400 Subject: [PATCH 12/49] Created an sqlite database handler --- src/dataset/database_handler.py | 104 ++++++++++++++++++++++++++++++++ 1 file changed, 104 insertions(+) create mode 100644 src/dataset/database_handler.py diff --git a/src/dataset/database_handler.py b/src/dataset/database_handler.py new file mode 100644 index 0000000..c766339 --- /dev/null +++ b/src/dataset/database_handler.py @@ -0,0 +1,104 @@ +import os +import sqlite3 +from typing import List, Dict, Any + +class DatabaseHandler: + def __init__( + self, + db_name: str = "data\\alignment_database.db", + ): + # Get the directory of this script + script_dir = os.path.dirname(os.path.realpath(__file__)) + + # Combine the script directory with the relative database path + self.db_name = os.path.join(script_dir, db_name) + + self.create_tables() + + def create_tables(self): + with sqlite3.connect(self.db_name) as conn: + cursor = conn.cursor() + try: + # Create entry table + query = """ + CREATE TABLE IF NOT EXISTS entry_database ( + id TEXT PRIMARY KEY, + source TEXT, + title TEXT, + text TEXT, + url TEXT, + date_published TEXT, + authors TEXT + ) + """ + cursor.execute(query) + + # Create chunk table + query = """ + CREATE TABLE IF NOT EXISTS chunk_database ( + id TEXT PRIMARY KEY, + text TEXT, + entry_id TEXT, + FOREIGN KEY (entry_id) REFERENCES entry_database(id) + ) + """ + cursor.execute(query) + + except sqlite3.Error as e: + print(f"The error '{e}' occurred.") + + def upsert_entry(self, entry: Dict[str, Any]) -> bool: + with sqlite3.connect(self.db_name) as conn: + cursor = conn.cursor() + try: + # Fetch existing data + cursor.execute("SELECT * FROM entry_database WHERE id=?", (entry['id'],)) + existing_entry = cursor.fetchone() + + new_entry = ( + entry['id'], + entry['source'], + entry['title'], + entry['text'], + entry['url'], + entry['date_published'], + ', '.join(entry['authors']) + ) + + if existing_entry != new_entry: + query = """ + INSERT OR REPLACE INTO entry_database + (id, source, title, text, url, date_published, authors) + VALUES (?, ?, ?, ?, ?, ?, ?) + """ + cursor.execute(query, new_entry) + return True + else: + return False + + except sqlite3.Error as e: + print(f"The error '{e}' occurred.") + return False + + finally: + conn.commit() + + def upsert_chunks(self, entry_id: str, chunks: List[str]) -> bool: + with sqlite3.connect(self.db_name) as conn: + cursor = conn.cursor() + try: + # Delete existing chunks + cursor.execute("DELETE FROM chunk_database WHERE entry_id=?", (entry_id,)) + + # Insert new chunks + for i, chunk in enumerate(chunks): + chunk_id = f"{entry_id}_{str(i).zfill(6)}" + cursor.execute("INSERT INTO chunk_database (id, text, entry_id) VALUES (?, ?, ?)", (chunk_id, chunk, entry_id)) + return True + + except sqlite3.Error as e: + print(f"The error '{e}' occurred.") + return False + + finally: + conn.commit() \ No newline at end of file From c76ad2df2807fd56a582efbb7271ad53c6ba00d0 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:23:10 -0400 Subject: [PATCH 13/49] Added a Pinecone database handler --- src/dataset/pinecone_handler.py | 79 +++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 src/dataset/pinecone_handler.py diff --git a/src/dataset/pinecone_handler.py b/src/dataset/pinecone_handler.py new file mode 100644 index 0000000..786567f --- /dev/null +++ b/src/dataset/pinecone_handler.py @@ -0,0 +1,79 @@ +# dataset/pinecone_handler.py + +import pinecone +import os +from typing import List + + +class PineconeHandler: + def __init__( + self, + index_name: str, + dimensions: int = 1536, + metric: str = "cosine", + location: str = "us-central1-gcp" + ): + self.index_name = index_name + self.dimensions = dimensions + self.metric = metric + + PINECONE_API_KEY = os.getenv("PINECONE_API_KEY") + assert PINECONE_API_KEY, "PINECONE_API_KEY environment variable not set." + + pinecone.init( + api_key = PINECONE_API_KEY, + environment = location, + ) + + self.index = pinecone.Index(index_name=self.index_name) + index_stats_response = self.index.describe_index_stats() + + print(f"Index info:\n\t{index_stats_response}\n\n") + + def insert_entry(self, entry, chunks, embeddings, upsert_size=100): + assert len(chunks) == len(embeddings), f"len(chunks) != len(embeddings) for {entry['title']} of {entry['source']}" + + chunk_len = len(chunks) + + vectors = [ + { + 'id': f"{entry['id']}_{str(i).zfill(6)}", + 'values': embeddings[i], + 'metadata': { + 'entry_id': entry['id'], + 'source': entry['source'], + 'title': entry['title'], + 'authors': entry['authors'] + } + } for i in range(chunk_len) + ] + + self.index.upsert( + vectors=vectors, + batch_size=upsert_size + ) + # print(f"Successfully inserted {chunk_len} chunks from {entry['source']} article \'{entry['title']}\'.") + + def delete_entry(self, id): + self.index.delete( + filter={"entry_id": {"$eq": id}} + ) + # print(f"Successfully deleted elements from id {id}.") + + def info(self): + info = pinecone.describe_index(self.index_name) + return info + + def create_index(self): + pinecone.create_index( + name=self.index_name, + dimension=self.dimensions, + metric=self.metric, + metadata_config = { + "indexed": ["title", "author", "date", "url", "source"] + } + ) + + def delete_index(self): + if self.index_name in pinecone.list_indexes(): + pinecone.delete_index(self.index_name) \ No newline at end of file From 1d8c025d24af7dc72997e070b3c043612f2146d9 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:26:45 -0400 Subject: [PATCH 14/49] Created update_dataset.py It contains a class, ARDUpdater, that deals with upserting entries to both sql and pinecone databases by checking diffs. It deals with splitting chunks and embedding them before upserting to pinecone. --- src/dataset/update_dataset.py | 161 ++++++++++++++++++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 src/dataset/update_dataset.py diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py new file mode 100644 index 0000000..efa39cb --- /dev/null +++ b/src/dataset/update_dataset.py @@ -0,0 +1,161 @@ +# dataset/update_dataset.py + +from typing import List +import numpy as np +import logging +from tqdm.auto import tqdm +from datasets import load_dataset +import openai + +from .text_splitter import TokenSplitter +from .pinecone_handler import PineconeHandler +from .database_handler import DatabaseHandler + +class ARDUpdater: + def __init__( + self, + min_tokens_per_block: int = 300, # Minimum number of tokens per block. + max_tokens_per_block: int = 400, # Maximum number of tokens per block. + rate_limit_per_minute: int = 3_500, # Rate limit for the OpenAI API. + embedding_model="text-embedding-ada-002", + embedding_dims=1536, + index_name="stampy-chat-embeddings-test", + update_all=False + ): + self.rate_limit_per_minute = rate_limit_per_minute + self.delay_in_seconds = 60.0 / self.rate_limit_per_minute + + self.embedding_model = embedding_model + self.embedding_dims = embedding_dims + + self.token_splitter = TokenSplitter( + min_tokens=min_tokens_per_block, + max_tokens=max_tokens_per_block + ) + self.db = DatabaseHandler() + self.pinecone_handler = PineconeHandler(index_name=index_name) + + self.update_all = update_all + + ### initialization code ### + + self.logger = logging.getLogger(__name__) + self.logger.setLevel(logging.INFO) + + # File handler + file_handler = logging.FileHandler(r'src/dataset/logs/ard_updater.log') + file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')) + + # Console handler + console_handler = logging.StreamHandler() + console_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')) + + self.logger.addHandler(file_handler) + self.logger.addHandler(console_handler) + + self.logger.info("ARDUpdater initialized.") + + + def update(self, custom_sources: List[str] = ['all']): + for source in custom_sources: + self.update_source(source) + + def update_source(self, source: str): + self.logger.info(f"Updating {source} entries...") + + iterable_data = load_dataset('StampyAI/alignment-research-dataset', source, split='train', streaming=True) + + for entry in tqdm(iterable_data): + try: + entry = self.process_entry(entry) + if entry is None: + continue + + # If the upsertion produces a change to the sql database, update the pinecone db accordingly + if self.update_all or self.db.upsert_entry(entry): + self.pinecone_handler.delete_entry(entry['id']) + + signature = f"Title: {entry['title']}, Authors: {get_authors_str(entry['authors'])}" + + chunks = self.token_splitter.split(entry['text'], signature) + + self.db.upsert_chunks(entry['id'], chunks) + + embeddings = self.get_embeddings(chunks) + + self.pinecone_handler.insert_entry(entry, chunks, embeddings) + + self.logger.info(f"Successfully modified entry {entry['id']}.") + + except Exception as e: + self.logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True) + return + + self.logger.info(f"Successfully updated {source} entries.") + + def process_entry(self, entry): + try: + self.validate_entry(entry) + + return { + 'id': entry['id'], + 'source': entry['source'], + 'title': entry['title'], + 'text': entry['text'], + 'url': entry['url'], + 'date_published': entry['date_published'], + 'authors': entry['authors'], + # summary is ignored for now + } + except ValueError as e: + self.logger.error(f"Entry validation failed: {str(e)}", exc_info=True) + return None + + def validate_entry(self, entry): + metadata_types = { + 'id': str, + 'source': str, + 'title': str, + 'url': str, + 'date_published': str, + 'authors': list, # unsure as of yet if this is a string or a list. TODO it is subject to change +# 'summary': list # see previous comment + } + + for metadata_type, metadata_type_type in metadata_types.items(): + if metadata_type not in entry: + raise ValueError(f"Entry is missing required metadata '{metadata_type}'.") + if not isinstance(entry[metadata_type], metadata_type_type): + raise ValueError(f"Entry metadata '{metadata_type}' is not of type '{metadata_type_type}'.") + + # if len(entry['text']) < 500: + # raise ValueError(f"Entry text is too short (< 500 tokens).") + + def get_embeddings(self, chunks): + embeddings = np.zeros((len(chunks), self.embedding_dims)) + + openai_output = openai.Embedding.create( + model=self.embedding_model, + input=chunks + )['data'] + + for i, embedding in enumerate(openai_output): + embeddings[i] = embedding['embedding'] + + return embeddings + + def show_stats(self): #TODO + # Show index + # Show database + pass + + +##### Helper functions ##### + +def get_authors_str(authors_lst: List[str]) -> str: + if authors_lst == []: return 'n/a' + if len(authors_lst) == 1: return authors_lst[0] + else: + authors_lst = authors_lst[:3] + authors_str = ", ".join(authors_lst[:-1]) + " and " + authors_lst[-1] + return authors_str \ No newline at end of file From 1c96ca4112894fc935a64771a09601aa195dd0fc Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:27:02 -0400 Subject: [PATCH 15/49] Updated main.py to call update_dataset.py --- src/main.py | 69 ++++++++++++++++------------------------------------- 1 file changed, 21 insertions(+), 48 deletions(-) diff --git a/src/main.py b/src/main.py index 87a68e6..fb932a6 100644 --- a/src/main.py +++ b/src/main.py @@ -1,58 +1,31 @@ +# main.py + +import os import openai -""" -import config -from assistant.semantic_search import AlignmentSearch -from dataset.create_dataset import Dataset -openai.api_key = config.OPENAI_API_KEY +from dataset.update_dataset import ARDUpdater -from settings import PATH_TO_RAW_DATA, PATH_TO_DATASET, EMBEDDING_MODEL, LEN_EMBEDDINGS -""" -from tenacity import ( - retry, - stop_after_attempt, - wait_random_exponential, -) +if os.path.exists('src/.env'): + from dotenv import load_dotenv + load_dotenv() +else: + raise Exception("'src/.env' not found. Rename the 'src/.env.example' file and fill in values.") -import numpy as np - -import sys -import pickle -from pathlib import Path -import random - -src_path = Path(__file__).resolve().parent -if str(src_path) not in sys.path: - sys.path.append(str(src_path)) - -from dataset import create_dataset -#from assistant import semantic_search -from settings import EMBEDDING_MODEL, PATH_TO_DATASET_DICT_PKL - -import numpy as np -import matplotlib.pyplot as plt +OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') +openai.api_key = OPENAI_API_KEY -def load_rawdata_into_pkl(): - dataset = create_dataset.ChunkedARD( - min_tokens_per_block=200, max_tokens_per_block=300 +def update_database_and_pinecone(): + updater = ARDUpdater( + min_tokens_per_block=200, + max_tokens_per_block=300, + index_name="stampy-chat-embeddings-test", + update_all=False ) - dataset.get_alignment_texts() - dataset.get_embeddings() - dataset.save_data() - dataset.show_stats() + updater.update(['gwern_blog', 'aisafety.info']) + updater.show_stats() -def load_pkl_and_display_stuff(): - with open(PATH_TO_DATASET_DICT_PKL, 'rb') as f: - dataset_dict = pickle.load(f) - print(dataset_dict) - -@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(4)) -def get_embedding(text: str) -> np.ndarray: - result = openai.Embedding.create(model=EMBEDDING_MODEL, input=text) - return np.array(result["data"][0]["embedding"]) if __name__ == "__main__": - load_rawdata_into_pkl() - #load_pkl_and_display_stuff() - \ No newline at end of file + print("\n"*10) + update_database_and_pinecone() \ No newline at end of file From 5ba1e0c22d0319feede4b7faaa108302f91b5e01 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:27:14 -0400 Subject: [PATCH 16/49] Removed unused files --- src/dataset/create_dataset.py | 231 ---------------------------------- src/dataset/settings.py | 10 -- src/settings.py | 10 -- 3 files changed, 251 deletions(-) delete mode 100644 src/dataset/create_dataset.py delete mode 100644 src/dataset/settings.py delete mode 100644 src/settings.py diff --git a/src/dataset/create_dataset.py b/src/dataset/create_dataset.py deleted file mode 100644 index ff3447a..0000000 --- a/src/dataset/create_dataset.py +++ /dev/null @@ -1,231 +0,0 @@ -import numpy as np -from typing import List, Dict, Tuple, DefaultDict, Any -from collections import defaultdict -import time -import random -import pickle -import os -import concurrent.futures -from pathlib import Path -from tqdm.auto import tqdm -from dateutil.parser import parse, ParserError -import openai -from datasets import load_dataset - - -try: - import config - openai.api_key = config.OPENAI_API_KEY -except ImportError: - openai.api_key = os.environ.get('OPENAI_API_KEY') - - -from .settings import PATH_TO_DATASET_DICT_PKL, EMBEDDING_MODEL, LEN_EMBEDDINGS - -from .text_splitter import TokenSplitter, split_into_sentences - - -class MissingDataException(Exception): - pass - - -class ChunkedARD: - def __init__(self, - min_tokens_per_block: int = 300, # Minimum number of tokens per block. - max_tokens_per_block: int = 400, # Maximum number of tokens per block. - custom_sources: List[str] = None, # List of sources to include, like "alignmentforum", "lesswrong", "arxiv",etc. - rate_limit_per_minute: int = 3_500, # Rate limit for the OpenAI API. - ): - self.min_tokens_per_block = min_tokens_per_block # for the text splitter - self.max_tokens_per_block = max_tokens_per_block # for the text splitter - - self.custom_sources = custom_sources - self.rate_limit_per_minute = rate_limit_per_minute - self.delay_in_seconds = 60.0 / self.rate_limit_per_minute - - self.metadata: List[Dict[str, Any]] = [] # List of dicts, each containing: id, entry_id, source, title, text, url, date_published, authors. - - self.entries_per_source_count: DefaultDict[str, int] = defaultdict(int) # Number of entries per source. E.g.: {'source1': 10, 'source2': 20, 'total': 30} - self.total_counts = { - 'chars': 0, - 'words': 0, - 'sentences': 0, - 'chunks': 0, - 'entries': 0, - } - - if self.custom_sources is not None: - for source in self.custom_sources: - self.entries_per_source_count[source] = 0 - - @staticmethod - def _validate_required_metadata(entry: Dict[str, Any]): - metadata_types = { - 'id': str, - 'source': str, - 'title': str, - 'url': str, - 'date_published': str, - 'authors': list, - 'summary': list - } - - # Check that the 8 primary metadata keys all have the correct type and the key exists - for key, key_type in metadata_types.items(): - if not entry[key]: - raise MissingDataException(f"Entry {entry['id']} has no {key}.") - - if not isinstance(entry[key], key_type): - raise MissingDataException(f"Entry {entry['id']} has key {key} of type {type(entry[key])} when it should be {key_type}.") - - def iterable_data(self): - if not self.custom_sources: - return tqdm(load_dataset('StampyAI/alignment-research-dataset', 'all', split='train', streaming=True)) - - data = (entry for source in self.custom_sources for entry in load_dataset('StampyAI/alignment-research-dataset', source, split='train', streaming=True)) - return tqdm(data) - - - def get_alignment_texts(self): - text_splitter = TokenSplitter(self.min_tokens_per_block, self.max_tokens_per_block) - - # Load the dataset. streaming allows you to load one entry at a time, - iterable_data = self.iterable_data(self) - - for entry in iterable_data: - """Checks""" - - # if we specified custom sources, only include entries from those sources - if (self.custom_sources is not None) and (entry['source'] not in self.custom_sources): - continue - - # raise error if the entry does not contain the required metadata - try: - self._validate_required_metadata(entry) - except MissingDataException as mde: - #logging.error(str(mde)) # Log the error message - print(str(mde)) - except Exception as e: - raise e - - - #if the text is too short, ignore this text - if len(entry['text']) < 500: - continue - - - """Checks are done, so we construct the metadata.""" - - # Get id, source, title, text, url, date_published, authors, and summary - entry_id: str = entry['id'] - source: str = entry['source'] - title: str = entry['title'] - text: str = entry['text'] - url: str = entry['url'] - date_published: str = entry['date_published'] - authors: list = entry['authors'] - # summary is ignored for now, since most sources lack one. TODO: add summary. see self.metadata code as well. - - # Get signature - if authors: - signature = f"Title: {title}, Authors: {get_authors_str(authors)}" - else: - signature = f"Title: {title}" - - # We use the text_splitter to get the chunks from the entry, - # and we add a metadata element for each new chunk we add to the dataset. - chunks = text_splitter.split(text, signature) - - for i, chunk in enumerate(chunks): - self.metadata.append({ - 'id': f"{entry_id}_{str(i+1).zfill(6)}", - 'entry_id': entry_id, - 'source': source, - 'title': title, - 'text': chunk, - 'url': url, - 'date_published': date_published, - 'authors': authors - }) - - # Update counts - self.entries_per_source_count[entry['source']] += 1 - self.total_counts['entries'] += 1 - self.total_counts['chars'] += len(text) - self.total_counts['words'] += len(text.split()) - self.total_counts['sentences'] += len(split_into_sentences(text)) - self.total_counts['chunks'] += len(chunks) - - def show_stats(self): - print(f'Number of entries by source: {self.entries_per_source_count}') - print(f'Total entries count: {self.total_counts["entries"]}') - print(f'Total character count: {self.total_counts["chars"]}') - print(f'Total word count: {self.total_counts["words"]}') - print(f'Total sentence count: {self.total_counts["sentences"]}') - print(f'Total chunk count: {self.total_counts["chunks"]}') - - def get_embeddings(self): - def get_embeddings_at_index(texts: str, batch_idx: int, batch_size: int = 200): # int, np.ndarray - embeddings = np.zeros((batch_size, 1536)) - openai_output = openai.Embedding.create( - model=EMBEDDING_MODEL, - input=texts - )['data'] - for i, embedding in enumerate(openai_output): - embeddings[i] = embedding['embedding'] - return batch_idx, embeddings - - batch_size = 500 - rate_limit = 3500 / 60 # Maximum embeddings per second - - start = time.time() - embedding_strings = [chunk_data['text'] for chunk_data in self.metadata] - self.embeddings = np.zeros((len(embedding_strings), LEN_EMBEDDINGS)) - - with concurrent.futures.ThreadPoolExecutor() as executor: - futures = [executor.submit( - get_embeddings_at_index, - embedding_strings[batch_idx:batch_idx+batch_size], - batch_idx, - len(embedding_strings[batch_idx:batch_idx+batch_size]) - ) for batch_idx in range(0, len(embedding_strings), batch_size)] - num_completed = 0 - for future in concurrent.futures.as_completed(futures): - batch_idx, embeddings = future.result() - num_completed += embeddings.shape[0] - self.embeddings[batch_idx:batch_idx+embeddings.shape[0]] = embeddings - - elapsed_time = time.time() - start - expected_time = num_completed / rate_limit - sleep_time = max(expected_time - elapsed_time, 0) - time.sleep(sleep_time) - - print(f"Completed {num_completed}/{len(embedding_strings)} embeddings in {elapsed_time:.2f} seconds.") - - def save_data(self, path: str = PATH_TO_DATASET_DICT_PKL): - # Save the data to a pickle file - print(f"Saving data to {path}...") - data = { - "metadata": self.metadata, - "embeddings": self.embeddings.astype(np.float32), - "total_counts": self.total_counts, - "entries_per_source_count": self.entries_per_source_count - } - with open(path, 'wb') as f: - pickle.dump(data, f) - - -def get_authors_str(authors_lst: List[str]) -> str: - if authors_lst == []: return 'n/a' - if len(authors_lst) == 1: return authors_lst[0] - else: - authors_lst = authors_lst[:3] - authors_str = ", ".join(authors_lst[:-1]) + " and " + authors_lst[-1] - return authors_str - -def standardize_date(date_string, default_date='n/a'): - try: - dt = parse(date_string) - return dt.date().isoformat() - except (ParserError, ValueError): - return default_date \ No newline at end of file diff --git a/src/dataset/settings.py b/src/dataset/settings.py deleted file mode 100644 index cf24c94..0000000 --- a/src/dataset/settings.py +++ /dev/null @@ -1,10 +0,0 @@ -from pathlib import Path - -EMBEDDING_MODEL = "text-embedding-ada-002" -COMPLETIONS_MODEL = "gpt-3.5-turbo" - -LEN_EMBEDDINGS = 1536 -MAX_LEN_PROMPT = 4095 # This may be 8191, unsure. - -current_file_path = Path(__file__).resolve() -PATH_TO_DATASET_DICT_PKL = str(current_file_path.parent / 'data' / 'dataset_dict.pkl') \ No newline at end of file diff --git a/src/settings.py b/src/settings.py deleted file mode 100644 index 94abeb5..0000000 --- a/src/settings.py +++ /dev/null @@ -1,10 +0,0 @@ -from pathlib import Path - -EMBEDDING_MODEL = "text-embedding-ada-002" -COMPLETIONS_MODEL = "gpt-3.5-turbo" - -LEN_EMBEDDINGS = 1536 -MAX_LEN_PROMPT = 4095 # This may be 8191, unsure. - -current_file_path = Path(__file__).resolve() -PATH_TO_DATASET_DICT_PKL = str(current_file_path.parent / 'dataset' / 'data' / 'dataset_dict.pkl') \ No newline at end of file From 61cb161ffc7a6a3b3605ac483c307f0c7d261184 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:48:56 -0400 Subject: [PATCH 17/49] Improved validate_entry --- src/dataset/update_dataset.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index efa39cb..cd21f1c 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -1,6 +1,6 @@ # dataset/update_dataset.py -from typing import List +from typing import Any, Dict, List import numpy as np import logging from tqdm.auto import tqdm @@ -89,7 +89,6 @@ class ARDUpdater: except Exception as e: self.logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True) - return self.logger.info(f"Successfully updated {source} entries.") @@ -104,22 +103,21 @@ class ARDUpdater: 'text': entry['text'], 'url': entry['url'], 'date_published': entry['date_published'], - 'authors': entry['authors'], - # summary is ignored for now + 'authors': entry['authors'] } except ValueError as e: self.logger.error(f"Entry validation failed: {str(e)}", exc_info=True) return None - def validate_entry(self, entry): + def validate_entry(self, entry: Dict[str, str | List[str]], len_lower_limit: int = 0): metadata_types = { 'id': str, 'source': str, 'title': str, + 'text': str, 'url': str, 'date_published': str, - 'authors': list, # unsure as of yet if this is a string or a list. TODO it is subject to change -# 'summary': list # see previous comment + 'authors': List[str] } for metadata_type, metadata_type_type in metadata_types.items(): @@ -128,8 +126,8 @@ class ARDUpdater: if not isinstance(entry[metadata_type], metadata_type_type): raise ValueError(f"Entry metadata '{metadata_type}' is not of type '{metadata_type_type}'.") - # if len(entry['text']) < 500: - # raise ValueError(f"Entry text is too short (< 500 tokens).") + if len(entry['text']) < len_lower_limit: + raise ValueError(f"Entry text is too short (< {len_lower_limit} tokens).") def get_embeddings(self, chunks): embeddings = np.zeros((len(chunks), self.embedding_dims)) From 22b912c8146658e21095a65d9664941c1abd0a66 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:55:57 -0400 Subject: [PATCH 18/49] Added convenient __str__ to the pinecone handler --- src/dataset/pinecone_handler.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/dataset/pinecone_handler.py b/src/dataset/pinecone_handler.py index 786567f..c62e07a 100644 --- a/src/dataset/pinecone_handler.py +++ b/src/dataset/pinecone_handler.py @@ -1,5 +1,6 @@ # dataset/pinecone_handler.py +import json import pinecone import os from typing import List @@ -26,9 +27,10 @@ class PineconeHandler: ) self.index = pinecone.Index(index_name=self.index_name) + + def __str__(self) -> str: index_stats_response = self.index.describe_index_stats() - - print(f"Index info:\n\t{index_stats_response}\n\n") + return f"{self.index_name}:\n{json.dumps(index_stats_response, indent=4)}" def insert_entry(self, entry, chunks, embeddings, upsert_size=100): assert len(chunks) == len(embeddings), f"len(chunks) != len(embeddings) for {entry['title']} of {entry['source']}" From d8c6665a5109552e61c1b19eb7746869a53000fa Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Fri, 23 Jun 2023 18:59:57 -0400 Subject: [PATCH 19/49] Deleted unnecessary lines and added replace_current_index functionality to create_index --- src/dataset/pinecone_handler.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/dataset/pinecone_handler.py b/src/dataset/pinecone_handler.py index c62e07a..0ddceff 100644 --- a/src/dataset/pinecone_handler.py +++ b/src/dataset/pinecone_handler.py @@ -3,7 +3,6 @@ import json import pinecone import os -from typing import List class PineconeHandler: @@ -54,19 +53,16 @@ class PineconeHandler: vectors=vectors, batch_size=upsert_size ) - # print(f"Successfully inserted {chunk_len} chunks from {entry['source']} article \'{entry['title']}\'.") def delete_entry(self, id): self.index.delete( filter={"entry_id": {"$eq": id}} ) - # print(f"Successfully deleted elements from id {id}.") - def info(self): - info = pinecone.describe_index(self.index_name) - return info - - def create_index(self): + def create_index(self, replace_current_index: bool = False): + if replace_current_index: + self.delete_index() + pinecone.create_index( name=self.index_name, dimension=self.dimensions, From b25777741d4c2410a4bacd483dba26cfbb8501c0 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Sat, 24 Jun 2023 00:21:12 -0400 Subject: [PATCH 20/49] Added main() for modularity --- src/main.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/main.py b/src/main.py index fb932a6..5035108 100644 --- a/src/main.py +++ b/src/main.py @@ -26,6 +26,9 @@ def update_database_and_pinecone(): updater.show_stats() +def main(): + update_database_and_pinecone() + + if __name__ == "__main__": - print("\n"*10) - update_database_and_pinecone() \ No newline at end of file + main() \ No newline at end of file From 049aa315e568f54fbe353fc15d7aa2f6c642fcff Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 01:45:38 -0400 Subject: [PATCH 21/49] Added ARD_LangChain_QA_Chat notebook to gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index c98677c..4112d28 100644 --- a/.gitignore +++ b/.gitignore @@ -149,5 +149,6 @@ api/dataset_300.pkl api/.env.backup src/dataset_tests.ipynb +src/ARD_LangChain_QA_Chat.ipynb src/dataset/data/* src/dataset/logs/* \ No newline at end of file From 4b64f3bae0af69ae61fb9f743cecce6d50712084 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 01:45:59 -0400 Subject: [PATCH 22/49] Updated .env.example file --- src/.env.example | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/.env.example b/src/.env.example index 4e40cd4..eda5459 100644 --- a/src/.env.example +++ b/src/.env.example @@ -1,3 +1,3 @@ OPENAI_API_KEY="sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" -PINECONE_API_KEY="" # leave blank to use our online API instead -LOGGING_URL="" # leave blank if you're not testing logging specifically +PINECONE_API_KEY="xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx" +PINECONE_ENVIRONMENT="xx-xxxxx-gcp" \ No newline at end of file From 0c97ea61c50687c8cf167309eee2d0efa49ec309 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 01:46:19 -0400 Subject: [PATCH 23/49] Added settings file --- src/dataset/settings.py | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 src/dataset/settings.py diff --git a/src/dataset/settings.py b/src/dataset/settings.py new file mode 100644 index 0000000..892408a --- /dev/null +++ b/src/dataset/settings.py @@ -0,0 +1,6 @@ + +### PINECONE ### +PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" +PINECONE_VALUES_DIMS = 1536 +PINECONE_METRIC = "cosine" +PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] \ No newline at end of file From cda0ce44e9fa901c135ab999382710f0ffa6f909 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 01:48:00 -0400 Subject: [PATCH 24/49] Simplified update_source method, fixed bugs --- src/dataset/update_dataset.py | 61 +++++++++++++++-------------------- 1 file changed, 26 insertions(+), 35 deletions(-) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index cd21f1c..a60dded 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -1,11 +1,12 @@ # dataset/update_dataset.py -from typing import Any, Dict, List +import time +from typing import Dict, List import numpy as np import logging from tqdm.auto import tqdm -from datasets import load_dataset import openai +from datasets import load_dataset from .text_splitter import TokenSplitter from .pinecone_handler import PineconeHandler @@ -14,13 +15,12 @@ from .database_handler import DatabaseHandler class ARDUpdater: def __init__( self, - min_tokens_per_block: int = 300, # Minimum number of tokens per block. + min_tokens_per_block: int = 200, # Minimum number of tokens per block. max_tokens_per_block: int = 400, # Maximum number of tokens per block. rate_limit_per_minute: int = 3_500, # Rate limit for the OpenAI API. embedding_model="text-embedding-ada-002", embedding_dims=1536, index_name="stampy-chat-embeddings-test", - update_all=False ): self.rate_limit_per_minute = rate_limit_per_minute self.delay_in_seconds = 60.0 / self.rate_limit_per_minute @@ -35,8 +35,6 @@ class ARDUpdater: self.db = DatabaseHandler() self.pinecone_handler = PineconeHandler(index_name=index_name) - self.update_all = update_all - ### initialization code ### self.logger = logging.getLogger(__name__) @@ -64,35 +62,30 @@ class ARDUpdater: self.logger.info(f"Updating {source} entries...") iterable_data = load_dataset('StampyAI/alignment-research-dataset', source, split='train', streaming=True) - + iterable_data = iterable_data.map(self.preprocess) + iterable_data = iterable_data.filter(lambda entry: entry is not None) + iterable_data = iterable_data.filter(lambda entry: self.db.upsert_entry(entry)) + for entry in tqdm(iterable_data): + t_entry_start = time.time() try: - entry = self.process_entry(entry) - if entry is None: - continue + self.pinecone_handler.delete_entry(entry['id']) + + signature = f"Title: {entry['title']}, Authors: {get_authors_str(entry['authors'])}" + chunks = self.token_splitter.split(entry['text'], signature) + embeddings = self.get_embeddings(chunks) - # If the upsertion produces a change to the sql database, update the pinecone db accordingly - if self.update_all or self.db.upsert_entry(entry): - self.pinecone_handler.delete_entry(entry['id']) - - signature = f"Title: {entry['title']}, Authors: {get_authors_str(entry['authors'])}" - - chunks = self.token_splitter.split(entry['text'], signature) - - self.db.upsert_chunks(entry['id'], chunks) - - embeddings = self.get_embeddings(chunks) - - self.pinecone_handler.insert_entry(entry, chunks, embeddings) - - self.logger.info(f"Successfully modified entry {entry['id']}.") - + self.db.upsert_chunks(entry['id'], chunks) + self.pinecone_handler.insert_entry(entry, chunks, embeddings) except Exception as e: self.logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True) + + t_entry_end = time.time() + self.logger.info(f"Time for processing one entry: {t_entry_end - t_entry_start} seconds") self.logger.info(f"Successfully updated {source} entries.") - - def process_entry(self, entry): + + def preprocess(self, entry): try: self.validate_entry(entry) @@ -109,7 +102,7 @@ class ARDUpdater: self.logger.error(f"Entry validation failed: {str(e)}", exc_info=True) return None - def validate_entry(self, entry: Dict[str, str | List[str]], len_lower_limit: int = 0): + def validate_entry(self, entry: Dict[str, str | list], len_lower_limit: int = 0): metadata_types = { 'id': str, 'source': str, @@ -117,15 +110,13 @@ class ARDUpdater: 'text': str, 'url': str, 'date_published': str, - 'authors': List[str] + 'authors': list } for metadata_type, metadata_type_type in metadata_types.items(): - if metadata_type not in entry: - raise ValueError(f"Entry is missing required metadata '{metadata_type}'.") - if not isinstance(entry[metadata_type], metadata_type_type): - raise ValueError(f"Entry metadata '{metadata_type}' is not of type '{metadata_type_type}'.") - + if not isinstance(entry.get(metadata_type), metadata_type_type): + raise ValueError(f"Entry metadata '{metadata_type}' is not of type '{metadata_type_type}' or is missing.") + if len(entry['text']) < len_lower_limit: raise ValueError(f"Entry text is too short (< {len_lower_limit} tokens).") From d9e4cb67e09998859b6d974ca339fca73d3177ad Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 01:50:52 -0400 Subject: [PATCH 25/49] Utilized settings file for pinecone handler --- src/dataset/pinecone_handler.py | 58 ++++++++++++++++----------------- src/dataset/update_dataset.py | 3 +- src/main.py | 5 +-- 3 files changed, 30 insertions(+), 36 deletions(-) diff --git a/src/dataset/pinecone_handler.py b/src/dataset/pinecone_handler.py index 0ddceff..a908c27 100644 --- a/src/dataset/pinecone_handler.py +++ b/src/dataset/pinecone_handler.py @@ -3,28 +3,31 @@ import json import pinecone import os +from settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES + class PineconeHandler: def __init__( self, index_name: str, - dimensions: int = 1536, - metric: str = "cosine", - location: str = "us-central1-gcp" + create_index: bool = False, ): self.index_name = index_name - self.dimensions = dimensions - self.metric = metric PINECONE_API_KEY = os.getenv("PINECONE_API_KEY") + PINECONE_ENVIRONMENT = os.getenv("PINECONE_ENVIRONMENT") assert PINECONE_API_KEY, "PINECONE_API_KEY environment variable not set." + assert PINECONE_ENVIRONMENT, "PINECONE_LOCATION environment variable not set." pinecone.init( api_key = PINECONE_API_KEY, - environment = location, + environment = PINECONE_ENVIRONMENT, ) + if create_index: + self.create_index() + self.index = pinecone.Index(index_name=self.index_name) def __str__(self) -> str: @@ -32,25 +35,22 @@ class PineconeHandler: return f"{self.index_name}:\n{json.dumps(index_stats_response, indent=4)}" def insert_entry(self, entry, chunks, embeddings, upsert_size=100): - assert len(chunks) == len(embeddings), f"len(chunks) != len(embeddings) for {entry['title']} of {entry['source']}" - - chunk_len = len(chunks) - - vectors = [ - { - 'id': f"{entry['id']}_{str(i).zfill(6)}", - 'values': embeddings[i], - 'metadata': { - 'entry_id': entry['id'], - 'source': entry['source'], - 'title': entry['title'], - 'authors': entry['authors'] - } - } for i in range(chunk_len) - ] - self.index.upsert( - vectors=vectors, + vectors=list( + zip( + [f"{entry['id']}_{str(i).zfill(6)}" for i in range(len(chunks))], + embeddings.tolist(), + [ + { + 'entry_id': entry['id'], + 'source': entry['source'], + 'title': entry['title'], + 'authors': entry['authors'], + 'text': chunk, + } for chunk in chunks + ] + ) + ), batch_size=upsert_size ) @@ -59,17 +59,15 @@ class PineconeHandler: filter={"entry_id": {"$eq": id}} ) - def create_index(self, replace_current_index: bool = False): + def create_index(self, replace_current_index: bool = True): if replace_current_index: self.delete_index() pinecone.create_index( name=self.index_name, - dimension=self.dimensions, - metric=self.metric, - metadata_config = { - "indexed": ["title", "author", "date", "url", "source"] - } + dimension=PINECONE_VALUES_DIMS, + metric=PINECONE_METRIC, + metadata_config = {"indexed": PINECONE_METADATA_ENTRIES} ) def delete_index(self): diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index a60dded..0bb8f0f 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -20,7 +20,6 @@ class ARDUpdater: rate_limit_per_minute: int = 3_500, # Rate limit for the OpenAI API. embedding_model="text-embedding-ada-002", embedding_dims=1536, - index_name="stampy-chat-embeddings-test", ): self.rate_limit_per_minute = rate_limit_per_minute self.delay_in_seconds = 60.0 / self.rate_limit_per_minute @@ -33,7 +32,7 @@ class ARDUpdater: max_tokens=max_tokens_per_block ) self.db = DatabaseHandler() - self.pinecone_handler = PineconeHandler(index_name=index_name) + self.pinecone_handler = PineconeHandler() ### initialization code ### diff --git a/src/main.py b/src/main.py index 5035108..cac6170 100644 --- a/src/main.py +++ b/src/main.py @@ -19,11 +19,8 @@ def update_database_and_pinecone(): updater = ARDUpdater( min_tokens_per_block=200, max_tokens_per_block=300, - index_name="stampy-chat-embeddings-test", - update_all=False ) - updater.update(['gwern_blog', 'aisafety.info']) - updater.show_stats() + updater.update(['yudkowsky_blog']) def main(): From 822d958700a65a780b08659aff69ac81157f5c65 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 01:54:21 -0400 Subject: [PATCH 26/49] bug fix --- src/dataset/pinecone_handler.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/dataset/pinecone_handler.py b/src/dataset/pinecone_handler.py index a908c27..3b27650 100644 --- a/src/dataset/pinecone_handler.py +++ b/src/dataset/pinecone_handler.py @@ -3,17 +3,16 @@ import json import pinecone import os -from settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES +from .settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES class PineconeHandler: def __init__( self, - index_name: str, create_index: bool = False, ): - self.index_name = index_name + self.index_name = PINECONE_INDEX_NAME PINECONE_API_KEY = os.getenv("PINECONE_API_KEY") PINECONE_ENVIRONMENT = os.getenv("PINECONE_ENVIRONMENT") From 9314de03e54279c568d4afb23d50cc161ecb16e7 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 01:54:52 -0400 Subject: [PATCH 27/49] Removed time logging and show_stats --- src/dataset/update_dataset.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 0bb8f0f..745f031 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -1,6 +1,5 @@ # dataset/update_dataset.py -import time from typing import Dict, List import numpy as np import logging @@ -66,7 +65,6 @@ class ARDUpdater: iterable_data = iterable_data.filter(lambda entry: self.db.upsert_entry(entry)) for entry in tqdm(iterable_data): - t_entry_start = time.time() try: self.pinecone_handler.delete_entry(entry['id']) @@ -79,9 +77,6 @@ class ARDUpdater: except Exception as e: self.logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True) - t_entry_end = time.time() - self.logger.info(f"Time for processing one entry: {t_entry_end - t_entry_start} seconds") - self.logger.info(f"Successfully updated {source} entries.") def preprocess(self, entry): @@ -132,11 +127,6 @@ class ARDUpdater: return embeddings - def show_stats(self): #TODO - # Show index - # Show database - pass - ##### Helper functions ##### From f45e0c034fe389e21ef787aa149ebdd819d4e220 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 02:06:05 -0400 Subject: [PATCH 28/49] Removed logging file, added reset_dbs method --- src/dataset/database_handler.py | 7 +++++- src/dataset/update_dataset.py | 42 ++++++++++++--------------------- src/main.py | 2 +- 3 files changed, 22 insertions(+), 29 deletions(-) diff --git a/src/dataset/database_handler.py b/src/dataset/database_handler.py index c766339..c6e81a8 100644 --- a/src/dataset/database_handler.py +++ b/src/dataset/database_handler.py @@ -15,10 +15,15 @@ class DatabaseHandler: self.create_tables() - def create_tables(self): + def create_tables(self, reset: bool = False): with sqlite3.connect(self.db_name) as conn: cursor = conn.cursor() try: + if reset: + # Drop the tables if reset is True + cursor.execute("DROP TABLE IF EXISTS entry_database") + cursor.execute("DROP TABLE IF EXISTS chunk_database") + # Create entry table query = """ CREATE TABLE IF NOT EXISTS entry_database ( diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 745f031..a8b5c49 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -2,7 +2,6 @@ from typing import Dict, List import numpy as np -import logging from tqdm.auto import tqdm import openai from datasets import load_dataset @@ -11,6 +10,10 @@ from .text_splitter import TokenSplitter from .pinecone_handler import PineconeHandler from .database_handler import DatabaseHandler +import logging +logger = logging.getLogger(__name__) + + class ARDUpdater: def __init__( self, @@ -31,33 +34,14 @@ class ARDUpdater: max_tokens=max_tokens_per_block ) self.db = DatabaseHandler() - self.pinecone_handler = PineconeHandler() - - ### initialization code ### - - self.logger = logging.getLogger(__name__) - self.logger.setLevel(logging.INFO) - - # File handler - file_handler = logging.FileHandler(r'src/dataset/logs/ard_updater.log') - file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')) - - # Console handler - console_handler = logging.StreamHandler() - console_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')) - - self.logger.addHandler(file_handler) - self.logger.addHandler(console_handler) - - self.logger.info("ARDUpdater initialized.") - + self.pinecone_db = PineconeHandler() def update(self, custom_sources: List[str] = ['all']): for source in custom_sources: self.update_source(source) def update_source(self, source: str): - self.logger.info(f"Updating {source} entries...") + logger.info(f"Updating {source} entries...") iterable_data = load_dataset('StampyAI/alignment-research-dataset', source, split='train', streaming=True) iterable_data = iterable_data.map(self.preprocess) @@ -66,18 +50,18 @@ class ARDUpdater: for entry in tqdm(iterable_data): try: - self.pinecone_handler.delete_entry(entry['id']) + self.pinecone_db.delete_entry(entry['id']) signature = f"Title: {entry['title']}, Authors: {get_authors_str(entry['authors'])}" chunks = self.token_splitter.split(entry['text'], signature) embeddings = self.get_embeddings(chunks) self.db.upsert_chunks(entry['id'], chunks) - self.pinecone_handler.insert_entry(entry, chunks, embeddings) + self.pinecone_db.insert_entry(entry, chunks, embeddings) except Exception as e: - self.logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True) + logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True) - self.logger.info(f"Successfully updated {source} entries.") + logger.info(f"Successfully updated {source} entries.") def preprocess(self, entry): try: @@ -93,7 +77,7 @@ class ARDUpdater: 'authors': entry['authors'] } except ValueError as e: - self.logger.error(f"Entry validation failed: {str(e)}", exc_info=True) + logger.error(f"Entry validation failed: {str(e)}", exc_info=True) return None def validate_entry(self, entry: Dict[str, str | list], len_lower_limit: int = 0): @@ -127,6 +111,10 @@ class ARDUpdater: return embeddings + def reset_dbs(self): + self.db.create_tables(True) + self.pinecone_db.create_index(True) + ##### Helper functions ##### diff --git a/src/main.py b/src/main.py index cac6170..23a91e4 100644 --- a/src/main.py +++ b/src/main.py @@ -20,7 +20,7 @@ def update_database_and_pinecone(): min_tokens_per_block=200, max_tokens_per_block=300, ) - updater.update(['yudkowsky_blog']) + updater.update(['gwern_blog']) def main(): From 2438a2b01f8279f91bc6e64ba23ec11ac97ed8af Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 02:15:35 -0400 Subject: [PATCH 29/49] Added logger to sql and pinecone handlers --- src/dataset/database_handler.py | 10 +++++++--- src/dataset/pinecone_handler.py | 6 +++++- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/dataset/database_handler.py b/src/dataset/database_handler.py index c6e81a8..c2a891f 100644 --- a/src/dataset/database_handler.py +++ b/src/dataset/database_handler.py @@ -2,6 +2,10 @@ import os import sqlite3 from typing import List, Dict, Any +import logging +logger = logging.getLogger(__name__) + + class DatabaseHandler: def __init__( self, @@ -50,7 +54,7 @@ class DatabaseHandler: cursor.execute(query) except sqlite3.Error as e: - print(f"The error '{e}' occurred.") + logger.error(f"The error '{e}' occurred.") def upsert_entry(self, entry: Dict[str, Any]) -> bool: with sqlite3.connect(self.db_name) as conn: @@ -82,7 +86,7 @@ class DatabaseHandler: return False except sqlite3.Error as e: - print(f"The error '{e}' occurred.") + logger.error(f"The error '{e}' occurred.") return False finally: @@ -102,7 +106,7 @@ class DatabaseHandler: return True except sqlite3.Error as e: - print(f"The error '{e}' occurred.") + logger.error(f"The error '{e}' occurred.") return False finally: diff --git a/src/dataset/pinecone_handler.py b/src/dataset/pinecone_handler.py index 3b27650..bca5273 100644 --- a/src/dataset/pinecone_handler.py +++ b/src/dataset/pinecone_handler.py @@ -3,8 +3,11 @@ import json import pinecone import os -from .settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES +import logging +logger = logging.getLogger(__name__) + +from .settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES class PineconeHandler: @@ -71,4 +74,5 @@ class PineconeHandler: def delete_index(self): if self.index_name in pinecone.list_indexes(): + logger.info(f"Deleting index '{self.index_name}'.") pinecone.delete_index(self.index_name) \ No newline at end of file From 17ca839d54b26f28bb21195be741b94fcc5717e8 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 02:39:17 -0400 Subject: [PATCH 30/49] Removed logs from gitignore --- .gitignore | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 4112d28..5bfc436 100644 --- a/.gitignore +++ b/.gitignore @@ -150,5 +150,5 @@ api/.env.backup src/dataset_tests.ipynb src/ARD_LangChain_QA_Chat.ipynb -src/dataset/data/* -src/dataset/logs/* \ No newline at end of file + +src/dataset/data/* \ No newline at end of file From 473b5ed9459fad44a1a79698d3d05381d20fc11e Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 02:39:49 -0400 Subject: [PATCH 31/49] Renamed class and file --- src/dataset/{pinecone_handler.py => pinecone_db_handler.py} | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) rename src/dataset/{pinecone_handler.py => pinecone_db_handler.py} (99%) diff --git a/src/dataset/pinecone_handler.py b/src/dataset/pinecone_db_handler.py similarity index 99% rename from src/dataset/pinecone_handler.py rename to src/dataset/pinecone_db_handler.py index bca5273..6f53daa 100644 --- a/src/dataset/pinecone_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -4,13 +4,13 @@ import json import pinecone import os +from .settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES + import logging logger = logging.getLogger(__name__) -from .settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES - -class PineconeHandler: +class PineconeDB: def __init__( self, create_index: bool = False, From 20346d5fa493f493abb5f97eb4af25e1eccf021c Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 02:40:00 -0400 Subject: [PATCH 32/49] Renamed file --- src/dataset/{database_handler.py => sql_db_handler.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename src/dataset/{database_handler.py => sql_db_handler.py} (100%) diff --git a/src/dataset/database_handler.py b/src/dataset/sql_db_handler.py similarity index 100% rename from src/dataset/database_handler.py rename to src/dataset/sql_db_handler.py From 5900da48bb381030e54e5abfdd2b7fcfff7e4c7f Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 02:45:50 -0400 Subject: [PATCH 33/49] Renamed classes, added settings --- src/dataset/pinecone_db_handler.py | 2 +- src/dataset/settings.py | 19 ++++++++++++++++-- src/dataset/sql_db_handler.py | 16 +++++---------- src/dataset/update_dataset.py | 32 ++++++++++++------------------ src/main.py | 11 +++++----- 5 files changed, 41 insertions(+), 39 deletions(-) diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py index 6f53daa..a0db67d 100644 --- a/src/dataset/pinecone_db_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -10,7 +10,7 @@ import logging logger = logging.getLogger(__name__) -class PineconeDB: +class PineconeDBHandler: def __init__( self, create_index: bool = False, diff --git a/src/dataset/settings.py b/src/dataset/settings.py index 892408a..f8045a3 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -1,6 +1,21 @@ +from pathlib import Path + +### FILE PATHS ### +current_file_path = Path(__file__).resolve() +SQL_DB_PATH = str(current_file_path.parent / 'data' / 'ARD.db') + +### DATASET ### +ARD_DATASET_NAME = "StampyAI/alignment-research-dataset" + +### EMBEDDINGS ### +EMBEDDINGS_MODEL = "text-embedding-ada-002" +EMBEDDING_DIMS = 1536 ### PINECONE ### PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" -PINECONE_VALUES_DIMS = 1536 +PINECONE_VALUES_DIMS = EMBEDDING_DIMS PINECONE_METRIC = "cosine" -PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] \ No newline at end of file +PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] + +### MISC ### +CUSTOM_SOURCES = ['gwern_blog'] \ No newline at end of file diff --git a/src/dataset/sql_db_handler.py b/src/dataset/sql_db_handler.py index c2a891f..05afbf3 100644 --- a/src/dataset/sql_db_handler.py +++ b/src/dataset/sql_db_handler.py @@ -1,21 +1,15 @@ -import os import sqlite3 from typing import List, Dict, Any +from .settings import SQL_DB_PATH + import logging logger = logging.getLogger(__name__) -class DatabaseHandler: - def __init__( - self, - db_name: str = "data\\alignment_database.db", - ): - # Get the directory of this script - script_dir = os.path.dirname(os.path.realpath(__file__)) - - # Combine the script directory with the relative database path - self.db_name = os.path.join(script_dir, db_name) +class SQLDBHandler: + def __init__(self): + self.db_name = SQL_DB_PATH self.create_tables() diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index a8b5c49..d9d2e8d 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -7,8 +7,10 @@ import openai from datasets import load_dataset from .text_splitter import TokenSplitter -from .pinecone_handler import PineconeHandler -from .database_handler import DatabaseHandler +from .sql_db_handler import SQLDBHandler +from .pinecone_db_handler import PineconeDBHandler + +from .settings import EMBEDDINGS_MODEL, EMBEDDING_DIMS, ARD_DATASET_NAME import logging logger = logging.getLogger(__name__) @@ -20,21 +22,13 @@ class ARDUpdater: min_tokens_per_block: int = 200, # Minimum number of tokens per block. max_tokens_per_block: int = 400, # Maximum number of tokens per block. rate_limit_per_minute: int = 3_500, # Rate limit for the OpenAI API. - embedding_model="text-embedding-ada-002", - embedding_dims=1536, ): self.rate_limit_per_minute = rate_limit_per_minute self.delay_in_seconds = 60.0 / self.rate_limit_per_minute - self.embedding_model = embedding_model - self.embedding_dims = embedding_dims - - self.token_splitter = TokenSplitter( - min_tokens=min_tokens_per_block, - max_tokens=max_tokens_per_block - ) - self.db = DatabaseHandler() - self.pinecone_db = PineconeHandler() + self.token_splitter = TokenSplitter(min_tokens_per_block, max_tokens_per_block) + self.sql_db = SQLDBHandler() + self.pinecone_db = PineconeDBHandler() def update(self, custom_sources: List[str] = ['all']): for source in custom_sources: @@ -43,10 +37,10 @@ class ARDUpdater: def update_source(self, source: str): logger.info(f"Updating {source} entries...") - iterable_data = load_dataset('StampyAI/alignment-research-dataset', source, split='train', streaming=True) + iterable_data = load_dataset(ARD_DATASET_NAME, source, split='train', streaming=True) iterable_data = iterable_data.map(self.preprocess) iterable_data = iterable_data.filter(lambda entry: entry is not None) - iterable_data = iterable_data.filter(lambda entry: self.db.upsert_entry(entry)) + iterable_data = iterable_data.filter(lambda entry: self.sql_db.upsert_entry(entry)) for entry in tqdm(iterable_data): try: @@ -56,7 +50,7 @@ class ARDUpdater: chunks = self.token_splitter.split(entry['text'], signature) embeddings = self.get_embeddings(chunks) - self.db.upsert_chunks(entry['id'], chunks) + self.sql_db.upsert_chunks(entry['id'], chunks) self.pinecone_db.insert_entry(entry, chunks, embeddings) except Exception as e: logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True) @@ -99,10 +93,10 @@ class ARDUpdater: raise ValueError(f"Entry text is too short (< {len_lower_limit} tokens).") def get_embeddings(self, chunks): - embeddings = np.zeros((len(chunks), self.embedding_dims)) + embeddings = np.zeros((len(chunks), EMBEDDING_DIMS)) openai_output = openai.Embedding.create( - model=self.embedding_model, + model=EMBEDDINGS_MODEL, input=chunks )['data'] @@ -112,7 +106,7 @@ class ARDUpdater: return embeddings def reset_dbs(self): - self.db.create_tables(True) + self.sql_db.create_tables(True) self.pinecone_db.create_index(True) diff --git a/src/main.py b/src/main.py index 23a91e4..eb6f418 100644 --- a/src/main.py +++ b/src/main.py @@ -1,18 +1,16 @@ # main.py import os -import openai - -from dataset.update_dataset import ARDUpdater - if os.path.exists('src/.env'): from dotenv import load_dotenv load_dotenv() else: raise Exception("'src/.env' not found. Rename the 'src/.env.example' file and fill in values.") -OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') -openai.api_key = OPENAI_API_KEY +import openai +openai.api_key = os.environ.get('OPENAI_API_KEY') + +from dataset.update_dataset import ARDUpdater def update_database_and_pinecone(): @@ -20,6 +18,7 @@ def update_database_and_pinecone(): min_tokens_per_block=200, max_tokens_per_block=300, ) + updater.reset_dbs() updater.update(['gwern_blog']) From 0204810aa54286ab72b4c33da244099ed93d5700 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 03:09:51 -0400 Subject: [PATCH 34/49] Improving naming --- src/dataset/pinecone_db_handler.py | 4 ++-- src/dataset/settings.py | 7 ++++--- src/dataset/sql_db_handler.py | 2 +- src/dataset/update_dataset.py | 19 ++++++++----------- src/main.py | 9 +++++---- 5 files changed, 20 insertions(+), 21 deletions(-) diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py index a0db67d..e62e91d 100644 --- a/src/dataset/pinecone_db_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -10,7 +10,7 @@ import logging logger = logging.getLogger(__name__) -class PineconeDBHandler: +class PineconeDB: def __init__( self, create_index: bool = False, @@ -36,7 +36,7 @@ class PineconeDBHandler: index_stats_response = self.index.describe_index_stats() return f"{self.index_name}:\n{json.dumps(index_stats_response, indent=4)}" - def insert_entry(self, entry, chunks, embeddings, upsert_size=100): + def upsert_entry(self, entry, chunks, embeddings, upsert_size=100): self.index.upsert( vectors=list( zip( diff --git a/src/dataset/settings.py b/src/dataset/settings.py index f8045a3..ec04ef2 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -9,13 +9,14 @@ ARD_DATASET_NAME = "StampyAI/alignment-research-dataset" ### EMBEDDINGS ### EMBEDDINGS_MODEL = "text-embedding-ada-002" -EMBEDDING_DIMS = 1536 +EMBEDDINGS_DIMS = 1536 +EMBEDDINGS_RATE_LIMIT = 3500 ### PINECONE ### PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" -PINECONE_VALUES_DIMS = EMBEDDING_DIMS +PINECONE_VALUES_DIMS = EMBEDDINGS_DIMS PINECONE_METRIC = "cosine" PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] ### MISC ### -CUSTOM_SOURCES = ['gwern_blog'] \ No newline at end of file +RESET_DBS = False \ No newline at end of file diff --git a/src/dataset/sql_db_handler.py b/src/dataset/sql_db_handler.py index 05afbf3..9ba223d 100644 --- a/src/dataset/sql_db_handler.py +++ b/src/dataset/sql_db_handler.py @@ -7,7 +7,7 @@ import logging logger = logging.getLogger(__name__) -class SQLDBHandler: +class SQLDB: def __init__(self): self.db_name = SQL_DB_PATH diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index d9d2e8d..4c4dac5 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -7,10 +7,10 @@ import openai from datasets import load_dataset from .text_splitter import TokenSplitter -from .sql_db_handler import SQLDBHandler -from .pinecone_db_handler import PineconeDBHandler +from .sql_db_handler import SQLDB +from .pinecone_db_handler import PineconeDB -from .settings import EMBEDDINGS_MODEL, EMBEDDING_DIMS, ARD_DATASET_NAME +from .settings import EMBEDDINGS_MODEL, EMBEDDINGS_DIMS, EMBEDDINGS_RATE_LIMIT, ARD_DATASET_NAME import logging logger = logging.getLogger(__name__) @@ -21,14 +21,10 @@ class ARDUpdater: self, min_tokens_per_block: int = 200, # Minimum number of tokens per block. max_tokens_per_block: int = 400, # Maximum number of tokens per block. - rate_limit_per_minute: int = 3_500, # Rate limit for the OpenAI API. ): - self.rate_limit_per_minute = rate_limit_per_minute - self.delay_in_seconds = 60.0 / self.rate_limit_per_minute - self.token_splitter = TokenSplitter(min_tokens_per_block, max_tokens_per_block) - self.sql_db = SQLDBHandler() - self.pinecone_db = PineconeDBHandler() + self.sql_db = SQLDB() + self.pinecone_db = PineconeDB() def update(self, custom_sources: List[str] = ['all']): for source in custom_sources: @@ -51,7 +47,7 @@ class ARDUpdater: embeddings = self.get_embeddings(chunks) self.sql_db.upsert_chunks(entry['id'], chunks) - self.pinecone_db.insert_entry(entry, chunks, embeddings) + self.pinecone_db.upsert_entry(entry, chunks, embeddings) except Exception as e: logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True) @@ -93,7 +89,8 @@ class ARDUpdater: raise ValueError(f"Entry text is too short (< {len_lower_limit} tokens).") def get_embeddings(self, chunks): - embeddings = np.zeros((len(chunks), EMBEDDING_DIMS)) + embeddings = np.zeros((len(chunks), EMBEDDINGS_DIMS)) + rate_limit = EMBEDDINGS_RATE_LIMIT #TODO: use this rate_limit openai_output = openai.Embedding.create( model=EMBEDDINGS_MODEL, diff --git a/src/main.py b/src/main.py index eb6f418..770ac39 100644 --- a/src/main.py +++ b/src/main.py @@ -11,19 +11,20 @@ import openai openai.api_key = os.environ.get('OPENAI_API_KEY') from dataset.update_dataset import ARDUpdater +from dataset.settings import RESET_DBS - -def update_database_and_pinecone(): +def update_sql_and_pinecone_dbs(): updater = ARDUpdater( min_tokens_per_block=200, max_tokens_per_block=300, ) - updater.reset_dbs() + if RESET_DBS: + updater.reset_dbs() updater.update(['gwern_blog']) def main(): - update_database_and_pinecone() + update_sql_and_pinecone_dbs() if __name__ == "__main__": From abb506511e6fc99dfdd70a11eb3fe39839cbdddd Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 06:41:21 -0400 Subject: [PATCH 35/49] Chaining iterable_data calls --- src/dataset/update_dataset.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 4c4dac5..d720815 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -33,10 +33,11 @@ class ARDUpdater: def update_source(self, source: str): logger.info(f"Updating {source} entries...") - iterable_data = load_dataset(ARD_DATASET_NAME, source, split='train', streaming=True) - iterable_data = iterable_data.map(self.preprocess) - iterable_data = iterable_data.filter(lambda entry: entry is not None) - iterable_data = iterable_data.filter(lambda entry: self.sql_db.upsert_entry(entry)) + iterable_data = load_dataset( + ARD_DATASET_NAME, source, split='train', streaming=True + ).map(self.preprocess).filter( + lambda entry: entry is not None + ).filter(lambda entry: self.sql_db.upsert_entry(entry)) for entry in tqdm(iterable_data): try: From 560b132e6778ccf514b5d6dcc1a32f7af04508d1 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 07:21:45 -0400 Subject: [PATCH 36/49] Added filenames on line 1 and removed RESET_DBS, fixed logging bug --- src/dataset/pinecone_db_handler.py | 2 +- src/dataset/settings.py | 7 +++---- src/dataset/sql_db_handler.py | 2 ++ src/dataset/text_splitter.py | 2 ++ src/main.py | 9 +++++---- 5 files changed, 13 insertions(+), 9 deletions(-) diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py index e62e91d..7ecef06 100644 --- a/src/dataset/pinecone_db_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -1,4 +1,4 @@ -# dataset/pinecone_handler.py +# dataset/pinecone_db_handler.py import json import pinecone diff --git a/src/dataset/settings.py b/src/dataset/settings.py index ec04ef2..dcd4c91 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -1,3 +1,5 @@ +# dataset/settings.py + from pathlib import Path ### FILE PATHS ### @@ -16,7 +18,4 @@ EMBEDDINGS_RATE_LIMIT = 3500 PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" PINECONE_VALUES_DIMS = EMBEDDINGS_DIMS PINECONE_METRIC = "cosine" -PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] - -### MISC ### -RESET_DBS = False \ No newline at end of file +PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] \ No newline at end of file diff --git a/src/dataset/sql_db_handler.py b/src/dataset/sql_db_handler.py index 9ba223d..008a252 100644 --- a/src/dataset/sql_db_handler.py +++ b/src/dataset/sql_db_handler.py @@ -1,3 +1,5 @@ +# dataset/sql_db_handler.py + import sqlite3 from typing import List, Dict, Any diff --git a/src/dataset/text_splitter.py b/src/dataset/text_splitter.py index 221a7b2..2a99f6a 100644 --- a/src/dataset/text_splitter.py +++ b/src/dataset/text_splitter.py @@ -1,3 +1,5 @@ +# dataset/text_splitter.py + import re from typing import List import tiktoken diff --git a/src/main.py b/src/main.py index 770ac39..a38c918 100644 --- a/src/main.py +++ b/src/main.py @@ -10,17 +10,18 @@ else: import openai openai.api_key = os.environ.get('OPENAI_API_KEY') +import logging +logging.basicConfig(level=logging.INFO) + from dataset.update_dataset import ARDUpdater -from dataset.settings import RESET_DBS + def update_sql_and_pinecone_dbs(): updater = ARDUpdater( min_tokens_per_block=200, max_tokens_per_block=300, ) - if RESET_DBS: - updater.reset_dbs() - updater.update(['gwern_blog']) + updater.update() def main(): From 11dc2dd6c16ae9962c863ac0d321d584f2afe29d Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 07:26:41 -0400 Subject: [PATCH 37/49] Fixed entry typing and renamed var for clarity --- src/dataset/update_dataset.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index d720815..580c509 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -1,6 +1,6 @@ # dataset/update_dataset.py -from typing import Dict, List +from typing import Dict, List, Union import numpy as np from tqdm.auto import tqdm import openai @@ -71,7 +71,7 @@ class ARDUpdater: logger.error(f"Entry validation failed: {str(e)}", exc_info=True) return None - def validate_entry(self, entry: Dict[str, str | list], len_lower_limit: int = 0): + def validate_entry(self, entry: Dict[str, Union[str, list]], char_len_lower_limit: int = 0): metadata_types = { 'id': str, 'source': str, @@ -86,8 +86,8 @@ class ARDUpdater: if not isinstance(entry.get(metadata_type), metadata_type_type): raise ValueError(f"Entry metadata '{metadata_type}' is not of type '{metadata_type_type}' or is missing.") - if len(entry['text']) < len_lower_limit: - raise ValueError(f"Entry text is too short (< {len_lower_limit} tokens).") + if len(entry['text']) < char_len_lower_limit: + raise ValueError(f"Entry text is too short (< {char_len_lower_limit} characters).") def get_embeddings(self, chunks): embeddings = np.zeros((len(chunks), EMBEDDINGS_DIMS)) From 49e9f08cf5bb0ffe28c9487d0095dcc6a8c500e0 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 07:34:14 -0400 Subject: [PATCH 38/49] Removed magic number 3 by using a settings const --- src/dataset/settings.py | 5 ++++- src/dataset/update_dataset.py | 6 +++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/dataset/settings.py b/src/dataset/settings.py index dcd4c91..063f22d 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -18,4 +18,7 @@ EMBEDDINGS_RATE_LIMIT = 3500 PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" PINECONE_VALUES_DIMS = EMBEDDINGS_DIMS PINECONE_METRIC = "cosine" -PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] \ No newline at end of file +PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] + +### MISCELLANEOUS ### +MAX_NUM_AUTHORS_IN_SIGNATURE = 3 \ No newline at end of file diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 580c509..5761cc2 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -10,7 +10,7 @@ from .text_splitter import TokenSplitter from .sql_db_handler import SQLDB from .pinecone_db_handler import PineconeDB -from .settings import EMBEDDINGS_MODEL, EMBEDDINGS_DIMS, EMBEDDINGS_RATE_LIMIT, ARD_DATASET_NAME +from .settings import EMBEDDINGS_MODEL, EMBEDDINGS_DIMS, EMBEDDINGS_RATE_LIMIT, ARD_DATASET_NAME, MAX_NUM_AUTHORS_IN_SIGNATURE import logging logger = logging.getLogger(__name__) @@ -114,6 +114,6 @@ def get_authors_str(authors_lst: List[str]) -> str: if authors_lst == []: return 'n/a' if len(authors_lst) == 1: return authors_lst[0] else: - authors_lst = authors_lst[:3] - authors_str = ", ".join(authors_lst[:-1]) + " and " + authors_lst[-1] + authors_lst = authors_lst[:MAX_NUM_AUTHORS_IN_SIGNATURE] + authors_str = f"{', '.join(authors_lst[:-1])} and {authors_lst[-1]}" return authors_str \ No newline at end of file From 67f0dafd14fdcc7fa28debd64434117fd237a060 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 07:48:18 -0400 Subject: [PATCH 39/49] Make env var defs fail-fast --- src/dataset/pinecone_db_handler.py | 6 ++---- src/main.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py index 7ecef06..f5d7187 100644 --- a/src/dataset/pinecone_db_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -17,10 +17,8 @@ class PineconeDB: ): self.index_name = PINECONE_INDEX_NAME - PINECONE_API_KEY = os.getenv("PINECONE_API_KEY") - PINECONE_ENVIRONMENT = os.getenv("PINECONE_ENVIRONMENT") - assert PINECONE_API_KEY, "PINECONE_API_KEY environment variable not set." - assert PINECONE_ENVIRONMENT, "PINECONE_LOCATION environment variable not set." + PINECONE_API_KEY = os.environ["PINECONE_API_KEY"] + PINECONE_ENVIRONMENT = os.environ["PINECONE_ENVIRONMENT"] pinecone.init( api_key = PINECONE_API_KEY, diff --git a/src/main.py b/src/main.py index a38c918..73a79e7 100644 --- a/src/main.py +++ b/src/main.py @@ -8,7 +8,7 @@ else: raise Exception("'src/.env' not found. Rename the 'src/.env.example' file and fill in values.") import openai -openai.api_key = os.environ.get('OPENAI_API_KEY') +openai.api_key = os.environ['OPENAI_API_KEY'] import logging logging.basicConfig(level=logging.INFO) From 0ec3a8f5e460d2c724d975a16451baf04e4440af Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 07:50:16 -0400 Subject: [PATCH 40/49] Reorder imports --- src/dataset/pinecone_db_handler.py | 2 +- src/dataset/sql_db_handler.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py index f5d7187..07c3286 100644 --- a/src/dataset/pinecone_db_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -1,8 +1,8 @@ # dataset/pinecone_db_handler.py +import os import json import pinecone -import os from .settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES diff --git a/src/dataset/sql_db_handler.py b/src/dataset/sql_db_handler.py index 008a252..d48d46d 100644 --- a/src/dataset/sql_db_handler.py +++ b/src/dataset/sql_db_handler.py @@ -1,7 +1,7 @@ # dataset/sql_db_handler.py +from typing import List, Dict, Union import sqlite3 -from typing import List, Dict, Any from .settings import SQL_DB_PATH @@ -52,7 +52,7 @@ class SQLDB: except sqlite3.Error as e: logger.error(f"The error '{e}' occurred.") - def upsert_entry(self, entry: Dict[str, Any]) -> bool: + def upsert_entry(self, entry: Dict[str, Union[str, list]]) -> bool: with sqlite3.connect(self.db_name) as conn: cursor = conn.cursor() try: From 084b7139a1e53834e486011c4f922223aa1b1177 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 10:19:41 -0400 Subject: [PATCH 41/49] Reformatting --- src/dataset/pinecone_db_handler.py | 28 ++++++++++++++-------------- src/dataset/update_dataset.py | 2 +- src/main.py | 8 ++------ 3 files changed, 17 insertions(+), 21 deletions(-) diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py index 07c3286..b1eecff 100644 --- a/src/dataset/pinecone_db_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -37,20 +37,20 @@ class PineconeDB: def upsert_entry(self, entry, chunks, embeddings, upsert_size=100): self.index.upsert( vectors=list( - zip( - [f"{entry['id']}_{str(i).zfill(6)}" for i in range(len(chunks))], - embeddings.tolist(), - [ - { - 'entry_id': entry['id'], - 'source': entry['source'], - 'title': entry['title'], - 'authors': entry['authors'], - 'text': chunk, - } for chunk in chunks - ] - ) - ), + zip( + [f"{entry['id']}_{str(i).zfill(6)}" for i in range(len(chunks))], + embeddings.tolist(), + [ + { + 'entry_id': entry['id'], + 'source': entry['source'], + 'title': entry['title'], + 'authors': entry['authors'], + 'text': chunk, + } for chunk in chunks + ] + ) + ), batch_size=upsert_size ) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 5761cc2..23dde50 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -3,8 +3,8 @@ from typing import Dict, List, Union import numpy as np from tqdm.auto import tqdm -import openai from datasets import load_dataset +import openai from .text_splitter import TokenSplitter from .sql_db_handler import SQLDB diff --git a/src/main.py b/src/main.py index 73a79e7..25772dc 100644 --- a/src/main.py +++ b/src/main.py @@ -21,12 +21,8 @@ def update_sql_and_pinecone_dbs(): min_tokens_per_block=200, max_tokens_per_block=300, ) - updater.update() - - -def main(): - update_sql_and_pinecone_dbs() + updater.update(['gwern_blog']) if __name__ == "__main__": - main() \ No newline at end of file + update_sql_and_pinecone_dbs() \ No newline at end of file From 2025940db9d8a7785c80a47510122e7a1df8af28 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 10:36:15 -0400 Subject: [PATCH 42/49] Moving Pinecone env variables assignments to settings --- src/dataset/pinecone_db_handler.py | 5 +---- src/dataset/settings.py | 3 +++ 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py index b1eecff..39e695a 100644 --- a/src/dataset/pinecone_db_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -4,7 +4,7 @@ import os import json import pinecone -from .settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES +from .settings import PINECONE_INDEX_NAME, PINECONE_VALUES_DIMS, PINECONE_METRIC, PINECONE_METADATA_ENTRIES, PINECONE_API_KEY, PINECONE_ENVIRONMENT import logging logger = logging.getLogger(__name__) @@ -17,9 +17,6 @@ class PineconeDB: ): self.index_name = PINECONE_INDEX_NAME - PINECONE_API_KEY = os.environ["PINECONE_API_KEY"] - PINECONE_ENVIRONMENT = os.environ["PINECONE_ENVIRONMENT"] - pinecone.init( api_key = PINECONE_API_KEY, environment = PINECONE_ENVIRONMENT, diff --git a/src/dataset/settings.py b/src/dataset/settings.py index 063f22d..4438ca8 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -1,5 +1,6 @@ # dataset/settings.py +import os from pathlib import Path ### FILE PATHS ### @@ -19,6 +20,8 @@ PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" PINECONE_VALUES_DIMS = EMBEDDINGS_DIMS PINECONE_METRIC = "cosine" PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] +PINECONE_API_KEY = os.environ["PINECONE_API_KEY"] +PINECONE_ENVIRONMENT = os.environ["PINECONE_ENVIRONMENT"] ### MISCELLANEOUS ### MAX_NUM_AUTHORS_IN_SIGNATURE = 3 \ No newline at end of file From 987d565b11e5055c5b7df1728a373a825f7e3bbe Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Tue, 27 Jun 2023 11:20:41 -0400 Subject: [PATCH 43/49] Added retry to the embeddings method --- src/dataset/update_dataset.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 23dde50..cbc65ea 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -2,6 +2,7 @@ from typing import Dict, List, Union import numpy as np +from tenacity import retry, stop_after_attempt from tqdm.auto import tqdm from datasets import load_dataset import openai @@ -88,7 +89,8 @@ class ARDUpdater: if len(entry['text']) < char_len_lower_limit: raise ValueError(f"Entry text is too short (< {char_len_lower_limit} characters).") - + + @retry(stop=stop_after_attempt(3)) def get_embeddings(self, chunks): embeddings = np.zeros((len(chunks), EMBEDDINGS_DIMS)) rate_limit = EMBEDDINGS_RATE_LIMIT #TODO: use this rate_limit From f1f0c8a21b44672eb8d27559553f0aae7fd678ba Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Sun, 2 Jul 2023 06:01:03 -0400 Subject: [PATCH 44/49] Implemented upsert_entries in the Pinecone handler --- src/dataset/pinecone_db_handler.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py index 39e695a..b0090d3 100644 --- a/src/dataset/pinecone_db_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -50,11 +50,38 @@ class PineconeDB: ), batch_size=upsert_size ) - + + def upsert_entries(self, entries_batch, chunks_batch, chunks_ids_batch, embeddings, upsert_size=100): + self.index.upsert( + vectors=list( + zip( + chunks_ids_batch, + embeddings.tolist(), + [ + { + 'entry_id': entry['id'], + 'source': entry['source'], + 'title': entry['title'], + 'authors': entry['authors'], + 'text': chunk, + } + for entry in entries_batch + for chunk in chunks_batch + ] + ) + ), + batch_size=upsert_size + ) + def delete_entry(self, id): self.index.delete( filter={"entry_id": {"$eq": id}} ) + + def delete_entries(self, ids): + self.index.delete( + filter={"entry_id": {"$in": ids}} + ) def create_index(self, replace_current_index: bool = True): if replace_current_index: From 5321af981384b59134f035495802aee356fb612b Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Sun, 2 Jul 2023 06:01:31 -0400 Subject: [PATCH 45/49] Implemented upsert_chunks in the sql handler --- src/dataset/sql_db_handler.py | 23 +++++++++-------------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/src/dataset/sql_db_handler.py b/src/dataset/sql_db_handler.py index d48d46d..747d7b4 100644 --- a/src/dataset/sql_db_handler.py +++ b/src/dataset/sql_db_handler.py @@ -87,23 +87,18 @@ class SQLDB: finally: conn.commit() - - def upsert_chunks(self, entry_id: str, chunks: List[str]) -> bool: + + def upsert_chunks(self, chunks_ids_batch: List[str], chunks_batch: List[str]) -> bool: with sqlite3.connect(self.db_name) as conn: cursor = conn.cursor() try: - # Delete existing chunks - cursor.execute("DELETE FROM chunk_database WHERE entry_id=?", (entry_id,)) - - # Insert new chunks - for i, chunk in enumerate(chunks): - chunk_id = f"{entry_id}_{str(i).zfill(6)}" - cursor.execute("INSERT INTO chunk_database (id, text, entry_id) VALUES (?, ?, ?)", (chunk_id, chunk, entry_id)) - return True - + for chunk_id, chunk in zip(chunks_ids_batch, chunks_batch): + cursor.execute(""" + INSERT OR REPLACE INTO chunk_database + (id, text) + VALUES (?, ?) + """, (chunk_id, chunk)) except sqlite3.Error as e: logger.error(f"The error '{e}' occurred.") - return False - finally: - conn.commit() \ No newline at end of file + conn.commit() From 0cf2fdad74ca19eefc2cf0d504cf3dddddac8965 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Sun, 2 Jul 2023 06:03:05 -0400 Subject: [PATCH 46/49] Implemented batch updating --- src/dataset/update_dataset.py | 79 ++++++++++++++++++++++++++--------- 1 file changed, 60 insertions(+), 19 deletions(-) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index cbc65ea..6efb9e4 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -30,32 +30,64 @@ class ARDUpdater: def update(self, custom_sources: List[str] = ['all']): for source in custom_sources: self.update_source(source) - - def update_source(self, source: str): + + def update_source(self, source: str, chunk_size: int = 100): logger.info(f"Updating {source} entries...") - - iterable_data = load_dataset( - ARD_DATASET_NAME, source, split='train', streaming=True - ).map(self.preprocess).filter( - lambda entry: entry is not None - ).filter(lambda entry: self.sql_db.upsert_entry(entry)) - for entry in tqdm(iterable_data): - try: - self.pinecone_db.delete_entry(entry['id']) + streamed_dataset = load_dataset( + ARD_DATASET_NAME, source, split='train', streaming=True + ).map(self.preprocess_and_validate).filter( + self.is_valid_entry + ).filter( + self.is_sql_entry_upserted + ) + + for batch in self.batchify(streamed_dataset, chunk_size): + entries_batch = batch['entries_batch'] + chunks_batch = batch['chunks_batch'] + chunks_ids_batch = batch['chunks_ids_batch'] - signature = f"Title: {entry['title']}, Authors: {get_authors_str(entry['authors'])}" - chunks = self.token_splitter.split(entry['text'], signature) - embeddings = self.get_embeddings(chunks) + try: + embeddings = self.get_embeddings(chunks_batch) + + self.sql_db.upsert_chunks(chunks_ids_batch, chunks_batch) + self.pinecone_db.delete_entries([entry['id'] for entry in entries_batch]) + self.pinecone_db.upsert_entries(entries_batch, chunks_batch, chunks_ids_batch, embeddings) - self.sql_db.upsert_chunks(entry['id'], chunks) - self.pinecone_db.upsert_entry(entry, chunks, embeddings) + logger.info(f"Successfully updated {len(entries_batch)} {source} entries with {len(chunks_batch)} total chunks.") except Exception as e: logger.error(f"An error occurred while updating source {source}: {str(e)}", exc_info=True) - + logger.info(f"Successfully updated {source} entries.") + + def batchify(self, iterable, chunk_size): + entries_batch = [] + chunks_batch = [] + chunks_ids_batch = [] + + for entry in iterable: + chunks = self.token_splitter.split(entry['text'], f"Title: {entry['title']}, Authors: {get_authors_str(entry['authors'])}") + chunks_ids = [f"{entry['id']}_{str(i).zfill(6)}" for i in range(len(chunks))] + + # Add this entry's chunks to the current batch, even if it causes the batch size to exceed chunk_size. + entries_batch.append(entry) + chunks_batch.extend(chunks) + chunks_ids_batch.extend(chunks_ids) + + # If this batch is large enough, yield it and start a new one. + if len(chunks_batch) >= chunk_size: + yield {'entries_batch': entries_batch, 'chunks_batch': chunks_batch, 'chunks_ids_batch': chunks_ids_batch} + + entries_batch = [] + chunks_batch = [] + chunks_ids_batch = [] + + # Yield any remaining items. + if entries_batch: + yield {'entries_batch': entries_batch, 'chunks_batch': chunks_batch, 'chunks_ids_batch': chunks_ids_batch} - def preprocess(self, entry): + def preprocess_and_validate(self, entry): + """Preprocesses and validates the entry data""" try: self.validate_entry(entry) @@ -89,11 +121,20 @@ class ARDUpdater: if len(entry['text']) < char_len_lower_limit: raise ValueError(f"Entry text is too short (< {char_len_lower_limit} characters).") + + @staticmethod + def is_valid_entry(entry): + """Checks if the entry is valid""" + return entry is not None + + def is_sql_entry_upserted(self, entry): + """Upserts an entry to the SQL database and returns the success status""" + return self.sql_db.upsert_entry(entry) @retry(stop=stop_after_attempt(3)) def get_embeddings(self, chunks): embeddings = np.zeros((len(chunks), EMBEDDINGS_DIMS)) - rate_limit = EMBEDDINGS_RATE_LIMIT #TODO: use this rate_limit + rate_limit = EMBEDDINGS_RATE_LIMIT # TODO: use this rate_limit openai_output = openai.Embedding.create( model=EMBEDDINGS_MODEL, From cdfb5a52a774d5364942cc9741f3bf5780330680 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Sun, 2 Jul 2023 07:09:06 -0400 Subject: [PATCH 47/49] Added sentence transformer option --- src/dataset/settings.py | 6 ++++-- src/dataset/update_dataset.py | 23 ++++++++++++++++------- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/src/dataset/settings.py b/src/dataset/settings.py index 4438ca8..f524191 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -11,9 +11,11 @@ SQL_DB_PATH = str(current_file_path.parent / 'data' / 'ARD.db') ARD_DATASET_NAME = "StampyAI/alignment-research-dataset" ### EMBEDDINGS ### -EMBEDDINGS_MODEL = "text-embedding-ada-002" +USE_OPENAI_EMBEDDINGS = False +OPENAI_EMBEDDINGS_MODEL = "text-embedding-ada-002" +SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL = "sentence-transformers/multi-qa-mpnet-base-cos-v1" EMBEDDINGS_DIMS = 1536 -EMBEDDINGS_RATE_LIMIT = 3500 +OPENAI_EMBEDDINGS_RATE_LIMIT = 3500 ### PINECONE ### PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 6efb9e4..671a8d4 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -11,7 +11,7 @@ from .text_splitter import TokenSplitter from .sql_db_handler import SQLDB from .pinecone_db_handler import PineconeDB -from .settings import EMBEDDINGS_MODEL, EMBEDDINGS_DIMS, EMBEDDINGS_RATE_LIMIT, ARD_DATASET_NAME, MAX_NUM_AUTHORS_IN_SIGNATURE +from .settings import USE_OPENAI_EMBEDDINGS, OPENAI_EMBEDDINGS_MODEL, SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, EMBEDDINGS_DIMS, OPENAI_EMBEDDINGS_RATE_LIMIT, ARD_DATASET_NAME, MAX_NUM_AUTHORS_IN_SIGNATURE import logging logger = logging.getLogger(__name__) @@ -26,6 +26,12 @@ class ARDUpdater: self.token_splitter = TokenSplitter(min_tokens_per_block, max_tokens_per_block) self.sql_db = SQLDB() self.pinecone_db = PineconeDB() + + if not USE_OPENAI_EMBEDDINGS: + from langchain.embeddings import HuggingFaceEmbeddings + self.hf_embeddings = HuggingFaceEmbeddings( + model_name=SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, + ) def update(self, custom_sources: List[str] = ['all']): for source in custom_sources: @@ -48,11 +54,14 @@ class ARDUpdater: chunks_ids_batch = batch['chunks_ids_batch'] try: - embeddings = self.get_embeddings(chunks_batch) + if USE_OPENAI_EMBEDDINGS: + embeddings = self.get_openai_embeddings(chunks_batch) + else: + embeddings = np.array(self.hf_embeddings.embed_documents(chunks_batch)) self.sql_db.upsert_chunks(chunks_ids_batch, chunks_batch) - self.pinecone_db.delete_entries([entry['id'] for entry in entries_batch]) - self.pinecone_db.upsert_entries(entries_batch, chunks_batch, chunks_ids_batch, embeddings) + # self.pinecone_db.delete_entries([entry['id'] for entry in entries_batch]) + # self.pinecone_db.upsert_entries(entries_batch, chunks_batch, chunks_ids_batch, embeddings) logger.info(f"Successfully updated {len(entries_batch)} {source} entries with {len(chunks_batch)} total chunks.") except Exception as e: @@ -132,12 +141,12 @@ class ARDUpdater: return self.sql_db.upsert_entry(entry) @retry(stop=stop_after_attempt(3)) - def get_embeddings(self, chunks): + def get_openai_embeddings(self, chunks): embeddings = np.zeros((len(chunks), EMBEDDINGS_DIMS)) - rate_limit = EMBEDDINGS_RATE_LIMIT # TODO: use this rate_limit + rate_limit = OPENAI_EMBEDDINGS_RATE_LIMIT # TODO: use this rate_limit openai_output = openai.Embedding.create( - model=EMBEDDINGS_MODEL, + model=OPENAI_EMBEDDINGS_MODEL, input=chunks )['data'] From aeba2fbfb05c10f74295e92837081130b3040443 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Sun, 2 Jul 2023 07:22:38 -0400 Subject: [PATCH 48/49] bug-fix and fixed device option --- src/dataset/settings.py | 4 +++- src/dataset/update_dataset.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/dataset/settings.py b/src/dataset/settings.py index f524191..951e138 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -1,6 +1,7 @@ # dataset/settings.py import os +import torch from pathlib import Path ### FILE PATHS ### @@ -13,9 +14,10 @@ ARD_DATASET_NAME = "StampyAI/alignment-research-dataset" ### EMBEDDINGS ### USE_OPENAI_EMBEDDINGS = False OPENAI_EMBEDDINGS_MODEL = "text-embedding-ada-002" -SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL = "sentence-transformers/multi-qa-mpnet-base-cos-v1" EMBEDDINGS_DIMS = 1536 OPENAI_EMBEDDINGS_RATE_LIMIT = 3500 +SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL = "sentence-transformers/multi-qa-mpnet-base-cos-v1" +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" ### PINECONE ### PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 671a8d4..279b20a 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -11,7 +11,7 @@ from .text_splitter import TokenSplitter from .sql_db_handler import SQLDB from .pinecone_db_handler import PineconeDB -from .settings import USE_OPENAI_EMBEDDINGS, OPENAI_EMBEDDINGS_MODEL, SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, EMBEDDINGS_DIMS, OPENAI_EMBEDDINGS_RATE_LIMIT, ARD_DATASET_NAME, MAX_NUM_AUTHORS_IN_SIGNATURE +from .settings import USE_OPENAI_EMBEDDINGS, OPENAI_EMBEDDINGS_MODEL, SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, EMBEDDINGS_DIMS, OPENAI_EMBEDDINGS_RATE_LIMIT, DEVICE, ARD_DATASET_NAME, MAX_NUM_AUTHORS_IN_SIGNATURE import logging logger = logging.getLogger(__name__) @@ -31,6 +31,8 @@ class ARDUpdater: from langchain.embeddings import HuggingFaceEmbeddings self.hf_embeddings = HuggingFaceEmbeddings( model_name=SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, + model_kwargs={'device': DEVICE}, + encode_kwargs={'show_progress_bar': True} ) def update(self, custom_sources: List[str] = ['all']): From 511f231872cdb7923a4abe08f328a2cf82545d66 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Sun, 2 Jul 2023 07:23:13 -0400 Subject: [PATCH 49/49] Uncommented pinecone functionality --- src/dataset/update_dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 279b20a..f07b6c1 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -62,8 +62,8 @@ class ARDUpdater: embeddings = np.array(self.hf_embeddings.embed_documents(chunks_batch)) self.sql_db.upsert_chunks(chunks_ids_batch, chunks_batch) - # self.pinecone_db.delete_entries([entry['id'] for entry in entries_batch]) - # self.pinecone_db.upsert_entries(entries_batch, chunks_batch, chunks_ids_batch, embeddings) + self.pinecone_db.delete_entries([entry['id'] for entry in entries_batch]) + self.pinecone_db.upsert_entries(entries_batch, chunks_batch, chunks_ids_batch, embeddings) logger.info(f"Successfully updated {len(entries_batch)} {source} entries with {len(chunks_batch)} total chunks.") except Exception as e: