mirror of
https://github.com/wassname/stampy-chat.git
synced 2026-09-11 12:50:34 +08:00
Merge pull request #39 from StampyAI/update-sql-and-pinecone
Update sql and pinecone
This commit is contained in:
Henri Lemoine
@@ -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/*
|
||||
@@ -0,0 +1,3 @@
|
||||
OPENAI_API_KEY="sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
|
||||
PINECONE_API_KEY="xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
|
||||
PINECONE_ENVIRONMENT="xx-xxxxx-gcp"
|
||||
@@ -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"))
|
||||
"""
|
||||
|
||||
@@ -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)
|
||||
+28
-9
@@ -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')
|
||||
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
|
||||
@@ -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()
|
||||
@@ -1,3 +1,5 @@
|
||||
# dataset/text_splitter.py
|
||||
|
||||
import re
|
||||
from typing import List
|
||||
import tiktoken
|
||||
|
||||
@@ -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
|
||||
+19
-154
@@ -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)
|
||||
update_sql_and_pinecone_dbs()
|
||||
@@ -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')
|
||||
Reference in New Issue
Block a user