chunk unification

This commit is contained in:
Fraser
2023-03-29 00:15:27 -04:00
parent 07b85d6021
commit 72378efdd7
4 changed files with 54 additions and 185 deletions
+9 -49
View File
@@ -17,22 +17,15 @@ class handler(BaseHTTPRequestHandler):
self.wfile.write(chat(data['query'], data['history']).encode('utf-8'))
# ------------------------------- chat gpt stuff -------------------------------
# --------------------------------- chat stuff ---------------------------------
from api.get_blocks import get_top_k_blocks, Block
from typing import List, Dict
import openai
import os
import requests
from typing import List, Dict
try:
import tiktoken
except ImportError as e:
print(e)
print("Please install tiktoken with `pip install tiktoken`")
import openai
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
openai.api_key = OPENAI_API_KEY
from api.get_blocks import get_top_k_blocks, Block
import tiktoken
# OpenAI models
EMBEDDING_MODEL = "text-embedding-ada-002"
@@ -44,39 +37,10 @@ LEN_EMBEDDINGS = 1536
MAX_TOKEN_LEN_PROMPT = 4095 # This may be 8191, unsure.
TRUNCATE_CONTEXT = 2000
def moderate_query(query: str) -> List[str]:
"""This function uses the OpenAI Moderation API to check if a query contains any offensive language.
# OpenAI API key
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
openai.api_key = OPENAI_API_KEY
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)
@@ -197,10 +161,6 @@ def chat(query: str, history: List[Dict[str, str]] = [], k: str = 10, mode: str
"""
# 1. Check if the query is offensive
flagged_categories: List[str] = moderate_query(query)
if len(flagged_categories) > 0:
return f"Your query contains offensive language. Please try again."
# 2. Find the top-k most relevant blocks from the Alignment Research Dataset
top_k_blocks: List[Block] = get_top_k_blocks(query, k, HyDE)
+37 -11
View File
@@ -1,8 +1,10 @@
import time
import numpy as np
import json
from typing import List
import dataclasses
import itertools
import json
import numpy as np
import openai
import time
EMBEDDING_MODEL = "text-embedding-ada-002"
COMPLETIONS_MODEL = "gpt-3.5-turbo"
@@ -32,14 +34,14 @@ class Dataset:
self.info_types = dataset_dict['info_types']
self.embeddings = np.array(dataset_dict['embeddings'])
@dataclasses.dataclass
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
title: str
author: str
date: str
url: str
tags: str
text: str
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.
@@ -111,4 +113,28 @@ def get_top_k_blocks(user_query: str, k: int = 10, HyDE: bool = False) -> List[B
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
return unify(blocks)
# for all blocks that are "the same" (same title, author, date, url, tags),
# combine their text with "\n\n...\n\n" in between, returning the list.
def unify(blocks: List[Block]) -> List[Block]:
key = lambda block: (block.title, block.author, block.date, block.url, block.tags)
blocks.sort(key=key)
unified_blocks: List[Block] = []
for k, g in itertools.groupby(blocks, key=key):
text = "\n\n\n[...]\n\n\n".join([block.text for block in g])
unified_blocks.append(Block(k[0], k[1], k[2], k[3], k[4], text))
return unified_blocks
+1 -123
View File
@@ -2,6 +2,7 @@
import json
import dataclasses
from api.get_blocks import get_top_k_blocks
from http.server import BaseHTTPRequestHandler
@dataclasses.dataclass
@@ -30,127 +31,4 @@ class handler(BaseHTTPRequestHandler):
data = json.loads(post_data)
self.wfile.write(json.dumps(get_top_k_blocks(data['query']), cls = Encoder).encode('utf-8'))
# -------------------------------- 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'])
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
+7 -2
View File
@@ -40,8 +40,13 @@ const ShowSemanticEntry: React.FC<{entry: SemanticEntry}> = ({entry}) => {
<h3 className="text-xl flex-1">{entry.title}</h3>
<p className="flex-1 text-right my-0">{entry.author} - {entry.date}</p>
</div>
<p className="text-sm">{entry.text}</p>
{ entry.text.split("\n").map((paragraph, i) => {
const p = paragraph.trim();
if (p === "") return <></>;
if (p === "[...]") return <hr key={i} />;
return <p className="text-sm" key={i}> {paragraph} </p>
})
}
<a href={entry.url}>Read more</a>
</div>