diff --git a/.gitignore b/.gitignore index c557758..481e4c1 100644 --- a/.gitignore +++ b/.gitignore @@ -138,3 +138,4 @@ src/tmp.py *.DS_Store .vercel/ + diff --git a/package-lock.json b/package-lock.json new file mode 100644 index 0000000..cbb41ab --- /dev/null +++ b/package-lock.json @@ -0,0 +1,6 @@ +{ + "name": "AlignmentSearch", + "lockfileVersion": 3, + "requires": true, + "packages": {} +} diff --git a/requirements.txt b/requirements.txt index a9a238e..f2bd0d9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,4 +5,5 @@ langchain requests tiktoken tqdm -nltk \ No newline at end of file +nltk +dateutil \ No newline at end of file diff --git a/src/semantic_search.py b/src/assistant/semantic_search.py similarity index 99% rename from src/semantic_search.py rename to src/assistant/semantic_search.py index 151eb55..98829d5 100644 --- a/src/semantic_search.py +++ b/src/assistant/semantic_search.py @@ -10,7 +10,7 @@ from tenacity import ( import tiktoken import config -from dataset import Dataset +from dataset.create_dataset import Dataset from settings import PATH_TO_DATASET, EMBEDDING_MODEL, COMPLETIONS_MODEL openai.api_key = config.OPENAI_API_KEY @@ -120,6 +120,7 @@ class AlignmentSearch: return explanation_prompt def get_user_prompt(self, user_query: str, mode: str) -> str: + pass #TODO: fix junk def create_messages(self, system_prompt: str, user_prompt: str, context: str, question: str): return [ diff --git a/src/dataset/__init__.py b/src/dataset/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/dataset.py b/src/dataset/create_dataset.py similarity index 82% rename from src/dataset.py rename to src/dataset/create_dataset.py index acfdbe3..e0fbfd9 100644 --- a/src/dataset.py +++ b/src/dataset/create_dataset.py @@ -8,6 +8,8 @@ import pickle import os import concurrent.futures from pathlib import Path +from tqdm.auto import tqdm +from dateutil.parser import parse, ParserError from tenacity import ( retry, @@ -15,10 +17,6 @@ from tenacity import ( wait_random_exponential, ) # for exponential backoff -from text_splitter import TokenSplitter, split_into_sentences -from settings import PATH_TO_DATA, PATH_TO_EMBEDDINGS, PATH_TO_DATASET, EMBEDDING_MODEL, LEN_EMBEDDINGS -import os -from tqdm.auto import tqdm import openai try: @@ -28,6 +26,12 @@ except ImportError: openai.api_key = os.environ.get('OPENAI_API_KEY') +from .settings import PATH_TO_RAW_DATA, PATH_TO_DATASET, 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, @@ -43,11 +47,11 @@ class MissingDataException(Exception): class Dataset: def __init__(self, - jsonl_data_path: str, # Path to the dataset .jsonl file. + 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 = 400, # Minimum number of tokens per block. - max_tokens_per_block: int = 600, # Maximum number of tokens per block. + 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 @@ -115,6 +119,8 @@ class Dataset: 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'] @@ -163,12 +169,17 @@ class Dataset: if (self.custom_sources is not None) and (entry['source'] not in self.custom_sources): continue - self.articles_count[entry['source']] += 1 - self.total_articles_count += 1 - # Get title, author, date, URL, tags, and text title, author, date_published, url, tags, text = self.extract_info_from_article(entry) - + + #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}, " @@ -185,7 +196,7 @@ class Dataset: 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] * len(blocks)) + self.embeddings_metadata_index.extend([self.total_articles_count-1] * len(blocks)) # Update counts self.total_char_count += len(text) @@ -199,50 +210,41 @@ class Dataset: error_count_dict[str(e)] += 1 def get_embeddings(self): - # Get an embedding for each text, with retries if necessary - #TODO: check batch size stuff at https://github.com/openai/openai-cookbook/blob/main/examples/vector_databases/pinecone/Gen_QA.ipynb - # to speed up the process - # @retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(5)) - def get_embedding_at_index(text: str, i: int, delay_in_seconds: float = 0) -> np.ndarray: - time.sleep(delay_in_seconds) - embedding = openai.Embedding.create( + 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=text - ) - return i, embedding["data"][0]["embedding"] - + input=texts + )['data'] + for i, embedding in enumerate(openai_output): + embeddings[i] = embedding['embedding'] + return batch_idx, embeddings + + batch_size = 200 + 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_embedding_at_index, text, i) for i, text in enumerate(self.embedding_strings)] + 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): - i, embedding = future.result() - self.embeddings[i] = embedding - num_completed += 1 - if num_completed % 50 == 0: - print(f"Completed {num_completed}/{len(self.embedding_strings)} embeddings in {time.time() - start:.2f} seconds.") - print(f"Completed {num_completed}/{len(self.embedding_strings)} embeddings in {time.time() - start:.2f} seconds.") + batch_idx, embeddings = future.result() + num_completed += embeddings.shape[0] + self.embeddings[batch_idx:batch_idx+embeddings.shape[0]] = embeddings - #TODO: complete this to speed up embeddings - """ def get_embeddings_in_batches(self): - # Get an embedding for each text, with retries if necessary - - @retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(5)) - def get_embedding_in_batches(batch: List[str], i: int, delay_in_seconds: float = 0) -> np.ndarray: - try: - res = openai.Embedding.create(input=batch, engine=EMBEDDING_MODEL) - except: - done = False - while not done: - time.sleep(5) - try: - res = openai.Embedding.create(input=batch, engine=EMBEDDING_MODEL) - done = True - except: - pass - """ + 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) @@ -250,7 +252,7 @@ class Dataset: def load_embeddings(self, path: str): self.embeddings = np.load(path) - def save_class(self, path: str): + def save_class(self, path: str = PATH_TO_DATASET): # Save the class to a pickle file print(f"Saving class to {path}...") with open(path, 'wb') as f: @@ -272,7 +274,17 @@ def get_authors_list(authors_string: str) -> List[str]: 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 @@ -311,7 +323,7 @@ if __name__ == "__main__": ] dataset = Dataset( - jsonl_data_path=PATH_TO_DATA.resolve(), + 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, @@ -323,5 +335,5 @@ if __name__ == "__main__": 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/data/dataset.pkl b/src/dataset/data/dataset.pkl new file mode 100644 index 0000000..405c7c0 Binary files /dev/null and b/src/dataset/data/dataset.pkl differ diff --git a/src/dataset/settings.py b/src/dataset/settings.py new file mode 100644 index 0000000..611493e --- /dev/null +++ b/src/dataset/settings.py @@ -0,0 +1,24 @@ +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. + + +def get_rawdata_file_path(): + current_file_path = Path(__file__).resolve() + data_file_path = current_file_path.parent / 'data' / 'alignment_texts.jsonl' + return str(data_file_path) + +def get_dataset_file_path(): + current_file_path = Path(__file__).resolve() + data_file_path = current_file_path.parent / 'data' / 'dataset.pkl' + return str(data_file_path) + + +PATH_TO_RAW_DATA = get_rawdata_file_path() + +PATH_TO_DATASET = get_dataset_file_path() + diff --git a/src/text_splitter.py b/src/dataset/text_splitter.py similarity index 98% rename from src/text_splitter.py rename to src/dataset/text_splitter.py index 67f130b..d4229ae 100644 --- a/src/text_splitter.py +++ b/src/dataset/text_splitter.py @@ -7,9 +7,13 @@ from typing import List import nltk -# Download the Punkt tokenizer if you haven't already +# 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 split_into_sentences(text: str) -> List[str]: """ Splits the input text into sentences. @@ -87,7 +91,7 @@ class TokenSplitter: last_block = dec(enc(latest_plus_current)[-max_tokens:]) blocks.append(last_block) - return blocks + return [block.strip() for block in blocks] def split(self, text: str, signature: str = None) -> List[str]: if signature is None: diff --git a/src/main.py b/src/main.py index 0f88830..4024626 100644 --- a/src/main.py +++ b/src/main.py @@ -1,21 +1,158 @@ -import numpy as np -import pickle import openai - +""" import config -from semantic_search import AlignmentSearch -from settings import PATH_TO_DATASET +from assistant.semantic_search import AlignmentSearch +from dataset.create_dataset import Dataset openai.api_key = config.OPENAI_API_KEY +from settings import PATH_TO_RAW_DATA, PATH_TO_DATASET, EMBEDDING_MODEL, LEN_EMBEDDINGS +""" +from tenacity import ( + retry, + stop_after_attempt, + wait_random_exponential, +) -def main(): - with open(PATH_TO_DATASET, 'rb') as f: - dataset = pickle.load(f) +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 PATH_TO_DATASET, EMBEDDING_MODEL + + +import numpy as np +import matplotlib.pyplot as plt + + +def load_rawdata_into_pkl(): + """with open(PATH_TO_DATASET, 'rb') as f: + dataset = pickle.load(f) AS = AlignmentSearch(dataset=dataset) prompt = "What would be an idea to solve the Alignment Problem? Name the Lesswrong post by Quintin Pope that discusses this idea." answer = AS.search_and_answer(prompt, 3, HyDE=False) print(answer) + """ + # List of possible sources: + all_sources = ["https://aipulse.org", "ebook", "https://qualiacomputing.com", "alignment forum", "lesswrong", "manual", "arxiv", "https://deepmindsafetyresearch.medium.com", "waitbutwhy.com", "GitHub", "https://aiimpacts.org", "arbital.com", "carado.moe", "nonarxiv_papers", "https://vkrakovna.wordpress.com", "https://jsteinhardt.wordpress.com", "audio-transcripts", "https://intelligence.org", "youtube", "reports", "https://aisafety.camp", "curriculum", "https://www.yudkowsky.net", "distill", "Cold Takes", "printouts", "gwern.net", "generative.ink", "greaterwrong.com"] # These 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 = 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/2000 + ) + dataset.get_alignment_texts() + print(len(dataset.embedding_strings)) + dataset.get_embeddings() + # dataset.save_embeddings("data/embeddings.npy") + + dataset.save_class() + # # dataset = pickle.load(open("dataset.pkl", "rb")) + +@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, 'rb') as f: + dataset = pickle.load(f) + + embeddings_len = len(dataset.embedding_strings) + i1 = random.randint(0,embeddings_len-1) + #i2 = random.randint(0,embeddings_len-1) + #print(len(dataset.embeddings)) + #print(len(dataset.embedding_strings)) + #embedding_test = get_embedding(dataset.embedding_strings[i]) + #print(np.dot(embedding_test,dataset.embeddings[i])) + + metadata_i1 = dataset.embeddings_metadata_index[i1] + print("metadata:",dataset.metadata[metadata_i1]) + print("embedding_string:",dataset.embedding_strings[i1]) + #print("embedding_vector:",dataset.embeddings[i1]) + embedding_of_string1 = get_embedding(dataset.embedding_strings[i1]) + + #metadata_i2 = dataset.embeddings_metadata_index[i2] + #print("metadata:",dataset.metadata[metadata_i2]) + #print("embedding_string:",dataset.embedding_strings[i2]) + #print("embedding_vector:",dataset.embeddings[i1]) + #embedding_of_string2 = get_embedding(dataset.embedding_strings[i2]) + #embedding_of_string2 = get_embedding("000000000000000000000000000000000000000000000000000000000000000000000000000000000") + + #print(len(embedding_of_string1)) + vector = dataset.embeddings[i1] + plot_likelihood(vector) + #plot_likelihood(get_embedding("tst")) + print(max(vector), min(vector)) + print(sum([x**2 for x in vector])) + + + + + #print(np.dot(embedding_of_string1, embedding_of_string2)) + +def plot_likelihood(embeddings, num_buckets=200): + # Calculate the histogram + histogram, bin_edges = np.histogram(embeddings.flatten(), bins=num_buckets, range=(embeddings.min(), embeddings.max())) + + # Normalize the histogram to get likelihoods + likelihoods = histogram / embeddings.flatten().size + + # Plot the likelihoods + plt.bar(bin_edges[:-1], likelihoods, width=(bin_edges[1] - bin_edges[0]), edgecolor="k", alpha=0.7) + plt.xlabel("Value") + plt.ylabel("Likelihood") + plt.title("Likelihood of Floats in the Vector Embedding") + plt.savefig("bla.png") + + + + + + if __name__ == "__main__": - main() \ No newline at end of file + #load_rawdata_into_pkl() + print_out_dataset_stuff() \ No newline at end of file diff --git a/src/settings.py b/src/settings.py index 571faff..addaf8b 100644 --- a/src/settings.py +++ b/src/settings.py @@ -1,12 +1,24 @@ from pathlib import Path EMBEDDING_MODEL = "text-embedding-ada-002" -COMPLETIONS_MODEL = "text-davinci-003" +COMPLETIONS_MODEL = "gpt-3.5-turbo" LEN_EMBEDDINGS = 1536 MAX_LEN_PROMPT = 4095 # This may be 8191, unsure. -project_path = Path(__file__).parent.parent#.parent -PATH_TO_DATA = project_path / "src" / "data" / "alignment_texts.jsonl" # Path to the dataset .jsonl file. -PATH_TO_EMBEDDINGS = project_path / "src" / "data" / "embeddings.npy" # Path to the saved embeddings (.npy) file. -PATH_TO_DATASET = project_path / "src" / "data" / "dataset.pkl" # Path to the saved dataset (.pkl) file, containing the dataset class object. + +def get_rawdata_file_path(): + current_file_path = Path(__file__).resolve() + data_file_path = current_file_path.parent / 'dataset' / 'data' / 'alignment_texts.jsonl' + return str(data_file_path) + +def get_dataset_file_path(): + current_file_path = Path(__file__).resolve() + data_file_path = current_file_path.parent / 'dataset' / 'data' / 'dataset.pkl' + return str(data_file_path) + + +PATH_TO_RAW_DATA = get_rawdata_file_path() + +PATH_TO_DATASET = get_dataset_file_path() + diff --git a/web/api/informed_assistant.py b/web/api/informed_assistant.py index df1e49c..09b011c 100644 --- a/web/api/informed_assistant.py +++ b/web/api/informed_assistant.py @@ -35,7 +35,7 @@ import tiktoken import asyncio import config -from semantic_search import get_top_k_blocks +from assistant.semantic_search import get_top_k_blocks # OpenAI API key diff --git a/web/api/semantic_search.py b/web/api/semantic_search.py index e740f8c..17d1ca1 100644 --- a/web/api/semantic_search.py +++ b/web/api/semantic_search.py @@ -53,10 +53,20 @@ MAX_LEN_PROMPT = 4095 # This may be 8191, unsure. # Paths from pathlib import Path -project_path = Path(__file__).parent.parent.parent -PATH_TO_DATA = project_path / "web" / "api" / "data" / "alignment_texts.jsonl" # Path to the dataset .jsonl file. -PATH_TO_EMBEDDINGS = project_path / "web" / "api" / "data" / "embeddings.npy" # Path to the saved embeddings (.npy) file. -PATH_TO_DATASET = project_path / "web" / "api" / "data" / "dataset.pkl" # Path to the saved dataset (.pkl) file, containing the dataset class object. +#current_file = Path(__file__).parent.parent.parent +#PATH_TO_DATA = project_path / "web" / "api" / "data" / "alignment_texts.jsonl" # Path to the dataset .jsonl file. +#PATH_TO_EMBEDDINGS = project_path / "web" / "api" / "data" / "embeddings.npy" # Path to the saved embeddings (.npy) file. +#PATH_TO_DATASET = project_path / "web" / "api" / "data" / "dataset.pkl" # Path to the saved dataset (.pkl) file, containing the dataset class object. + + +# Get the path from the environment variable +PATH_TO_DATASET = os.environ.get("PATH_TO_DATASET") + +# Fallback to the local path if the environment variable is not set +if PATH_TO_DATASET is None: + PATH_TO_DATASET = Path(__file__).parent / "data" / "dataset.pkl" +else: + PATH_TO_DATASET = Path(PATH_TO_DATASET) class Dataset: