wassname/stampy-chat
1
0
Fork 0
mirror of https://github.com/wassname/stampy-chat.git synced 2026-09-11 12:50:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into temp

Browse Source
This commit is contained in:
Fraser
2023-03-28 15:53:18 -04:00
parent 53725e5612 31a74b2882
commit 9fc0fc9072
11 changed files with 580 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+7 -2
View File
@@ -133,8 +133,13 @@ dmypy.json
*config.py
*embeddings.npy
*.pickle
*.pkl
src/data/dataset.pkl
src/tmp.py
*.DS_Store
src/*
src/semantic_search.py
src/settings.py
src/testing.ipynb
src/text_splitter.py
.vercel/
.vercel/
src/dataset.py
+81 -50
View File
@@ -8,20 +8,39 @@ import pickle
import os
import concurrent.futures
from pathlib import Path
from tenacity import (
retry,
stop_after_attempt,
wait_random_exponential,
) # for exponential backoff
import json
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
openai.api_key = os.environ.get('OPENAI_API_KEY')
try:
import config
openai.api_key = config.OPENAI_API_KEY
except ImportError:
openai.api_key = os.environ.get('OPENAI_API_KEY')
from pathlib import Path
EMBEDDING_MODEL = "text-embedding-ada-002"
COMPLETIONS_MODEL = "text-davinci-003"
LEN_EMBEDDINGS = 1536
MAX_LEN_PROMPT = 4095 # This may be 8191, unsure.
project_path = Path(__file__).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_PKL = project_path / "src" / "data" / "dataset.pkl" # Path to the saved dataset (.pkl) file, containing the dataset class object.
PATH_TO_DATASET_JSON = project_path / "src" / "data" / "dataset.json" # Path to the saved dataset (.json) file, containing the dataset class object.
# print(f"PATH_TO_DATA: {PATH_TO_DATA}")
# print(f"PATH_TO_EMBEDDINGS: {PATH_TO_EMBEDDINGS}")
# print(f"PATH_TO_DATASET_PKL: {PATH_TO_DATASET_PKL}")
# print(f"PATH_TO_DATASET_JSON: {PATH_TO_DATASET_JSON}")
error_count_dict = {
@@ -55,8 +74,8 @@ class Dataset:
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 of an article, its URL, and text. E.g.: [('title', 'url', 'text'), ...]
self.embedding_strings: List[str] = [] # List of strings, each being a few paragraphs from a single article (not exceeding 1000 words).
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}
@@ -181,7 +200,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)
@@ -198,7 +217,6 @@ class Dataset:
# 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(
@@ -217,28 +235,32 @@ class Dataset:
i, embedding = future.result()
self.embeddings[i] = embedding
num_completed += 1
if num_completed % 20 == 0:
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.")
#TODO: complete this to speed up embeddings
""" def get_embeddings_in_batches(self):
# Get an embedding for each text, with retries if necessary
# #TODO: complete this to speed up embeddings
# def get_embeddings_in_batches(self):
# # Get an embedding for each text, with retries if necessary
# batch_size = 100
@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
"""
# def get_embedding_in_batches(batch: List[str], i: int, delay_in_seconds: float = 0) -> np.ndarray:
# res = openai.Embedding.create(input=batch, engine=EMBEDDING_MODEL)
# return i, res["data"][0]["embedding"]
# start = time.time()
# self.embeddings = np.zeros((len(self.embedding_strings), LEN_EMBEDDINGS))
# with concurrent.futures.ThreadPoolExecutor() as executor:
# futures = [executor.submit(get_embedding_in_batches, batch, i) for i, batch in enumerate(self.embedding_strings)]
# 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.")
def save_embeddings(self, path: str):
np.save(path, self.embeddings)
@@ -247,11 +269,31 @@ class Dataset:
self.embeddings = np.load(path)
def save_class(self, path: str):
# 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_json(self, path: str):
# Save the class to a json file
dataset_dict = {
'metadata': self.metadata,
'embedding_strings': self.embedding_strings,
'embeddings_metadata_index': self.embeddings_metadata_index,
'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,
'sources_so_far': self.sources_so_far,
'info_types': self.info_types,
'embeddings': self.embeddings.tolist()
}
print(f"Saving class to {path}...")
with open(path, 'w') as f:
json.dump(dataset_dict, f)
def get_authors_list(authors_string: str) -> List[str]:
@@ -270,13 +312,9 @@ def get_authors_list(authors_string: str) -> List[str]:
return authors
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
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 = [
@@ -288,8 +326,8 @@ if __name__ == "__main__":
"manual",
# "arxiv",
# "https://deepmindsafetyresearch.medium.com",
"waitbutwhy.com",
"GitHub",
# "waitbutwhy.com",
# "GitHub",
# "https://aiimpacts.org",
# "arbital.com",
# "carado.moe",
@@ -300,8 +338,8 @@ if __name__ == "__main__":
# "https://intelligence.org",
# "youtube",
# "reports",
# "https://aisafety.camp",
"curriculum",
"https://aisafety.camp",
# "curriculum",
# "https://www.yudkowsky.net",
# "distill",
# "Cold Takes",
@@ -310,7 +348,6 @@ if __name__ == "__main__":
# "generative.ink",
# "greaterwrong.com"
]
dataset = Dataset(
jsonl_data_path=PATH_TO_DATA.resolve(),
@@ -319,15 +356,9 @@ if __name__ == "__main__":
min_tokens_per_block=200, max_tokens_per_block=300,
# fraction_of_articles_to_use=1/2000
)
# Test get_embedding
dataset.embedding_strings = ["This is a test", "This is another test"]
dataset.get_alignment_texts()
dataset.get_embeddings()
# dataset.get_alignment_texts()
# dataset.get_embeddings()
# dataset.save_embeddings("data/embeddings.npy")
# dataset.save_class("data/dataset.pkl")
# # dataset = pickle.load(open("dataset.pkl", "rb"))
dataset.save_json(PATH_TO_DATASET_JSON.resolve())
src/testing.ipynb
+18 -3
View File
@@ -1221,10 +1221,25 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 2,
"metadata": {},
"outputs": [],
"outputs": [
{
"ename": "RuntimeError",
"evalue": "asyncio.run() cannot be called from a running event loop",
"output_type": "error",
"traceback": [
"\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[1;31mRuntimeError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[1;32mIn[2], line 24\u001b[0m\n\u001b[0;32m 22\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39m__name__\u001b[39m \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39m__main__\u001b[39m\u001b[39m'\u001b[39m:\n\u001b[0;32m 23\u001b[0m question \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mWhat is the Natural Abstraction Hypothesis?\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m---> 24\u001b[0m asyncio\u001b[39m.\u001b[39;49mrun(test_stream(question))\n",
"File \u001b[1;32mc:\\Python310\\lib\\asyncio\\runners.py:33\u001b[0m, in \u001b[0;36mrun\u001b[1;34m(main, debug)\u001b[0m\n\u001b[0;32m 9\u001b[0m \u001b[39m\"\"\"Execute the coroutine and return the result.\u001b[39;00m\n\u001b[0;32m 10\u001b[0m \n\u001b[0;32m 11\u001b[0m \u001b[39mThis function runs the passed coroutine, taking care of\u001b[39;00m\n\u001b[1;32m (...)\u001b[0m\n\u001b[0;32m 30\u001b[0m \u001b[39m asyncio.run(main())\u001b[39;00m\n\u001b[0;32m 31\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[0;32m 32\u001b[0m \u001b[39mif\u001b[39;00m events\u001b[39m.\u001b[39m_get_running_loop() \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m---> 33\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\n\u001b[0;32m 34\u001b[0m \u001b[39m\"\u001b[39m\u001b[39masyncio.run() cannot be called from a running event loop\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m 36\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m coroutines\u001b[39m.\u001b[39miscoroutine(main):\n\u001b[0;32m 37\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\u001b[39m\"\u001b[39m\u001b[39ma coroutine was expected, got \u001b[39m\u001b[39m{!r}\u001b[39;00m\u001b[39m\"\u001b[39m\u001b[39m.\u001b[39mformat(main))\n",
"\u001b[1;31mRuntimeError\u001b[0m: asyncio.run() cannot be called from a running event loop"
]
}
],
"source": [
"import asyncio\n",
"\n",
"async def test_stream(question: str):\n",
" assistant_prompt = \"You are a helpful assistant, and you help users by answering questions and providing information about AI Alignment, on which you are extremely knowledgeable. Answer the user's question even if you are not certain of the answer; it is supremely important that you do attempt to offer an answer related to the user's query.\"\n",
" \n",
@@ -1498,7 +1513,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.7"
"version": "3.10.6"
},
"orig_nbformat": 4,
"vscode": {
web/api/data/dataset.pkl
BIN
View File
Binary file not shown.
web/api/embeddings.py
-127
View File
@@ -1,127 +0,0 @@
# ---------------------------------- web code ----------------------------------
import json
from http.server import BaseHTTPRequestHandler
class handler(BaseHTTPRequestHandler):
# post request = calculate factorial of passed number
def do_POST(self):
self.send_response(200)
self.send_header('Content-type', 'application/json')
self.end_headers()
content_length = int(self.headers['Content-Length'])
post_data = self.rfile.read(content_length)
data = json.loads(post_data)
results = {};
for i, link in enumerate(embeddings(data['query'])):
results[i] = json.dumps(link.__dict__)
self.wfile.write(json.dumps(results).encode('utf-8'))
class Link:
def __init__(self, url, title):
self.url = url
self.title = title
# -------------------------------- non-web-code --------------------------------
import pickle
from typing import List, Tuple
import os
import numpy as np # TODO: Add to requirements.txt
from tenacity import ( # TODO: Add to requirements.txt
retry,
stop_after_attempt,
wait_random_exponential,
)
import openai # TODO: Add to requirements.txt
os.environ.get('OPENAI_API_KEY')
openai.api_key = os.environ.get('OPENAI_API_KEY')
from pathlib import Path # BAD
EMBEDDING_MODEL = "text-embedding-ada-002" # BAD
COMPLETIONS_MODEL = "text-davinci-003" # BAD
LEN_EMBEDDINGS = 1536 # BAD
MAX_LEN_PROMPT = 4095 # This may be 8191, unsure. # BAD
project_path = Path(__file__).parent.parent.parent
PATH_TO_DATA = project_path / "src" / "data" / "alignment_texts.jsonl" # Path to the dataset .jsonl file. # BAD
PATH_TO_EMBEDDINGS = project_path / "src" / "data" / "embeddings.npy" # Path to the saved embeddings (.npy) file. # BAD
PATH_TO_DATASET = project_path / "src" / "data" / "dataset.pkl" # Path to the saved dataset (.pkl) file, containing the dataset class object. # BAD
@retry(wait=wait_random_exponential(min=1, max=10), stop=stop_after_attempt(4))
def get_embedding(text: str) -> np.ndarray:
result = openai.Embedding.create(model=EMBEDDING_MODEL, input=text)
return result["data"][0]["embedding"]
def get_top_k_blocks(user_query: str, k: int, HyDE: bool = False) -> List[str]:
"""Get the top k blocks that are most semantically similar to the query, using the provided dataset.
Args:
query (str): The query to be searched for.
k (int): The number of blocks to return.
HyDE (bool, optional): Whether to use HyDE or not. Defaults to False.
Returns:
List[str]: A list of the top k blocks that are most semantically similar to the query.
"""
# Get the dataset
with open(PATH_TO_DATASET, "rb") as f:
dataset = pickle.load(f)
# Get the embedding for the query.
query_embedding = get_embedding(user_query)
# If HyDE is enabled, produce a no-context ChatCompletion to the query.
if HyDE:
messages = [
{"role": "system", "content": "You are a knowledgeable AI Alignment assistant. Do your best to answer the user's question, even if you don't know the answer for sure."},
{"role": "user", "content": user_query},
]
HyDE_completion = openai.ChatCompletion.create(
model=COMPLETIONS_MODEL,
messages=messages,
temperature=0.0,
max_tokens=200
)["choices"][0]["text"]
HyDe_completion_embedding = get_embedding(f"Question: {user_query}\n\nAnswer: {HyDE_completion}")
similarity_scores = np.dot(dataset.metadataset.embeddings, HyDe_completion_embedding)
else:
similarity_scores = np.dot(dataset.metadataset.embeddings, query_embedding)
ordered_blocks = np.argsort(similarity_scores)[::-1] # Sort the blocks by similarity score
top_k_indices = ordered_blocks[:k] # Get the top k indices
top_k = [dataset.metadataset.embedding_strings[i] for i in top_k_indices] # Get the top k strings
# Get associated links
return top_k
def embeddings(query):
# write a function here that takes a query, returns a bunch of semantically similar links
return [ \
Link('https://www.lesswrong.com/posts/FinfRNLMfbq5ESxB9/microsoft-research-paper-claims-sparks-of-artificial', \
'Microsoft Research Paper Claims Sparks of Artificial Intelligence'), \
Link('https://www.lesswrong.com/posts/XhfBRM7oRcpNZwjm8/abstracts-should-be-either-actually-short-tm-or-broken-into', \
'Abstracts should be either actually short™ or broken into'), \
Link('https://www.lesswrong.com/posts/ohXcBjGvazPAxq2ex/continue-working-on-hard-alignment-don-t-give-up', \
'Continue working on hard alignment, don\'t give up'), \
Link('https://www.lesswrong.com/posts/' + query + '/this-is-a-test', \
'This is a test of ' + query), \
]
web/api/informed_assistant.py
+300
View File
@@ -0,0 +1,300 @@
# # ---------------------------------- web code ----------------------------------
# import json
# from http.server import BaseHTTPRequestHandler
# class handler(BaseHTTPRequestHandler):
# # post request = calculate factorial of passed number
# def do_POST(self):
# self.send_response(200)
# self.send_header('Content-type', 'application/json')
# self.end_headers()
# content_length = int(self.headers['Content-Length'])
# post_data = self.rfile.read(content_length)
# data = json.loads(post_data)
# results = {}
# for i, link in enumerate(informed_assistant(data['query'])):
# results[i] = json.dumps(link.__dict__)
# self.wfile.write(json.dumps(results).encode('utf-8'))
# # -------------------------------- non-web-code --------------------------------
# import time
# import os
# import openai
# import requests
# from typing import List, Dict
# import openai
# import tiktoken
# import asyncio
# import config
# from semantic_search import get_top_k_blocks
# # OpenAI API key
# try:
# import config
# OPENAI_API_KEY = config.OPENAI_API_KEY
# except ImportError:
# OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
# openai.api_key = OPENAI_API_KEY
# # OpenAI models
# EMBEDDING_MODEL = "text-embedding-ada-002"
# COMPLETIONS_MODEL = "gpt-3.5-turbo"
# # OpenAI parameters
# LEN_EMBEDDINGS = 1536
# 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.
# class Dataset:
# pass
# class Block:
# def __init__(self, title: str, author: str, date: str, url: str, tags: str, text: str):
# self.title = title
# self.author = author
# self.date = date
# self.url = url
# self.tags = tags
# self.text = text
# MODERATION_ENDPOINT = "https://api.openai.com/v1/moderations"
# def moderate_query(query: str) -> List[str]:
# """This function uses the OpenAI Moderation API to check if a query contains any offensive language.
# Args:
# query (str): The query to be checked.
# Raises:
# Exception: If the API call fails.
# Returns:
# List[str]: A list of categories that the query was flagged for.
# """
# headers = {"Content-Type": "application/json","Authorization": f"Bearer {OPENAI_API_KEY}"}
# data = {"input": query}
# response = requests.post(MODERATION_ENDPOINT, headers=headers, data=json.dumps(data))
# flagged_categories = []
# if response.status_code == 200:
# moderation_results = response.json()
# flagged = moderation_results['results'][0]['flagged']
# categories = moderation_results['results'][0]['categories']
# if flagged:
# for category, is_flagged in categories.items():
# if is_flagged:
# flagged_categories.append(category)
# else:
# raise Exception(f"Error: {response.status_code} {response.reason}")
# return flagged_categories
# def limit_tokens(text: str, max_tokens: int, encoding_name: str = "cl100k_base") -> str:
# encoding = tiktoken.get_encoding(encoding_name)
# tokens = encoding.encode(text)[:max_tokens]
# return encoding.decode(tokens)
# def generate_prompt(user_query: str, previous_dialogue: List[Dict[str, str]] = [], blocks: List[Block] = [], mode: str = "standard") -> List[Dict[str, str]]:
# """
# This function generates a prompt in messages format for the OpenAI ChatCompletions API.
# First, it picks a system description using the mode.
# Second, it adds the previous dialogue to the prompt.
# Third, it adds an instruction to the prompt based on the mode.
# Fourth, it adds the context from the top-k most relevant blocks from the Alignment Research Dataset to the prompt.
# Fifth, it adds the user query to the prompt.
# Messages take the following format:
# messages=[
# {"role": "system", "content": "You are a helpful assistant."},
# {"role": "user", "content": "Who won the world series in 2020?"},
# {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."},
# {"role": "user", "content": "Where was it played?"}
# ]
# Args:
# user_query (str): The user query.
# previous_dialogue (List[Dict[str, str]]): The previous dialogue. Defaults to [].
# blocks (List[Block]): The top-k most relevant blocks from the Alignment Research Dataset. Defaults to [].
# mode (str): The mode of the assistant. Can be "standard", etc. Defaults to "standard".
# Returns:
# List[Dict[str, str]]: The prompt in messages format.
# """
# # Initialize prompt
# prompt = []
# # Generate system description
# if mode == "standard":
# prompt.append({"role": "system", "content": "You are a helpful assistant knowledgeable about AI Alignment and Safety."})
# # elif mode == "other":
# else:
# raise Exception(f"Invalid mode: {mode}")
# # Add previous dialogue
# for message in previous_dialogue:
# prompt.append(message)
# # Add instruction
# if mode == "standard":
# instruction_prompt = "Please answer my question (after the Q:) using the provided context."
# prompt.append({"role": "assistant", "content": instruction_prompt})
# # elif mode == "other":
# else:
# raise Exception(f"Invalid mode: {mode}")
# # Add context from top-k blocks
# if blocks is None:
# return "Context missing."
# context_prompt = "Context:\n\n"
# for i, block in enumerate(blocks):
# context_prompt += f"[{i}] {block.text}\n\n"
# context_prompt = context_prompt[:-2]
# context_prompt = limit_tokens(context_prompt, 2000)
# prompt.append({"role": "user", "content": f"{context_prompt}"})
# # Add user query
# prompt.append({"role": "user", "content": f"Q: {user_query}"})
# return prompt
# def normal_completion(prompt: List[Dict[str, str]]) -> str:
# """
# This function uses the OpenAI ChatCompletions API to answer a user query.
# Args:
# messages (Dict[str, str]): A dictionary containing the system prompt and user prompt, in addition to any previous dialogue.
# Returns:
# str: The answer generated by the API.
# Raises:
# Exception: If the API call fails.
# """
# try:
# return openai.ChatCompletion.create(
# model=COMPLETIONS_MODEL,
# messages=prompt
# )["choices"][0]["message"]["content"]
# except Exception as e:
# print(e)
# return "I'm sorry, I failed to process your query. Please try again. If the problem persists, please contact the administrator."
# async def stream_completion(prompt: List[Dict[str, str]], stream_delay: float = 0.1) -> str:
# """
# This function uses the OpenAI ChatCompletions API to answer a user query, streaming the response.
# Args:
# messages (Dict[str, str]): A dictionary containing the system prompt and user prompt, in addition to any previous dialogue.
# Returns:
# str: The answer generated by the API.
# Raises:
# Exception: If the API call fails.
# """
# try:
# async for part in await openai.ChatCompletion.acreate(
# model=COMPLETIONS_MODEL,
# messages=prompt,
# stream=True
# ):
# finish_reason = part["choices"][0]["finish_reason"]
# if "content" in part["choices"][0]["delta"]:
# content = part["choices"][0]["delta"]["content"]
# yield content
# elif finish_reason:
# print(f"Stream finished: {finish_reason}")
# break
# except Exception as e:
# print(e)
# response = "I'm sorry, I failed to process your query. Please try again. If the problem persists, please contact the administrator."
# for word in response.split():
# time.sleep(stream_delay)
# yield f"{word} "
# def informed_assistant(user_query: str, previous_dialogue: List[Dict[str, str]] = [], k: str = 10, mode: str = "standard", HyDE: bool = False, stream: bool = True, stream_delay: float = 0.1) -> str:
# """
# This function uses the OpenAI ChatCompletions API to answer a user query.
# It first checks if the query is offensive, and if so, raises an exception.
# Then, it finds the top-k most relevant blocks from the Alignment Research Dataset and uses them as context for the ChatCompletions API.
# It uses the blocks to generate a prompt for the ChatCompletions API.
# Finally, it uses the ChatCompletions API to generate an answer to the user query.
# Args:
# user_query (str): The user query.
# previous_dialogue (List[Dict[str, str]]): The previous dialogue. Defaults to [].
# k (str): The number of blocks to use as context.
# mode (str): The mode to use for the ChatCompletions API. Defaults to "standard".
# HyDE (bool): Whether to use the HyDE technique for semantic search. This makes search slower, but better. Defaults to False.
# stream (bool): Whether to stream the results from the ChatCompletions API. Defaults to True.
# stream_delay (float): The delay between each word in the streamed response when streaming a hard-coded response. Defaults to 0.1.
# Returns:
# str: The answer to the user query.
# Raises:
# Exception: If the query is offensive.
# """
# # 1. Check if the query is offensive
# flagged_categories: List[str] = moderate_query(user_query)
# if len(flagged_categories) > 0:
# response = f"Your query contains offensive language. Please try again."
# if stream:
# for word in response.split():
# time.sleep(stream_delay)
# yield f"{word} "
# else:
# return response
# # 2. Find the top-k most relevant blocks from the Alignment Research Dataset
# top_k_blocks: List[Block] = get_top_k_blocks(user_query, k, HyDE)
# # 3. Generate a prompt for the ChatCompletions API
# prompt: List[Dict[str, str]] = generate_prompt(user_query, previous_dialogue, top_k_blocks, mode)
# # 4. Use the top-k most relevant blocks as context for the ChatCompletions API, and generate an answer to the user query
# if stream:
# return stream_completion(prompt)
# else:
# return normal_completion(prompt)
# if __name__ == "__main__":
# # Test the question answering function
# user_query = "Within the area of mitigating AI risk, there are several broad classes of action being taken. What does Technical safety research focus on?"
# previous_dialogue = [
# {"role": "assistant", "content": "Hi! I know all about AI Alignment. Ask me a question!"},
# ]
# k = 10
# mode = "standard"
# HyDE = True
# stream = False # Doesn't quite work yet
# print(asyncio.run(informed_assistant(user_query, previous_dialogue, k, mode, HyDE, stream)))
# # if stream:
# # for part in informed_assistant(user_query, previous_dialogue, k, mode, HyDE, stream):
# # print(part, end="")
# # else:
# # print(informed_assistant(user_query, previous_dialogue, k, mode, HyDE, stream))
web/api/requirements.txt
+1
View File
@@ -1,3 +1,4 @@
openai==0.27.2
numpy==1.24.2
tenacity==8.2.2
# aiohttp==3.8.3
web/api/semantic_search.py
+162
View File
@@ -0,0 +1,162 @@
# -------------------------------- non-web-code --------------------------------
import time
import os
import json
import numpy as np
from typing import List
import openai
from openai.error import RateLimitError
try:
import config
openai.api_key = config.OPENAI_API_KEY
except ImportError:
openai.api_key = os.environ.get('OPENAI_API_KEY')
# OpenAI models
EMBEDDING_MODEL = "text-embedding-ada-002"
COMPLETIONS_MODEL = "gpt-3.5-turbo"
# OpenAI parameters
LEN_EMBEDDINGS = 1536
MAX__TOKEN_LEN_PROMPT = 4095 # This may be 8191, unsure.
# Paths
import pathlib
project_path = pathlib.Path(__file__).parent
PATH_TO_DATASET_JSON = project_path / "data" / "dataset.json" # Path to the saved dataset (.json) file, containing the dataset class object.
class Dataset:
def __init__(self, path_to_dataset: str = PATH_TO_DATASET_JSON):
self.path_to_dataset = path_to_dataset # .json
self.load_dataset()
def load_dataset(self): # Load the dataset from the saved .json file
with open(self.path_to_dataset, 'rb') as f:
dataset_dict = json.load(f)
self.metadata = dataset_dict['metadata']
self.embedding_strings = dataset_dict['embedding_strings']
self.embeddings_metadata_index = dataset_dict['embeddings_metadata_index']
self.articles_count = dataset_dict['articles_count']
self.total_articles_count = dataset_dict['total_articles_count']
self.total_char_count = dataset_dict['total_char_count']
self.total_word_count = dataset_dict['total_word_count']
self.total_sentence_count = dataset_dict['total_sentence_count']
self.total_block_count = dataset_dict['total_block_count']
self.sources_so_far = dataset_dict['sources_so_far']
self.info_types = dataset_dict['info_types']
self.embeddings = np.array(dataset_dict['embeddings'])
class Block:
def __init__(self, title: str, author: str, date: str, url: str, tags: str, text: str):
self.title = title
self.author = author
self.date = date
self.url = url
self.tags = tags
self.text = text
def get_embedding(text: str) -> np.ndarray:
"""Get the embedding for a given text. The function will retry with exponential backoff if the API rate limit is reached, up to 4 times.
Args:
text (str): The text to get the embedding for.
Returns:
np.ndarray: The embedding for the given text.
"""
max_retries = 4
max_wait_time = 10
for attempt in range(max_retries):
try:
result = openai.Embedding.create(
model=EMBEDDING_MODEL,
input=text
)
return result["data"][0]["embedding"]
except RateLimitError as e:
if attempt + 1 == max_retries:
raise e
wait_time = min(max_wait_time, (2 ** attempt)) # Exponential backoff
time.sleep(wait_time)
def get_top_k_blocks(user_query: str, k: int = 10, HyDE: bool = False) -> List[Block]:
"""Get the top k blocks that are most semantically similar to the query, using the provided dataset.
Args:
query (str): The query to be searched for.
k (int, optional): The number of blocks to return.
HyDE (bool, optional): Whether to use HyDE or not. Defaults to False.
Returns:
List[Block]: A list of the top k blocks that are most semantically similar to the query.
"""
# Get the dataset (in data/dataset.json)
metadataset = Dataset()
# Get the embedding for the query.
query_embedding = get_embedding(user_query)
# If HyDE is enabled, produce a no-context ChatCompletion to the query.
if HyDE:
messages = [
{"role": "system", "content": "You are a knowledgeable AI Alignment assistant."},
{"role": "user", "content": f"Do your best to answer the question/instruction, even if you don't know the correct answer or action for sure.\nQ: {user_query}"},
]
HyDE_completion = openai.ChatCompletion.create(
model=COMPLETIONS_MODEL,
messages=messages
)["choices"][0]["message"]["content"]
HyDe_completion_embedding = get_embedding(f"Question: {user_query}\n\nAnswer: {HyDE_completion}")
similarity_scores = np.dot(metadataset.embeddings, HyDe_completion_embedding)
else:
similarity_scores = np.dot(metadataset.embeddings, query_embedding)
ordered_blocks = np.argsort(similarity_scores)[::-1] # Sort the blocks by similarity score
top_k_block_indices = ordered_blocks[:k] # Get the top k indices of the blocks
top_k_metadata_indexes = [metadataset.embeddings_metadata_index[i] for i in top_k_block_indices]
# Get the top k blocks (title, author, date, url, tags, text)
top_k_texts = [metadataset.embedding_strings[i] for i in top_k_block_indices] # Get the top k texts
top_k_metadata = [metadataset.metadata[i] for i in top_k_metadata_indexes] # Get the top k metadata (title, author, date, url, tags)
# Combine the top k texts and metadata into a list of Block objects
top_k_metadata_and_text = [list(top_k_metadata[i]) + [top_k_texts[i]] for i in range(k)]
blocks = [Block(*block) for block in top_k_metadata_and_text]
return blocks
# ---------------------------------- web code ----------------------------------
from http.server import BaseHTTPRequestHandler
class handler(BaseHTTPRequestHandler):
def do_POST(self):
self.send_response(200)
self.send_header('Content-type', 'application/json')
self.end_headers()
content_length = int(self.headers['Content-Length'])
post_data = self.rfile.read(content_length)
data = json.loads(post_data)
results = {}
query = data['query']
if 'k' in data:
k = data['k']
else:
k=10
if 'HyDE' in data:
HyDE = data['HyDE']
else:
HyDE = False
for i, block in enumerate(get_top_k_blocks(query, k=k, HyDE=HyDE)):
results[i] = json.dumps(block.__dict__)
self.wfile.write(json.dumps(results).encode('utf-8'))
web/package-lock.json
Generated
+1 -1
View File
@@ -8,7 +8,7 @@
"name": "alignment_search",
"version": "0.1.0",
"dependencies": {
"next": "^13.2.1",
"next": "^13.2.4",
"react": "18.2.0",
"react-dom": "18.2.0",
"zod": "^3.20.6"
web/package.json
+1 -1
View File
@@ -9,7 +9,7 @@
"start": "next start"
},
"dependencies": {
"next": "^13.2.1",
"next": "^13.2.4",
"react": "18.2.0",
"react-dom": "18.2.0",
"zod": "^3.20.6"
web/src/pages/index.tsx
+9 -8
View File
@@ -42,14 +42,14 @@ const Home: NextPage = () => {
const SearchBox: React.FC = () => {
const [query, setQuery] = useState("");
const [results, setResults] = useState<{title: string, url: string}[]>([]);
const [results, setResults] = useState<{title: string, author: string, date: string, url: string, tags: string, text: string}[]>([]);
const [loading, setLoading] = useState(false);
const embeddings = async (query: String) => {
const semantic_search = async (query: String) => {
setLoading(true);
const res = await fetch("/api/embeddings", {
const res = await fetch("/api/semantic_search", {
method: "POST",
headers: { "Content-Type": "application/json", },
body: JSON.stringify({query: query}),
@@ -60,18 +60,19 @@ const SearchBox: React.FC = () => {
setLoading(false);
// data looks like
// { 0: "{'url' : 'https://foo.com', 'title' : 'foo'}",
// 1: "{'url' : 'https://bar.com', 'title' : 'bar'}" }
// { 0: "{'title': 'First Title', 'author': 'Bob Miles', 'date': 'March 1st, 2023', 'url': 'https://example.com', 'tags': ['tag1', 'tag2'], 'text': 'This is the content of the article'}",
// 1: "{'title': 'Second Title', 'author': 'Frank Ocean', 'date': 'March 6th, 2023', 'url': 'https://ai.com', 'tags': ['tag3', 'tag4'], 'text': 'This is the content of the article'}"
// }
// so we need to convert it to a list of objects
return Object.keys(data).map((key) => JSON.parse(data[key])) || [{title: "error", url: "error"}];
return Object.keys(data).map((key) => JSON.parse(data[key])) || [{title: "error", author: "error", date: "error", url: "error", tags: ["error"], text: "error"}];
};
return (
<>
<form className="flex mb-2" onSubmit={async (e) => { // store in a form so that <enter> submits
e.preventDefault();
setResults(await embeddings(query));
setResults(await semantic_search(query));
}}>
<input
@@ -89,7 +90,7 @@ const SearchBox: React.FC = () => {
<ul>
{results.map((result) => (
<li key={result.url} className="my-1">
<a href={result.url}>{result.title}</a>
<a href={result.url}>{result.text}</a>
</li>
))}
</ul>