diff --git a/.gitignore b/.gitignore index 293af7e..5bfc436 100644 --- a/.gitignore +++ b/.gitignore @@ -129,6 +129,8 @@ dmypy.json .pyre/ # Other +*test.py +*test.ipynb *alignment_texts.jsonl *config.py *.DS_Store @@ -145,3 +147,8 @@ 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 diff --git a/src/.env.example b/src/.env.example new file mode 100644 index 0000000..eda5459 --- /dev/null +++ b/src/.env.example @@ -0,0 +1,3 @@ +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/create_dataset.py b/src/dataset/create_dataset.py deleted file mode 100644 index 2456e44..0000000 --- a/src/dataset/create_dataset.py +++ /dev/null @@ -1,401 +0,0 @@ -import jsonlines -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 - -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_RAW_DATA, 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 source.": 0, - "Entry has no title.": 0, - "Entry has no text.": 0, - "Entry has no URL.": 0, - "Entry has wrong citation level.": 0 -} - - -class MissingDataException(Exception): - pass - - -class Dataset: - 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 - 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.articles_count: DefaultDict[str, int] = defaultdict(int) # Number of articles per source. E.g.: {'source1': 10, 'source2': 20, 'total': 30} - - if self.custom_sources is not None: - for source in self.custom_sources: - self.articles_count[source] = 0 - self.total_articles_count = 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. - - 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] - - # 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] - - # 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 - - # 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 - - #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 - - 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() - self.embeddings = np.zeros((len(self.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], - batch_idx, - len(self.embedding_strings[batch_idx:batch_idx+batch_size]) - ) for batch_idx in range(0, len(self.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(self.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 - } - 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(",")] - else: - authors = [authors_string.strip()] - return authors - -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 diff --git a/src/dataset/pinecone_db_handler.py b/src/dataset/pinecone_db_handler.py new file mode 100644 index 0000000..b0090d3 --- /dev/null +++ b/src/dataset/pinecone_db_handler.py @@ -0,0 +1,100 @@ +# 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, + create_index: bool = False, + ): + self.index_name = PINECONE_INDEX_NAME + + 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) + + 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)}" + + 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=PINECONE_VALUES_DIMS, + metric=PINECONE_METRIC, + metadata_config = {"indexed": PINECONE_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 index 28cb3da..951e138 100644 --- a/src/dataset/settings.py +++ b/src/dataset/settings.py @@ -1,12 +1,31 @@ +# dataset/settings.py + +import os +import torch 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. - +### FILE PATHS ### 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 +SQL_DB_PATH = str(current_file_path.parent / 'data' / 'ARD.db') + +### DATASET ### +ARD_DATASET_NAME = "StampyAI/alignment-research-dataset" + +### EMBEDDINGS ### +USE_OPENAI_EMBEDDINGS = False +OPENAI_EMBEDDINGS_MODEL = "text-embedding-ada-002" +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" +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 diff --git a/src/dataset/sql_db_handler.py b/src/dataset/sql_db_handler.py new file mode 100644 index 0000000..747d7b4 --- /dev/null +++ b/src/dataset/sql_db_handler.py @@ -0,0 +1,104 @@ +# dataset/sql_db_handler.py + +from typing import List, Dict, Union +import sqlite3 + +from .settings import SQL_DB_PATH + +import logging +logger = logging.getLogger(__name__) + + +class SQLDB: + def __init__(self): + self.db_name = SQL_DB_PATH + + 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, + 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]) -> bool: + with sqlite3.connect(self.db_name) as conn: + cursor = conn.cursor() + try: + 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.") + finally: + conn.commit() 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/dataset/update_dataset.py b/src/dataset/update_dataset.py new file mode 100644 index 0000000..f07b6c1 --- /dev/null +++ b/src/dataset/update_dataset.py @@ -0,0 +1,173 @@ +# dataset/update_dataset.py + +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 .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 + +import logging +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. + ): + 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, + model_kwargs={'device': DEVICE}, + encode_kwargs={'show_progress_bar': True} + ) + + 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): + 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, chunk_size): + entries_batch = batch['entries_batch'] + chunks_batch = batch['chunks_batch'] + chunks_ids_batch = batch['chunks_ids_batch'] + + try: + 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) + + 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_and_validate(self, entry): + """Preprocesses and validates the entry data""" + 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'] + } + except ValueError as e: + 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): + 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 + + 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): + embeddings = np.zeros((len(chunks), 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'] + + 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 diff --git a/src/main.py b/src/main.py index 900a14f..25772dc 100644 --- a/src/main.py +++ b/src/main.py @@ -1,163 +1,28 @@ +# 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 -""" -import config -from assistant.semantic_search import AlignmentSearch -from dataset.create_dataset import Dataset +openai.api_key = os.environ['OPENAI_API_KEY'] -openai.api_key = config.OPENAI_API_KEY +import logging +logging.basicConfig(level=logging.INFO) -from settings import PATH_TO_RAW_DATA, PATH_TO_DATASET, EMBEDDING_MODEL, LEN_EMBEDDINGS -""" -from tenacity import ( - retry, - stop_after_attempt, - wait_random_exponential, -) - -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 +from dataset.update_dataset import ARDUpdater -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 - - # 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.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, +def update_sql_and_pinecone_dbs(): + updater = ARDUpdater( + 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() - -@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") - - - - + updater.update(['gwern_blog']) 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 + update_sql_and_pinecone_dbs() \ No newline at end of file diff --git a/src/settings.py b/src/settings.py deleted file mode 100644 index 2362cc7..0000000 --- a/src/settings.py +++ /dev/null @@ -1,12 +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_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