From 689a9a7ae410ae5916a998d1f456db941e520939 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Mon, 3 Jul 2023 01:02:41 -0400 Subject: [PATCH 1/8] clean up --- .gitignore | 7 +------ src/dataset/settings.py | 9 +++++++-- src/dataset/update_dataset.py | 16 ++++++++-------- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/.gitignore b/.gitignore index 5bfc436..950aea0 100644 --- a/.gitignore +++ b/.gitignore @@ -142,13 +142,8 @@ dmypy.json temp/ *tmp.py -api/dataset.pkl -api/dataset_big.pkl -api/dataset_300.pkl - api/.env.backup src/dataset_tests.ipynb src/ARD_LangChain_QA_Chat.ipynb - -src/dataset/data/* \ No newline at end of file +src/dataset/data/ARD.db \ No newline at end of file diff --git a/src/dataset/settings.py b/src/dataset/settings.py index 951e138..405fe9a 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -13,19 +13,24 @@ ARD_DATASET_NAME = "StampyAI/alignment-research-dataset" ### EMBEDDINGS ### USE_OPENAI_EMBEDDINGS = False + OPENAI_EMBEDDINGS_MODEL = "text-embedding-ada-002" -EMBEDDINGS_DIMS = 1536 +OPENAI_EMBEDDINGS_DIMS = 1536 OPENAI_EMBEDDINGS_RATE_LIMIT = 3500 + SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL = "sentence-transformers/multi-qa-mpnet-base-cos-v1" +SENTENCE_TRANSFORMER_EMBEDDINGS_DIMS = 768 + DEVICE = "cuda" if torch.cuda.is_available() else "cpu" ### PINECONE ### PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" -PINECONE_VALUES_DIMS = EMBEDDINGS_DIMS +PINECONE_VALUES_DIMS = OPENAI_EMBEDDINGS_DIMS if USE_OPENAI_EMBEDDINGS else SENTENCE_TRANSFORMER_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 ### +CHUNK_SIZE = 5000 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 f07b6c1..5092947 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, DEVICE, ARD_DATASET_NAME, MAX_NUM_AUTHORS_IN_SIGNATURE +from .settings import USE_OPENAI_EMBEDDINGS, OPENAI_EMBEDDINGS_MODEL, OPENAI_EMBEDDINGS_DIMS, OPENAI_EMBEDDINGS_RATE_LIMIT, SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, SENTENCE_TRANSFORMER_EMBEDDINGS_DIMS, DEVICE, ARD_DATASET_NAME, CHUNK_SIZE, MAX_NUM_AUTHORS_IN_SIGNATURE import logging logger = logging.getLogger(__name__) @@ -32,14 +32,14 @@ class ARDUpdater: self.hf_embeddings = HuggingFaceEmbeddings( model_name=SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, model_kwargs={'device': DEVICE}, - encode_kwargs={'show_progress_bar': True} + encode_kwargs={'show_progress_bar': False} ) def update(self, custom_sources: List[str] = ['all']): for source in custom_sources: self.update_source(source) - def update_source(self, source: str, chunk_size: int = 100): + def update_source(self, source: str): logger.info(f"Updating {source} entries...") streamed_dataset = load_dataset( @@ -50,7 +50,7 @@ class ARDUpdater: self.is_sql_entry_upserted ) - for batch in self.batchify(streamed_dataset, chunk_size): + for batch in self.batchify(streamed_dataset): entries_batch = batch['entries_batch'] chunks_batch = batch['chunks_batch'] chunks_ids_batch = batch['chunks_ids_batch'] @@ -71,7 +71,7 @@ class ARDUpdater: logger.info(f"Successfully updated {source} entries.") - def batchify(self, iterable, chunk_size): + def batchify(self, iterable): entries_batch = [] chunks_batch = [] chunks_ids_batch = [] @@ -80,13 +80,13 @@ class ARDUpdater: 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. + # 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: + if len(chunks_batch) >= CHUNK_SIZE: yield {'entries_batch': entries_batch, 'chunks_batch': chunks_batch, 'chunks_ids_batch': chunks_ids_batch} entries_batch = [] @@ -144,7 +144,7 @@ class ARDUpdater: @retry(stop=stop_after_attempt(3)) def get_openai_embeddings(self, chunks): - embeddings = np.zeros((len(chunks), EMBEDDINGS_DIMS)) + embeddings = np.zeros((len(chunks), OPENAI_EMBEDDINGS_DIMS)) rate_limit = OPENAI_EMBEDDINGS_RATE_LIMIT # TODO: use this rate_limit openai_output = openai.Embedding.create( From 07025fb24f1002645c5513b61b20c84e87f9dd03 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Mon, 17 Jul 2023 17:06:11 -0400 Subject: [PATCH 2/8] clean main --- src/main.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/main.py b/src/main.py index 25772dc..5017f42 100644 --- a/src/main.py +++ b/src/main.py @@ -17,11 +17,8 @@ from dataset.update_dataset import ARDUpdater def update_sql_and_pinecone_dbs(): - updater = ARDUpdater( - min_tokens_per_block=200, - max_tokens_per_block=300, - ) - updater.update(['gwern_blog']) + updater = ARDUpdater() + updater.update() if __name__ == "__main__": From e61d423714d023e54f5383956ec7ad8158772fed Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Mon, 17 Jul 2023 17:12:19 -0400 Subject: [PATCH 3/8] fixed pinecone info, added aisafety.info bias --- src/dataset/settings.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/dataset/settings.py b/src/dataset/settings.py index 405fe9a..d1e4aac 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -12,7 +12,7 @@ SQL_DB_PATH = str(current_file_path.parent / 'data' / 'ARD.db') ARD_DATASET_NAME = "StampyAI/alignment-research-dataset" ### EMBEDDINGS ### -USE_OPENAI_EMBEDDINGS = False +USE_OPENAI_EMBEDDINGS = True # If false, SentenceTransformer embeddings will be used. OPENAI_EMBEDDINGS_MODEL = "text-embedding-ada-002" OPENAI_EMBEDDINGS_DIMS = 1536 @@ -21,16 +21,17 @@ OPENAI_EMBEDDINGS_RATE_LIMIT = 3500 SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL = "sentence-transformers/multi-qa-mpnet-base-cos-v1" SENTENCE_TRANSFORMER_EMBEDDINGS_DIMS = 768 -DEVICE = "cuda" if torch.cuda.is_available() else "cpu" - ### PINECONE ### -PINECONE_INDEX_NAME = "stampy-chat-embeddings-test" +PINECONE_INDEX_NAME = "stampy-chat-ard" PINECONE_VALUES_DIMS = OPENAI_EMBEDDINGS_DIMS if USE_OPENAI_EMBEDDINGS else SENTENCE_TRANSFORMER_EMBEDDINGS_DIMS -PINECONE_METRIC = "cosine" +PINECONE_METRIC = "dotproduct" PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] -PINECONE_API_KEY = os.environ["PINECONE_API_KEY"] -PINECONE_ENVIRONMENT = os.environ["PINECONE_ENVIRONMENT"] +PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY", None) +PINECONE_ENVIRONMENT = os.environ.get("PINECONE_ENVIRONMENT", None) ### MISCELLANEOUS ### -CHUNK_SIZE = 5000 -MAX_NUM_AUTHORS_IN_SIGNATURE = 3 \ No newline at end of file +CHUNK_SIZE = 1750 +MAX_NUM_AUTHORS_IN_SIGNATURE = 3 +EMBEDDING_LENGTH_BIAS = { + "aisafety.info": 1.05, # In search, favor AISafety.info entries. +} \ No newline at end of file From d08477b036fe31d0ece3299221ef1228c6777546 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Mon, 17 Jul 2023 17:13:48 -0400 Subject: [PATCH 4/8] clean up code --- src/dataset/pinecone_db_handler.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py index b0090d3..c9dcdee 100644 --- a/src/dataset/pinecone_db_handler.py +++ b/src/dataset/pinecone_db_handler.py @@ -12,10 +12,18 @@ logger = logging.getLogger(__name__) class PineconeDB: def __init__( - self, + self, + index_name: str = PINECONE_INDEX_NAME, + values_dims: int = PINECONE_VALUES_DIMS, + metric: str = PINECONE_METRIC, + metadata_entries: list = PINECONE_METADATA_ENTRIES, create_index: bool = False, + log_index_stats: bool = True, ): - self.index_name = PINECONE_INDEX_NAME + self.index_name = index_name + self.values_dims = values_dims + self.metric = metric + self.metadata_entries = metadata_entries pinecone.init( api_key = PINECONE_API_KEY, @@ -27,9 +35,9 @@ class PineconeDB: self.index = pinecone.Index(index_name=self.index_name) - def __str__(self) -> str: - index_stats_response = self.index.describe_index_stats() - return f"{self.index_name}:\n{json.dumps(index_stats_response, indent=4)}" + if log_index_stats: + index_stats_response = self.index.describe_index_stats() + logger.info(f"{self.index_name}:\n{index_stats_response}") def upsert_entry(self, entry, chunks, embeddings, upsert_size=100): self.index.upsert( @@ -89,9 +97,9 @@ class PineconeDB: pinecone.create_index( name=self.index_name, - dimension=PINECONE_VALUES_DIMS, - metric=PINECONE_METRIC, - metadata_config = {"indexed": PINECONE_METADATA_ENTRIES} + dimension=self.values_dims, + metric=self.metric, + metadata_config = {"indexed": self.metadata_entries}, ) def delete_index(self): From 0fcaeb4a617fa83050de0a8eee23b149e5cca25a Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Mon, 17 Jul 2023 17:15:07 -0400 Subject: [PATCH 5/8] add embeddings to chunk table, add stream_chunks --- src/dataset/sql_db_handler.py | 42 ++++++++++++++++++++++++++++------- 1 file changed, 34 insertions(+), 8 deletions(-) diff --git a/src/dataset/sql_db_handler.py b/src/dataset/sql_db_handler.py index 747d7b4..eb91efd 100644 --- a/src/dataset/sql_db_handler.py +++ b/src/dataset/sql_db_handler.py @@ -1,6 +1,7 @@ # dataset/sql_db_handler.py from typing import List, Dict, Union +import numpy as np import sqlite3 from .settings import SQL_DB_PATH @@ -10,8 +11,8 @@ logger = logging.getLogger(__name__) class SQLDB: - def __init__(self): - self.db_name = SQL_DB_PATH + def __init__(self, db_name: str = SQL_DB_PATH): + self.db_name = db_name self.create_tables() @@ -43,6 +44,7 @@ class SQLDB: CREATE TABLE IF NOT EXISTS chunk_database ( id TEXT PRIMARY KEY, text TEXT, + embedding BLOB, entry_id TEXT, FOREIGN KEY (entry_id) REFERENCES entry_database(id) ) @@ -87,18 +89,42 @@ class SQLDB: finally: conn.commit() - - def upsert_chunks(self, chunks_ids_batch: List[str], chunks_batch: List[str]) -> bool: + + def upsert_chunks(self, chunks_ids_batch: List[str], chunks_batch: List[str], embeddings_batch: List[np.ndarray]) -> bool: with sqlite3.connect(self.db_name) as conn: cursor = conn.cursor() try: - for chunk_id, chunk in zip(chunks_ids_batch, chunks_batch): + for chunk_id, chunk, embedding in zip(chunks_ids_batch, chunks_batch, embeddings_batch): cursor.execute(""" INSERT OR REPLACE INTO chunk_database - (id, text) - VALUES (?, ?) - """, (chunk_id, chunk)) + (id, text, embedding) + VALUES (?, ?, ?) + """, (chunk_id, chunk, embedding.tobytes())) except sqlite3.Error as e: logger.error(f"The error '{e}' occurred.") finally: conn.commit() + + + def stream_chunks(self): + with sqlite3.connect(self.db_name) as conn: + cursor = conn.cursor() + + # Join entry_database and chunk_database tables and order by source + cursor.execute(""" + SELECT c.id, c.text, c.embedding, e.source + FROM chunk_database c + JOIN entry_database e ON c.entry_id = e.id + ORDER BY e.source + """) + + for row in cursor: + # Convert bytes back to numpy array + embedding = np.frombuffer(row[2], dtype=np.float64) if row[2] else None + + yield { + 'id': row[0], + 'text': row[1], + 'embedding': embedding, + 'source': row[3], + } \ No newline at end of file From 838732f0366e8d6dc06ce5e2f329065edf017493 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Mon, 17 Jul 2023 17:18:03 -0400 Subject: [PATCH 6/8] switched to custom langchain text_splitter --- src/dataset/text_splitter.py | 298 ++++++++++++----------------------- 1 file changed, 102 insertions(+), 196 deletions(-) diff --git a/src/dataset/text_splitter.py b/src/dataset/text_splitter.py index 2a99f6a..db5f802 100644 --- a/src/dataset/text_splitter.py +++ b/src/dataset/text_splitter.py @@ -1,222 +1,128 @@ # dataset/text_splitter.py -import re -from typing import List +from typing import List, Callable, Any +from langchain.text_splitter import TextSplitter +from nltk.tokenize import sent_tokenize import tiktoken -import re -from typing import List -import nltk + +def char_len(string: str) -> int: + return len(string) + +def word_len(string: str) -> int: + return len(string.split(" ")) + +def token_len(string: str) -> int: + return len(tiktoken.get_encoding("cl100k_base").encode(string)) -# Download the Punkt tokenizer if you haven't already. -# If you want to save a second everytime you run this file you can comment -# it out after the first time it was downloaded. -nltk.download("punkt") +def char_truncate(string: str, length: int, from_end: bool = False) -> str: + if from_end: + return string[-length:] + else: + return string[:length] + +def word_truncate(string: str, length: int, from_end: bool = False) -> str: + words = string.split(" ") + if from_end: + return " ".join(words[-length:]) + else: + return " ".join(words[:length]) + +def token_truncate(string: str, length: int, from_end: bool = False) -> str: + tokens = tiktoken.get_encoding("cl100k_base").encode(string) + if from_end: + return tiktoken.get_encoding("cl100k_base").decode(tokens[-length:]) + else: + return tiktoken.get_encoding("cl100k_base").decode(tokens[:length]) +class ParagraphSentenceUnitTextSplitter(TextSplitter): + """A custom TextSplitter that breaks text by paragraphs, sentences, and then units (chars/words/tokens/etc).""" + + DEFAULT_MIN_CHUNK_SIZE = 900 + DEFAULT_MAX_CHUNK_SIZE = 1100 + DEFAULT_TRUNCATE_FUNCTION = char_truncate -def split_into_sentences(text: str) -> List[str]: - """ - Splits the input text into sentences. + def __init__( + self, + min_chunk_size: int = DEFAULT_MIN_CHUNK_SIZE, + max_chunk_size: int = DEFAULT_MAX_CHUNK_SIZE, + truncate_function: Callable[[str, int], str] = DEFAULT_TRUNCATE_FUNCTION, + **kwargs: Any + ): + super().__init__(**kwargs) + self.min_chunk_size = min_chunk_size + self.max_chunk_size = max_chunk_size - :param text: The input text to be split. - :return: A list of sentences. - """ - text = text.replace("\n", " ") # Replace newline characters with spaces - sentences = nltk.sent_tokenize(text) # Use the Punkt tokenizer from the NLTK library to split the text into sentences - sentences = [s.strip() for s in sentences] # Strip leading and trailing whitespace from each sentence - return sentences - - -class TokenSplitter: - """Splits text into blocks of tokens according to chatgpt's tokenizer.""" - - def __init__(self, min_tokens: int = 200, max_tokens: int = 300): - self.encoding = tiktoken.get_encoding("cl100k_base") - self.min_tokens = min_tokens - self.max_tokens = max_tokens - self.default_signature = "{url, title, author} unknown" - - def _text_splitter(self, text: str, signature: str) -> List[str]: - """Splits text into blocks of tokens according to chatgpt's tokenizer.""" - # enc = self.encoding.encode # takes a string and returns a list of ints (tokens) - enc = self.encoding.encode_ordinary # takes a string and returns a list of ints (tokens) - dec = self.encoding.decode # takes a list of ints (tokens) and returns a string - tok_len = lambda x: len(enc(x)) # length of a string in tokens - - max_tokens = self.max_tokens - tok_len(signature) - 10 # 10 to be safe - assert max_tokens > 0, "max_tokens is too small for the signature" - - min_tokens = self.min_tokens - tok_len(signature) - 10 # 10 to be safe - assert min_tokens > 0, "min_tokens is too small for the signature" + self._truncate_function = truncate_function + def split_text(self, text: str) -> List[str]: blocks = [] current_block = "" - paragraphs = text.split("\n\n") - - for paragraph in paragraphs: - sentences = split_into_sentences(paragraph) - if current_block != "": - current_block += "\n\n" - for sentence in sentences: - potential_new_block = f"{current_block} {sentence}" - - if tok_len(potential_new_block) <= max_tokens: - current_block = potential_new_block - - else: - blocks.append(current_block) - if tok_len(sentence) < max_tokens: - current_block = sentence - else: - blocks.append(dec(enc(sentence)[:max_tokens])) - current_block = "" - - if tok_len(current_block) > min_tokens: + paragraphs = text.split("\n\n") + for paragraph in paragraphs: + current_block += "\n\n" + paragraph + block_length = self._length_function(current_block) + + if block_length > self.max_chunk_size: # current block is too large, truncate it + current_block = self._handle_large_paragraph(current_block, blocks, paragraph) + elif block_length >= self.min_chunk_size: blocks.append(current_block) current_block = "" - - if current_block != "": - if len(blocks) == 0: - blocks.append(current_block) - else: - latest_block = blocks[-1] - len_cur_block = tok_len(current_block) - latest_plus_current = latest_block + current_block - - if len_cur_block > min_tokens: - blocks.append(current_block) - - else: - # select the last self.max_tokens tokens from the latest block - last_block = dec(enc(latest_plus_current)[-max_tokens:]) - blocks.append(last_block) + else: # current block is too small, continue appending to it + continue + + blocks = self._handle_remaining_text(current_block, blocks) return [block.strip() for block in blocks] - def split(self, text: str, signature: str = None) -> List[str]: - if signature is None: - signature = self.default_signature + def _handle_large_paragraph(self, current_block, blocks, paragraph): + # Undo adding the whole paragraph + current_block = current_block[:-(len(paragraph)+2)] # +2 accounts for "\n\n" - blocks = self._text_splitter(text, signature) + sentences = sent_tokenize(paragraph) + for sentence in sentences: + current_block += f" {sentence}" + + block_length = self._length_function(current_block) + if block_length < self.min_chunk_size: + continue + elif block_length <= self.max_chunk_size: + blocks.append(current_block) + current_block = "" + else: + current_block = self._truncate_large_block(current_block, blocks, sentence) + + return current_block - # Check all block elements are strings - assert all([isinstance(block, str) for block in blocks]), "block elements are not strings" + def _truncate_large_block(self, current_block, blocks, sentence): + while self._length_function(current_block) > self.max_chunk_size: + # Truncate current_block to max size, set remaining sentence as next sentence + truncated_block = self._truncate_function(current_block, self.max_chunk_size) + blocks.append(truncated_block) - output = [f'"{block}"\n- {signature}' for block in blocks] - # Check all output elements are strings - assert all([isinstance(block, str) for block in output]), "output elements are not strings" + remaining_sentence = current_block[len(truncated_block):].lstrip() + current_block = sentence = remaining_sentence + + return current_block - return output + def _handle_remaining_text(self, current_block, blocks): + if blocks == []: # no blocks were added + return [current_block] + elif current_block: # any leftover text + len_current_block = self._length_function(current_block) + if len_current_block < self.min_chunk_size: + # it needs to take the last min_chunk_size-len_current_block units from the previous block + previous_block = blocks[-1] + required_units = self.min_chunk_size - len_current_block # calculate the required units -if __name__ == "__main__": - text = """This post has been recorded as part of the LessWrong Curated Podcast, and an be listened to on Spotify, Apple Podcasts, and Libsyn. + part_prev_block = self._truncate_function(previous_block, required_units, from_end=True) # get the required units from the previous block + last_block = part_prev_block + current_block -Over the last few years, deep-learning-based AI has progressed extremely rapidly in fields like natural language processing and image generation. However, self-driving cars seem stuck in perpetual beta mode, and aggressive predictions there have repeatedly been disappointing. Google's self-driving project started four years before AlexNet kicked off the deep learning revolution, and it still isn't deployed at large scale, thirteen years later. Why are these fields getting such different results? + blocks.append(last_block) + else: + blocks.append(current_block) -Right now, I think the biggest answer is that ML benchmarks judge models by average-case performance, while self-driving cars (and many other applications) require matching human worst-case performance. For MNIST, an easy handwriting recognition task, performance tops out at around 99.9% even for top models; it's not very practical to design for or measure higher reliability than that, because the test set is just 10,000 images and a handful are ambiguous. Redwood Research, which is exploring worst-case performance in the context of AI alignment, got reliability rates around 99.997% for their text generation models. - -By comparison, human drivers are ridiculously reliable. The US has around one traffic fatality per 100 million miles driven; if a human driver makes 100 decisions per mile, that gets you a worst-case reliability of ~1:10,000,000,000 or ~99.999999999%. That's around five orders of magnitude better than a very good deep learning model, and you get that even in an open environment, where data isn't pre-filtered and there are sometimes random mechanical failures. Matching that bar is hard! I'm sure future AI will get there, but each additional "nine" of reliability is typically another unit of engineering effort. (Note that current self-driving systems use a mix of different models embedded in a larger framework, not one model trained end-to-end like GPT-3.) - -(The numbers here are only rough Fermi estimates. I'm sure one could nitpick them by going into pre-pandemic vs. post-pandemic crash rates, laws in the US vs. other countries, what percentage of crashes are drunk drivers, do drunk drivers count, how often would a really bad decision be fatal, etc. But I'm confident that whichever way you do the math, you'll still find that humans are many orders of magnitude more reliable.) - -Other types of accidents are similarly rare. Eg. pre-pandemic, there were around 40 million commercial flights per year, but only a handful of fatal crashes. If each flight involves 100 chances for the pilot to crash the plane by screwing up, then that would get you a reliability rate around 1:1,000,000,000, or ~99.99999999%. - -Even obviously dangerous activities can have very low critical failure rates. For example, shooting is a popular hobby in the US; the US market buys around 10 billion rounds of ammunition per year. There are around 500 accidental gun deaths per year, so shooting a gun has a reliability rate against accidental death of ~1:20,000,000, or 99.999995%. In a military context, the accidental death rate was around ten per year against ~1 billion rounds fired, for a reliability rate of ~99.9999999%. Deaths by fire are very rare compared to how often humans use candles, stoves, and so on; New York subway deaths are rare compared to several billion annual rides; out of hundreds of millions of hikers, only a tiny percentage fall off of cliffs; and so forth. - -The 2016 AI Impacts survey asked hundreds of AI researchers when they thought AI would be capable of doing certain tasks, playing poker, proving theorems and so on. Some tasks have been solved or have a solution "in sight", but right now, we're nowhere close to an AI that can replace human surgeons; robot-assisted surgeries still have manual control by human operators. Cosmetic surgeries on healthy patients have a fatality rate around 1:300,000, even before excluding unpredictable problems like blood clots. If a typical procedure involves two hundred chances to kill the patient by messing up, then an AI surgeon would need a reliability rate of at least 99.999998%. - -One concern with GPT-3 has been that it might accidentally be racist or offensive. Humans are, of course, sometimes racist or offensive, but in a tightly controlled Western professional context, it's pretty rare. Eg., one McDonald's employee was fired for yelling racial slurs at a customer. But McDonald's serves 70 million people a day, ~1% of the world's population. Assuming that 10% of such incidents get a news story and there's about one story per year, a similar language model would need a reliability rate of around 1:2,500,000,000, or 99.99999996%, to match McDonald's workers. When I did AI for the McDonald's drive-thru, the language model wasn't allowed to generate text at all. All spoken dialog had to be pre-approved and then manually engineered in. Reliability is hard! - -On the one hand, this might seem slightly optimistic for AI alignment research, since commercial AI teams will have to get better worst-case bounds on AI behavior for immediate economic reasons. On the other hand, because so much of the risk of AI is concentrated into a small number of very bad outcomes, it seems like such engineering might get us AIs that appear safe, and almost always are safe, but will still cause catastrophic failure in conditions that weren't anticipated. That seems bad.""" - text = """Imagine it's late autumn of 332 BC. You're Alexander the Great, and your armies are marching toward Egypt from Gaza. There’s just one little problem: you need to cross the Sinai peninsula - 150 miles of hot, barren desert. How will you carry food and water for the troops? - - -Green triangle on the left is the Nile river delta in Egypt; green chunk in the upper right is Israel. The big desert peninsula between them is the Sinai. - -Option 1: carry it - -A physically-active human needs about 3 lbs of food per day. (Modern hikers can probably find lighter calorie-dense foodstuffs, but we’re talking ancient history here.) Water requirements vary; 5 lbs is a minimum, but the US Army Quartermaster Corps recommends 20 lbs/day when marching through a hot desert. Alexander’s army crossed the desert in 7 days. Food might be reasonable, but to carry the water would mean 7*20 = 140 lbs per person, plus 50+ lbs of armor, weapons, etc. - -When I go hiking, I aim for a 20-30 lb pack. US marines are apparently expected to be able to carry 150 lbs for 9 miles - quite a bit less than the 20+ miles/day Alexander’s army managed, and with no comment on how long the marine in question might need to rest afterwards. (Also, I’m not sure I trust that source - 150 lbs for 9 miles sounds unrealistic to me, and if it’s true then I’m very impressed by marines.) - -Suffice to say that carrying that much water across that much desert is not a realistic option, even if we drink it along the way. - -Option 2: horses - -A horse consumes 20 lbs of food (half of which may be forage) and 80 lbs of water per day. In exchange, it can carry about 200 lbs (surprisingly, my source claims that horses can carry more than they can pull). Of course, that 200 lbs has to include the horse’s own food and water, plus whatever useful load it’s carrying. So, marching through a desert, a horse can only transport (200 lbs)/(80+20 lbs/day) = 2 days of supplies for itself, and that’s before whatever useful things actually need to be transported. - -In other words, there’s a hard upper limit on how far goods can be transported by horse without refilling supplies along the way. That limit is around 2 days travel time without any refill, 10 days if there’s plenty of fresh water along the route, or 20 days if there’s both water and forage. At 20 miles/day, that’s 40, 200, or 400 miles. Realistically, if we want the number of horses to be reasonable, the limit is more like half that much - 20 miles, 100 miles, or 200 miles, respectively. - -So horses also won’t work. - -Option 2.5: camels or other pack animals - -Contrary to popular image, camels actually need more water than horses. They can go a couple days without, but then need to fill up all at once. They can also carry a bit more weight, but they eat more food. At the end of the day, the numbers end up quite similar. - -Mules also end up with similar numbers, and cattle are generally worse. - -Option 3: ships - -Assuming the army marches along the coast, a supply fleet can sail alongside. At the time, a single large merchant ship could carry 400 tons - in other words, as much as about 4000 horses. Presumably the ship would cost a lot less than the horses, too. - -Well then, there’s our answer. Ships are clearly a vastly superior way to move goods. Range is a non-issue, capacity is far larger, and they’re far cheaper. They’re perfect for crossing the Sinai, which runs right along the coast anyway. - -Fast forward a few years to 327 BC, and Alexander is marching his armies back from India. He plans to cross the Gedrosian desert, along the coast of modern-day Pakistan and Iran. The plan is much like the Sinai: a supply fleet will sail alongside the army. Unfortunately, neither Alexander nor his commanders knows about the monsoons: across most of south Asia, the wind blows consistently southwest for half the year, and consistently northeast for the other half. There is nothing like it in the Mediterranean. And so, Alexander marches out expecting the fleet to catch up as soon as the winds turn - not realizing that the winds will not turn for months. Three quarters of his soldiers die in the desert. - -Thus end the campaigns of Alexander. - -Generalization -The above numbers are drawn from Donald Engels’ book Alexander the Great and the Logistics of Macedonian Army. But it tells us a lot more about the world than just the logistics of one particular ancient army. - -First, this highlights the importance of naval dominance in premodern warfare. A fleet was a far superior supply train, capable of moving a high volume of food and water over long distance at relatively low cost. Without a fleet, transport of food became expensive at best, regular resupply became a strategic necessity, and long routes through arid terrain became altogether impassable. Destroying an enemy’s fleet meant starving the army. Likewise, controlling ports wasn’t just for show - without a port, feeding the army became a serious problem. - -Another interesting insight into premodern warfare: away from friendly seas and rivers, the only way to keep an army fed was to either seize grain from enemies, or buy it from allies, either of whom needed to already be nearby. In Alexander’s case, deals were often struck to establish supply caches along the army’s intended route. - -An interesting exercise: to what extent was transportation a binding constraint on the size of premodern towns/cities? (One number you may want: Braudel (pg 121) estimates that 5000 square meters of land growing wheat would provide one person-year of food, not accounting for crop rotation.) Leave a comment if you try a calculation here; I'm curious to see how other peoples' models compare to my own. - -Modern Day -Today we have trains and trucks and roads, so the transportation constraint has relaxed somewhat. But here’s an interesting comparison: a modern 18-wheeler in the US is legally limited to haul 40 tons, while a panamax ship could carry about 50k tons through the canal (prior to the opening of the new locks in 2016). That’s a ratio of a bit over 1000 - surprisingly similar to the ship/horse ratio of antiquity, especially considering the much larger new-panamax and super-panamax ships also in use today. - - -Can we get a quick-and-dirty feel for tautness of the transportation constraint today? Here are a few very different angles: - -This USDA study shows rates on produce transport, typically about 7-20 cents per pound (see figure 6). The Smart & Final grocery store near me sells the cheaper produce items looked at in that study (bell peppers, cantaloupes, tomatoes, oranges) for 70-100 cents per pound, so transport alone is roughly 10-20% of the cost-to-consumer. -What about transporting humans? Average commute in the US is ~30 minutes each way; driving is usually in the 20-30 minute range, while public transit is usually 30-50. Assuming 8 hr workdays, that means commutes are typically ~10-20% of our work-hours. -The bureau of transportation estimates transport at 5.6% of the US economy for a very narrow measure, or 8.9% with a broader measure (though this still excludes non-market transport costs like e.g. commute time). -My interpretation: the transportation constraint becomes taut when it accounts for 10-20% of cost. If it’s less than that, it usually doesn’t limit production - we see plenty of goods which aren’t transportation-dependent or which are higher-value-per-weight, and the transportation constraint is generally slack for those. But once transportation hits about 10-20%, people start looking for alternatives, i.e. producing the goods somewhere else or using alternative goods. Obviously this is not based on very much data, but I find it intuitively plausible. - -Compared to ancient times, transportation constraints have obviously relaxed quite a lot. Yet qualitatively, the world today still does not look like a world of fully slack transportation constraints. To wrap up, let’s discuss what that would look like. - -Extreme Slackness -In Material Goods as an Abundant Resource, we discussed the world of the duplicator - a device capable of copying any item placed on it. In such a world, material scarcity is removed as an economic constraint - all material constraints are completely slack. - -What would be a corresponding sci-fi device for transportation constraints, and what would that world look like? - -I suggest portals: imagine we can create pairs of devices capable of transporting things from one device to the other, across any distance, at the speed of light. (We could instead imagine teleporters, removing the need for a pre-installed device at either end, but then the entire discussion would be about security.) What does the world of the portal look like? - -First, there’s complete geographical decoupling of production from consumption. People have no need to live near where they work; companies can put offices and factories wherever real estate is cheap. We can enjoy miles of wilderness on the back porch and a downtown district on the front porch; a swimming pool can open right into the ocean. Buying direct from the farm or factory is standard for most material goods. - -What are now tourist destinations would become options for an evening activity. Disneyworld would sell a park-hopper ticket that includes Disneyland California, Paris, and Shanghai, but the price of that ticket would be high enough to prevent the parks from becoming unpleasantly crowded - probably quite a bit more expensive than today, though possibly cheaper than today’s flights to Orlando. - -Obviously roads would cease to exist. Huge amounts of land would revert from asphalt to wilderness, but buildings would also be much more spread out. Buildings would be built close together more for show than for function - e.g. to provide the ambiance of a downtown or a community to those who want it. Physical life, in general, would look more like the structure of the internet rather than the structure of geography; “cities” would be clusters very spread out in space but very tightly connected via the portal network. Filter bubbles would be a much more physically tangible phenomenon. - -Geographically-defined governments would likely be replaced by some other form of government - governments based around access to portal hubs/networks are one natural possibility. Security would be a priority, early on - carrying an unauthorized portal into an area would earn a facefull of high explosives. On the other hand, it would be hard to prevent a high degree of mobility between areas controlled by different governments; the implications for government behavior are conceptually similar to seasteading. - -The structure of space near portal networks would be different in a big-O sense; the amount of space at a distance of about -r - would increase exponentially, rather than like -r -2 -. A nuclear warhead could go off five hundred feet away and you’d feel a breeze through a fast-branching portal network. On the other hand, viruses could spread much more rapidly. - -Anyway, at this point we’re getting into specifics of portals, so I’ll cut off the speculation. The point is: if transportation continues to get cheaper and more efficient over time, then we will converge to the world of the portal, or at least something like it. The details do matter - portals are different from teleportation or whatever might actually happen - but any method of fully relaxing transportation constraints will have qualitatively similar results, to a large extent.""" - - signature = "Title: Humans are very reliable agents, Author: alyssavance" - - splitting = TokenSplitter(max_tokens=200, min_tokens=300) - blocks = splitting.split(text, signature) - context = "Context: " + "\n\n---\n\n".join(blocks) - print(context) + return blocks \ No newline at end of file From 34bd99af3a87696f48849627239bfad2f8b9ad2a Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Mon, 17 Jul 2023 17:19:45 -0400 Subject: [PATCH 7/8] comments, refactor, new text_splitter, +bias --- src/dataset/update_dataset.py | 135 ++++++++++++++++++++++------------ 1 file changed, 90 insertions(+), 45 deletions(-) diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py index 5092947..4b77700 100644 --- a/src/dataset/update_dataset.py +++ b/src/dataset/update_dataset.py @@ -3,15 +3,18 @@ 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 -from .text_splitter import TokenSplitter +from .text_splitter import ParagraphSentenceUnitTextSplitter from .sql_db_handler import SQLDB from .pinecone_db_handler import PineconeDB -from .settings import USE_OPENAI_EMBEDDINGS, OPENAI_EMBEDDINGS_MODEL, OPENAI_EMBEDDINGS_DIMS, OPENAI_EMBEDDINGS_RATE_LIMIT, SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, SENTENCE_TRANSFORMER_EMBEDDINGS_DIMS, DEVICE, ARD_DATASET_NAME, CHUNK_SIZE, MAX_NUM_AUTHORS_IN_SIGNATURE +from .settings import USE_OPENAI_EMBEDDINGS, OPENAI_EMBEDDINGS_MODEL, \ + OPENAI_EMBEDDINGS_DIMS, OPENAI_EMBEDDINGS_RATE_LIMIT, \ + SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, SENTENCE_TRANSFORMER_EMBEDDINGS_DIMS, \ + ARD_DATASET_NAME, CHUNK_SIZE, MAX_NUM_AUTHORS_IN_SIGNATURE, \ + EMBEDDING_LENGTH_BIAS import logging logger = logging.getLogger(__name__) @@ -20,31 +23,50 @@ logger = logging.getLogger(__name__) class ARDUpdater: def __init__( 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. + min_chunk_size: int = ParagraphSentenceUnitTextSplitter.DEFAULT_MIN_CHUNK_SIZE, + max_chunk_size: int = ParagraphSentenceUnitTextSplitter.DEFAULT_MAX_CHUNK_SIZE, ): - self.token_splitter = TokenSplitter(min_tokens_per_block, max_tokens_per_block) + self.text_splitter = ParagraphSentenceUnitTextSplitter( + min_chunk_size=min_chunk_size, + max_chunk_size=max_chunk_size, + ) self.sql_db = SQLDB() self.pinecone_db = PineconeDB() if not USE_OPENAI_EMBEDDINGS: + import torch from langchain.embeddings import HuggingFaceEmbeddings + self.hf_embeddings = HuggingFaceEmbeddings( model_name=SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, - model_kwargs={'device': DEVICE}, + model_kwargs={'device': "cuda" if torch.cuda.is_available() else "cpu"}, encode_kwargs={'show_progress_bar': False} ) def update(self, custom_sources: List[str] = ['all']): + """ + Update the given sources. If no sources are provided, updates all sources. + + :param custom_sources: List of sources to update. + """ + for source in custom_sources: self.update_source(source) def update_source(self, source: str): + """ + Updates the entries from the given source. + + :param source: The name of the source to update. + """ + logger.info(f"Updating {source} entries...") streamed_dataset = load_dataset( ARD_DATASET_NAME, source, split='train', streaming=True - ).map(self.preprocess_and_validate).filter( + ).map( + self.preprocess_and_validate + ).filter( self.is_valid_entry ).filter( self.is_sql_entry_upserted @@ -54,53 +76,85 @@ class ARDUpdater: entries_batch = batch['entries_batch'] chunks_batch = batch['chunks_batch'] chunks_ids_batch = batch['chunks_ids_batch'] + sources_batch = batch['sources_batch'] try: - if USE_OPENAI_EMBEDDINGS: - embeddings = self.get_openai_embeddings(chunks_batch) - else: - embeddings = np.array(self.hf_embeddings.embed_documents(chunks_batch)) + embeddings = self.extract_embeddings(chunks_batch, sources_batch) - self.sql_db.upsert_chunks(chunks_ids_batch, chunks_batch) + self.sql_db.upsert_chunks(chunks_ids_batch, chunks_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: 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): + """ + Divides the iterable into batches of size ~CHUNK_SIZE. + + :param iterable: The iterable to divide into batches. + :returns: A generator that yields batches from the iterable. + """ + entries_batch = [] chunks_batch = [] chunks_ids_batch = [] + sources_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))] + chunks, chunks_ids = self.create_chunk_ids_and_authors(entry) - # 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) + sources_batch.extend([entry['source']] * len(chunks)) # 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} + yield self._create_batch(entries_batch, chunks_batch, chunks_ids_batch, sources_batch) entries_batch = [] chunks_batch = [] chunks_ids_batch = [] + sources_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_and_validate(self, entry): + yield self._create_batch(entries_batch, chunks_batch, chunks_ids_batch, sources_batch) + + def create_chunk_ids_and_authors(self, entry): + signature = f"Title: {entry['title']}, Author(s): {self.get_authors_str(entry['authors'])}" + chunks = self.text_splitter.split_text(entry['text']) + chunks = [f"- {signature}\n\n{chunk}" for chunk in chunks] + chunks_ids = [f"{entry['id']}_{str(i).zfill(6)}" for i in range(len(chunks))] + return chunks, chunks_ids + + def _create_batch(self, entries_batch, chunks_batch, chunks_ids_batch, sources_batch): + return {'entries_batch': entries_batch, 'chunks_batch': chunks_batch, 'chunks_ids_batch': chunks_ids_batch, 'sources_batch': sources_batch} + + 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) + + def extract_embeddings(self, chunks_batch, sources_batch): + if USE_OPENAI_EMBEDDINGS: + return self.get_openai_embeddings(chunks_batch, sources_batch) + else: + return np.array(self.hf_embeddings.embed_documents(chunks_batch, sources_batch)) + + def reset_dbs(self): + self.sql_db.create_tables(True) + self.pinecone_db.create_index(True) + + @staticmethod + def preprocess_and_validate(entry): """Preprocesses and validates the entry data""" try: - self.validate_entry(entry) + ARDUpdater.validate_entry(entry) return { 'id': entry['id'], @@ -115,7 +169,8 @@ class ARDUpdater: logger.error(f"Entry validation failed: {str(e)}", exc_info=True) return None - def validate_entry(self, entry: Dict[str, Union[str, list]], char_len_lower_limit: int = 0): + @staticmethod + def validate_entry(entry: Dict[str, Union[str, list]], char_len_lower_limit: int = 0): metadata_types = { 'id': str, 'source': str, @@ -138,36 +193,26 @@ class ARDUpdater: """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_openai_embeddings(self, chunks): + @staticmethod + def get_openai_embeddings(chunks, sources=''): embeddings = np.zeros((len(chunks), OPENAI_EMBEDDINGS_DIMS)) - rate_limit = OPENAI_EMBEDDINGS_RATE_LIMIT # TODO: use this rate_limit openai_output = openai.Embedding.create( model=OPENAI_EMBEDDINGS_MODEL, input=chunks )['data'] - for i, embedding in enumerate(openai_output): - embeddings[i] = embedding['embedding'] + for i, (embedding, source) in enumerate(zip(openai_output, sources)): + bias = EMBEDDING_LENGTH_BIAS.get(source, 1.0) + embeddings[i] = bias * np.array(embedding['embedding']) return embeddings - def reset_dbs(self): - self.sql_db.create_tables(True) - self.pinecone_db.create_index(True) - - -##### 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[: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 + @staticmethod + 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[: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 ad4f865881e42828c2f52d8d7aa68eb31254606b Mon Sep 17 00:00:00 2001 From: Henri Lemoine Date: Tue, 25 Jul 2023 17:37:05 -0400 Subject: [PATCH 8/8] moved this to the ARD repo --- src/.env.example | 3 - src/dataset/pinecone_db_handler.py | 108 -------------- src/dataset/settings.py | 37 ----- src/dataset/sql_db_handler.py | 130 ----------------- src/dataset/text_splitter.py | 128 ----------------- src/dataset/update_dataset.py | 218 ----------------------------- src/main.py | 25 ---- 7 files changed, 649 deletions(-) delete mode 100644 src/.env.example delete mode 100644 src/dataset/pinecone_db_handler.py delete mode 100644 src/dataset/settings.py delete mode 100644 src/dataset/sql_db_handler.py delete mode 100644 src/dataset/text_splitter.py delete mode 100644 src/dataset/update_dataset.py delete mode 100644 src/main.py diff --git a/src/.env.example b/src/.env.example deleted file mode 100644 index eda5459..0000000 --- a/src/.env.example +++ /dev/null @@ -1,3 +0,0 @@ -OPENAI_API_KEY="sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" -PINECONE_API_KEY="xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx" -PINECONE_ENVIRONMENT="xx-xxxxx-gcp" \ No newline at end of file diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py deleted file mode 100644 index c9dcdee..0000000 --- a/src/dataset/pinecone_db_handler.py +++ /dev/null @@ -1,108 +0,0 @@ -# dataset/pinecone_db_handler.py - -import os -import json -import pinecone - -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__) - - -class PineconeDB: - def __init__( - self, - index_name: str = PINECONE_INDEX_NAME, - values_dims: int = PINECONE_VALUES_DIMS, - metric: str = PINECONE_METRIC, - metadata_entries: list = PINECONE_METADATA_ENTRIES, - create_index: bool = False, - log_index_stats: bool = True, - ): - self.index_name = index_name - self.values_dims = values_dims - self.metric = metric - self.metadata_entries = metadata_entries - - pinecone.init( - api_key = PINECONE_API_KEY, - environment = PINECONE_ENVIRONMENT, - ) - - if create_index: - self.create_index() - - self.index = pinecone.Index(index_name=self.index_name) - - if log_index_stats: - index_stats_response = self.index.describe_index_stats() - logger.info(f"{self.index_name}:\n{index_stats_response}") - - 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 - ] - ) - ), - 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: - self.delete_index() - - pinecone.create_index( - name=self.index_name, - dimension=self.values_dims, - metric=self.metric, - metadata_config = {"indexed": self.metadata_entries}, - ) - - 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 diff --git a/src/dataset/settings.py b/src/dataset/settings.py deleted file mode 100644 index d1e4aac..0000000 --- a/src/dataset/settings.py +++ /dev/null @@ -1,37 +0,0 @@ -# dataset/settings.py - -import os -import torch -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 ### -USE_OPENAI_EMBEDDINGS = True # If false, SentenceTransformer embeddings will be used. - -OPENAI_EMBEDDINGS_MODEL = "text-embedding-ada-002" -OPENAI_EMBEDDINGS_DIMS = 1536 -OPENAI_EMBEDDINGS_RATE_LIMIT = 3500 - -SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL = "sentence-transformers/multi-qa-mpnet-base-cos-v1" -SENTENCE_TRANSFORMER_EMBEDDINGS_DIMS = 768 - -### PINECONE ### -PINECONE_INDEX_NAME = "stampy-chat-ard" -PINECONE_VALUES_DIMS = OPENAI_EMBEDDINGS_DIMS if USE_OPENAI_EMBEDDINGS else SENTENCE_TRANSFORMER_EMBEDDINGS_DIMS -PINECONE_METRIC = "dotproduct" -PINECONE_METADATA_ENTRIES = ["entry_id", "source", "title", "authors", "text"] -PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY", None) -PINECONE_ENVIRONMENT = os.environ.get("PINECONE_ENVIRONMENT", None) - -### MISCELLANEOUS ### -CHUNK_SIZE = 1750 -MAX_NUM_AUTHORS_IN_SIGNATURE = 3 -EMBEDDING_LENGTH_BIAS = { - "aisafety.info": 1.05, # In search, favor AISafety.info entries. -} \ No newline at end of file diff --git a/src/dataset/sql_db_handler.py b/src/dataset/sql_db_handler.py deleted file mode 100644 index eb91efd..0000000 --- a/src/dataset/sql_db_handler.py +++ /dev/null @@ -1,130 +0,0 @@ -# dataset/sql_db_handler.py - -from typing import List, Dict, Union -import numpy as np -import sqlite3 - -from .settings import SQL_DB_PATH - -import logging -logger = logging.getLogger(__name__) - - -class SQLDB: - def __init__(self, db_name: str = SQL_DB_PATH): - self.db_name = db_name - - self.create_tables() - - 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 ( - 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, - embedding BLOB, - entry_id TEXT, - FOREIGN KEY (entry_id) REFERENCES entry_database(id) - ) - """ - cursor.execute(query) - - except sqlite3.Error as e: - logger.error(f"The error '{e}' occurred.") - - def upsert_entry(self, entry: Dict[str, Union[str, list]]) -> 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: - logger.error(f"The error '{e}' occurred.") - return False - - finally: - conn.commit() - - def upsert_chunks(self, chunks_ids_batch: List[str], chunks_batch: List[str], embeddings_batch: List[np.ndarray]) -> bool: - with sqlite3.connect(self.db_name) as conn: - cursor = conn.cursor() - try: - for chunk_id, chunk, embedding in zip(chunks_ids_batch, chunks_batch, embeddings_batch): - cursor.execute(""" - INSERT OR REPLACE INTO chunk_database - (id, text, embedding) - VALUES (?, ?, ?) - """, (chunk_id, chunk, embedding.tobytes())) - except sqlite3.Error as e: - logger.error(f"The error '{e}' occurred.") - finally: - conn.commit() - - - def stream_chunks(self): - with sqlite3.connect(self.db_name) as conn: - cursor = conn.cursor() - - # Join entry_database and chunk_database tables and order by source - cursor.execute(""" - SELECT c.id, c.text, c.embedding, e.source - FROM chunk_database c - JOIN entry_database e ON c.entry_id = e.id - ORDER BY e.source - """) - - for row in cursor: - # Convert bytes back to numpy array - embedding = np.frombuffer(row[2], dtype=np.float64) if row[2] else None - - yield { - 'id': row[0], - 'text': row[1], - 'embedding': embedding, - 'source': row[3], - } \ No newline at end of file diff --git a/src/dataset/text_splitter.py b/src/dataset/text_splitter.py deleted file mode 100644 index db5f802..0000000 --- a/src/dataset/text_splitter.py +++ /dev/null @@ -1,128 +0,0 @@ -# dataset/text_splitter.py - -from typing import List, Callable, Any -from langchain.text_splitter import TextSplitter -from nltk.tokenize import sent_tokenize -import tiktoken - - -def char_len(string: str) -> int: - return len(string) - -def word_len(string: str) -> int: - return len(string.split(" ")) - -def token_len(string: str) -> int: - return len(tiktoken.get_encoding("cl100k_base").encode(string)) - - -def char_truncate(string: str, length: int, from_end: bool = False) -> str: - if from_end: - return string[-length:] - else: - return string[:length] - -def word_truncate(string: str, length: int, from_end: bool = False) -> str: - words = string.split(" ") - if from_end: - return " ".join(words[-length:]) - else: - return " ".join(words[:length]) - -def token_truncate(string: str, length: int, from_end: bool = False) -> str: - tokens = tiktoken.get_encoding("cl100k_base").encode(string) - if from_end: - return tiktoken.get_encoding("cl100k_base").decode(tokens[-length:]) - else: - return tiktoken.get_encoding("cl100k_base").decode(tokens[:length]) - - -class ParagraphSentenceUnitTextSplitter(TextSplitter): - """A custom TextSplitter that breaks text by paragraphs, sentences, and then units (chars/words/tokens/etc).""" - - DEFAULT_MIN_CHUNK_SIZE = 900 - DEFAULT_MAX_CHUNK_SIZE = 1100 - DEFAULT_TRUNCATE_FUNCTION = char_truncate - - def __init__( - self, - min_chunk_size: int = DEFAULT_MIN_CHUNK_SIZE, - max_chunk_size: int = DEFAULT_MAX_CHUNK_SIZE, - truncate_function: Callable[[str, int], str] = DEFAULT_TRUNCATE_FUNCTION, - **kwargs: Any - ): - super().__init__(**kwargs) - self.min_chunk_size = min_chunk_size - self.max_chunk_size = max_chunk_size - - self._truncate_function = truncate_function - - def split_text(self, text: str) -> List[str]: - blocks = [] - current_block = "" - - paragraphs = text.split("\n\n") - for paragraph in paragraphs: - current_block += "\n\n" + paragraph - block_length = self._length_function(current_block) - - if block_length > self.max_chunk_size: # current block is too large, truncate it - current_block = self._handle_large_paragraph(current_block, blocks, paragraph) - elif block_length >= self.min_chunk_size: - blocks.append(current_block) - current_block = "" - else: # current block is too small, continue appending to it - continue - - blocks = self._handle_remaining_text(current_block, blocks) - - return [block.strip() for block in blocks] - - def _handle_large_paragraph(self, current_block, blocks, paragraph): - # Undo adding the whole paragraph - current_block = current_block[:-(len(paragraph)+2)] # +2 accounts for "\n\n" - - sentences = sent_tokenize(paragraph) - for sentence in sentences: - current_block += f" {sentence}" - - block_length = self._length_function(current_block) - if block_length < self.min_chunk_size: - continue - elif block_length <= self.max_chunk_size: - blocks.append(current_block) - current_block = "" - else: - current_block = self._truncate_large_block(current_block, blocks, sentence) - - return current_block - - def _truncate_large_block(self, current_block, blocks, sentence): - while self._length_function(current_block) > self.max_chunk_size: - # Truncate current_block to max size, set remaining sentence as next sentence - truncated_block = self._truncate_function(current_block, self.max_chunk_size) - blocks.append(truncated_block) - - remaining_sentence = current_block[len(truncated_block):].lstrip() - current_block = sentence = remaining_sentence - - return current_block - - def _handle_remaining_text(self, current_block, blocks): - if blocks == []: # no blocks were added - return [current_block] - elif current_block: # any leftover text - len_current_block = self._length_function(current_block) - if len_current_block < self.min_chunk_size: - # it needs to take the last min_chunk_size-len_current_block units from the previous block - previous_block = blocks[-1] - required_units = self.min_chunk_size - len_current_block # calculate the required units - - part_prev_block = self._truncate_function(previous_block, required_units, from_end=True) # get the required units from the previous block - last_block = part_prev_block + current_block - - blocks.append(last_block) - else: - blocks.append(current_block) - - return blocks \ No newline at end of file diff --git a/src/dataset/update_dataset.py b/src/dataset/update_dataset.py deleted file mode 100644 index 4b77700..0000000 --- a/src/dataset/update_dataset.py +++ /dev/null @@ -1,218 +0,0 @@ -# dataset/update_dataset.py - -from typing import Dict, List, Union -import numpy as np -from tenacity import retry, stop_after_attempt -from datasets import load_dataset -import openai - -from .text_splitter import ParagraphSentenceUnitTextSplitter -from .sql_db_handler import SQLDB -from .pinecone_db_handler import PineconeDB - -from .settings import USE_OPENAI_EMBEDDINGS, OPENAI_EMBEDDINGS_MODEL, \ - OPENAI_EMBEDDINGS_DIMS, OPENAI_EMBEDDINGS_RATE_LIMIT, \ - SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, SENTENCE_TRANSFORMER_EMBEDDINGS_DIMS, \ - ARD_DATASET_NAME, CHUNK_SIZE, MAX_NUM_AUTHORS_IN_SIGNATURE, \ - EMBEDDING_LENGTH_BIAS - -import logging -logger = logging.getLogger(__name__) - - -class ARDUpdater: - def __init__( - self, - min_chunk_size: int = ParagraphSentenceUnitTextSplitter.DEFAULT_MIN_CHUNK_SIZE, - max_chunk_size: int = ParagraphSentenceUnitTextSplitter.DEFAULT_MAX_CHUNK_SIZE, - ): - self.text_splitter = ParagraphSentenceUnitTextSplitter( - min_chunk_size=min_chunk_size, - max_chunk_size=max_chunk_size, - ) - self.sql_db = SQLDB() - self.pinecone_db = PineconeDB() - - if not USE_OPENAI_EMBEDDINGS: - import torch - from langchain.embeddings import HuggingFaceEmbeddings - - self.hf_embeddings = HuggingFaceEmbeddings( - model_name=SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, - model_kwargs={'device': "cuda" if torch.cuda.is_available() else "cpu"}, - encode_kwargs={'show_progress_bar': False} - ) - - def update(self, custom_sources: List[str] = ['all']): - """ - Update the given sources. If no sources are provided, updates all sources. - - :param custom_sources: List of sources to update. - """ - - for source in custom_sources: - self.update_source(source) - - def update_source(self, source: str): - """ - Updates the entries from the given source. - - :param source: The name of the source to update. - """ - - logger.info(f"Updating {source} entries...") - - 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): - entries_batch = batch['entries_batch'] - chunks_batch = batch['chunks_batch'] - chunks_ids_batch = batch['chunks_ids_batch'] - sources_batch = batch['sources_batch'] - - try: - embeddings = self.extract_embeddings(chunks_batch, sources_batch) - - self.sql_db.upsert_chunks(chunks_ids_batch, chunks_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: - 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): - """ - Divides the iterable into batches of size ~CHUNK_SIZE. - - :param iterable: The iterable to divide into batches. - :returns: A generator that yields batches from the iterable. - """ - - entries_batch = [] - chunks_batch = [] - chunks_ids_batch = [] - sources_batch = [] - - for entry in iterable: - chunks, chunks_ids = self.create_chunk_ids_and_authors(entry) - - entries_batch.append(entry) - chunks_batch.extend(chunks) - chunks_ids_batch.extend(chunks_ids) - sources_batch.extend([entry['source']] * len(chunks)) - - # If this batch is large enough, yield it and start a new one. - if len(chunks_batch) >= CHUNK_SIZE: - yield self._create_batch(entries_batch, chunks_batch, chunks_ids_batch, sources_batch) - - entries_batch = [] - chunks_batch = [] - chunks_ids_batch = [] - sources_batch = [] - - # Yield any remaining items. - if entries_batch: - yield self._create_batch(entries_batch, chunks_batch, chunks_ids_batch, sources_batch) - - def create_chunk_ids_and_authors(self, entry): - signature = f"Title: {entry['title']}, Author(s): {self.get_authors_str(entry['authors'])}" - chunks = self.text_splitter.split_text(entry['text']) - chunks = [f"- {signature}\n\n{chunk}" for chunk in chunks] - chunks_ids = [f"{entry['id']}_{str(i).zfill(6)}" for i in range(len(chunks))] - return chunks, chunks_ids - - def _create_batch(self, entries_batch, chunks_batch, chunks_ids_batch, sources_batch): - return {'entries_batch': entries_batch, 'chunks_batch': chunks_batch, 'chunks_ids_batch': chunks_ids_batch, 'sources_batch': sources_batch} - - 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) - - def extract_embeddings(self, chunks_batch, sources_batch): - if USE_OPENAI_EMBEDDINGS: - return self.get_openai_embeddings(chunks_batch, sources_batch) - else: - return np.array(self.hf_embeddings.embed_documents(chunks_batch, sources_batch)) - - def reset_dbs(self): - self.sql_db.create_tables(True) - self.pinecone_db.create_index(True) - - @staticmethod - def preprocess_and_validate(entry): - """Preprocesses and validates the entry data""" - try: - ARDUpdater.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'] - } - except ValueError as e: - logger.error(f"Entry validation failed: {str(e)}", exc_info=True) - return None - - @staticmethod - def validate_entry(entry: Dict[str, Union[str, list]], char_len_lower_limit: int = 0): - metadata_types = { - 'id': str, - 'source': str, - 'title': str, - 'text': str, - 'url': str, - 'date_published': str, - 'authors': list - } - - for metadata_type, metadata_type_type in metadata_types.items(): - 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']) < 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 - - @staticmethod - def get_openai_embeddings(chunks, sources=''): - embeddings = np.zeros((len(chunks), OPENAI_EMBEDDINGS_DIMS)) - - openai_output = openai.Embedding.create( - model=OPENAI_EMBEDDINGS_MODEL, - input=chunks - )['data'] - - for i, (embedding, source) in enumerate(zip(openai_output, sources)): - bias = EMBEDDING_LENGTH_BIAS.get(source, 1.0) - embeddings[i] = bias * np.array(embedding['embedding']) - - return embeddings - - @staticmethod - 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[: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 diff --git a/src/main.py b/src/main.py deleted file mode 100644 index 5017f42..0000000 --- a/src/main.py +++ /dev/null @@ -1,25 +0,0 @@ -# main.py - -import os -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 openai -openai.api_key = os.environ['OPENAI_API_KEY'] - -import logging -logging.basicConfig(level=logging.INFO) - -from dataset.update_dataset import ARDUpdater - - -def update_sql_and_pinecone_dbs(): - updater = ARDUpdater() - updater.update() - - -if __name__ == "__main__": - update_sql_and_pinecone_dbs() \ No newline at end of file