initial commit based on McGill's AlignmentSearch

.gitignore

@@ -0,0 +1,145 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Other
+*alignment_texts.jsonl
+*config.py
+*.DS_Store
+*.env
+*.npy
+*.pkl
+
+.vercel/
+temp/
+*tmp.py
+
+api/dataset.pkl
+api/dataset_big.pkl
+api/dataset_300.pkl

README.md

@@ -0,0 +1,11 @@
+# AlignmentSearch
+
+![img](resources/screenshot.png)
+
+## Project Layout
+
+- `src/` stuff around processing our dataset and constructing embeddings
+- `api/` a flask app serving as our backend
+- `web/` a NextJS app for the frontend
+
+

api/Procfile

@@ -0,0 +1 @@
+web: gunicorn main:app --worker-class eventlet --threads 4

api/chat.py

@@ -0,0 +1,175 @@
+# ------------------------------- env, constants -------------------------------
+
+from get_blocks import get_top_k_blocks, Block
+
+from typing import List, Dict
+import openai
+import tiktoken
+import time
+import re
+
+# OpenAI models
+EMBEDDING_MODEL = "text-embedding-ada-002"
+COMPLETIONS_MODEL = "gpt-3.5-turbo"
+# COMPLETIONS_MODEL = "gpt-4"
+
+STANDARD_K = 20 if COMPLETIONS_MODEL == 'gpt-4' else 10
+
+# parameters
+
+# NOTE: All this is approximate, there's bits I'm intentionally not counting. Leave a buffer beyond what you might expect.
+NUM_TOKENS = 8191 if COMPLETIONS_MODEL == 'gpt-4' else 4095
+HISTORY_FRACTION = 0.25 # the (approximate) fraction of num_tokens to use for history text before truncating
+CONTEXT_FRACTION = 0.5  # the (approximate) fraction of num_tokens to use for context text before truncating
+
+ENCODER = tiktoken.get_encoding("cl100k_base")
+
+DEBUG_PRINT = True
+
+# --------------------------------- prompt code --------------------------------
+
+
+
+# limit a string to a certain number of tokens
+def cap(text: str, max_tokens: int) -> str:
+
+    if max_tokens <= 0: return "..."
+
+    encoded_text = ENCODER.encode(text)
+
+    if len(encoded_text) <= max_tokens: return text
+    else: return ENCODER.decode(encoded_text[:max_tokens]) + " ..."
+
+
+
+
+def construct_prompt(query: str, history: List[Dict[str, str]], context: List[Block]) -> List[Dict[str, str]]:
+
+    prompt = []
+
+    # History takes the format: history=[
+    #     {"role": "user", "content": "Die monster. You don’t belong in this world!"},
+    #     {"role": "assistant", "content": "It was not by my hand I am once again given flesh. I was called here by humans who wished to pay me tribute."},
+    #     {"role": "user", "content": "Tribute!?! You steal men's souls and make them your slaves!"},
+    #     {"role": "assistant", "content": "Perhaps the same could be said of all religions..."},
+    #     {"role": "user", "content": "Your words are as empty as your soul! Mankind ill needs a savior such as you!"},
+    #     {"role": "assistant", "content": "What is a man? A miserable little pile of secrets. But enough talk... Have at you!"},
+    # ]
+
+    source_prompt = "You are a helpful assistant knowledgeable about AI Alignment and Safety. " \
+        "Please give a clear and coherent answer to the user's questions.(written after \"Q:\") " \
+        "using the following sources. Each source is labeled with a letter. Feel free to " \
+        "use the sources in any order, and try to use multiple sources in your answers.\n\n"
+
+    token_count = len(ENCODER.encode(source_prompt))
+
+    # Context from top-k blocks
+    for i, block in enumerate(context):
+        block_str = f"[{chr(ord('a') + i)}] {block.title} - {block.author} - {block.date}\n{block.text}\n\n"
+        block_tc = len(ENCODER.encode(block_str))
+
+        if token_count + block_tc > int(NUM_TOKENS * CONTEXT_FRACTION):
+            source_prompt += cap(block_str, int(NUM_TOKENS * CONTEXT_FRACTION) - token_count)
+            break
+        else:
+            source_prompt += block_str
+            token_count += block_tc
+
+    source_prompt = source_prompt.strip();
+    if len(history) > 0:
+        source_prompt += "\n\n"\
+            "Before the question (\"Q: \"), there will be a history of previous questions and answers. " \
+            "These sources only apply to the last question. Any sources used in previous answers " \
+            "are invalid."
+
+    prompt.append({"role": "system", "content": source_prompt.strip()})
+
+
+    # Write a version of the last 10 messages into history, cutting things off when we hit the token limit.
+    token_count = 0
+    history_trnc = []
+    for message in history[:-10:-1]:
+        if message["role"] == "user":
+            history_trnc.append({"role": "user", "content": "Q: " + message["content"]})
+            token_count += len(ENCODER.encode("Q: " + message["content"]))
+        else:
+            content = cap(message["content"], int(NUM_TOKENS * HISTORY_FRACTION) - token_count)
+
+            # censor all source letters into [x]
+            content = re.sub(r"\[[0-9]+\]", "[x]", content)
+
+            history_trnc.append({"role": "assistant", "content": content})
+            token_count += len(ENCODER.encode(content))
+
+        if token_count > int(NUM_TOKENS * HISTORY_FRACTION):
+            break
+
+    prompt.extend(history_trnc[::-1])
+
+
+    question_prompt = f"In your answer, please cite any claims you make back to each source " \
+                    f"using the format: [a], [b], etc. If you use multiple sources to make a claim " \
+                    f"cite all of them. For example: \"AGI is concerning [c, d, e].\"\n\nQ: " + query
+
+    prompt.append({"role": "user", "content": question_prompt})
+
+    return prompt
+
+# ------------------------------- completion code -------------------------------
+import time
+import json
+
+# returns either (True, reply string, top_k_blocks)) or (False, error message string, None)
+def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STANDARD_K):
+    try:
+        # 1. Find the most relevant blocks from the Alignment Research Dataset
+        yield json.dumps({"state": "loading", "phase": "semantic"})
+        top_k_blocks = get_top_k_blocks(index, query, k)
+
+        yield json.dumps({"state": "loading", "phase": "semantic", 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in top_k_blocks]})
+
+        # 2. Generate a prompt
+        yield json.dumps({"state": "loading", "phase": "prompt"})
+        prompt = construct_prompt(query, history, top_k_blocks)
+
+        # 3. Count number of tokens left for completion (-50 for a buffer)
+        max_tokens_completion = NUM_TOKENS - sum([len(ENCODER.encode(message["content"]) + ENCODER.encode(message["role"])) for message in prompt]) - 50
+
+        # 4. Answer the user query
+        yield json.dumps({"state": "loading", "phase": "llm"})
+        t1 = time.time()
+        response = ''
+
+        for chunk in openai.ChatCompletion.create(
+            model=COMPLETIONS_MODEL,
+            messages=prompt,
+            max_tokens=max_tokens_completion,
+            stream=True
+        ):
+            res = chunk["choices"][0]["delta"]
+            if res is not None and res.get("content") is not None:
+                response += res["content"]
+                yield json.dumps({"state": "streaming", "content": res["content"]})
+
+
+        t2 = time.time()
+        print("Time to get response: ", t2 - t1)
+
+        if DEBUG_PRINT:
+            print('\n' * 10)
+            print(" ------------------------------ prompt: -----------------------------")
+            for message in prompt:
+                print(f"----------- {message['role']}: ------------------")
+                print(message['content'])
+
+            print('\n' * 10)
+
+            print(" ------------------------------ response: -----------------------------")
+            print(response)
+
+        yield json.dumps({"state": "done"})
+
+    except Exception as e:
+        print(e)
+        yield json.dumps({"state": "error", "error": str(e)})
+

api/dataset_dl.py

@@ -0,0 +1,23 @@
+# run this script to download the dataset
+
+import pickle
+import requests
+import os
+
+print('Downloading dataset...')
+
+url = os.environ.get('DATASET_URL')
+if url is None:
+    print('No dataset url provided.')
+    exit()
+
+dataset_dict_bytes = requests.get(url).content
+
+print('Unpacking dataset...')
+dataset_dict = pickle.loads(dataset_dict_bytes)
+
+print('Writing dataset to disk...')
+with open('dataset.pkl', 'wb') as f:
+    pickle.dump(dataset_dict, f)
+
+print('Done!')

api/get_blocks.py

@@ -0,0 +1,122 @@
+from typing import List, Tuple
+import dataclasses
+import datetime
+import itertools
+import numpy as np
+import openai
+import regex as re
+import time
+
+# ---------------------------------- constants ---------------------------------
+
+EMBEDDING_MODEL = "text-embedding-ada-002"
+
+# ------------------------------------ types -----------------------------------
+
+@dataclasses.dataclass
+class Block:
+    title: str
+    author: str
+    date: str
+    url: str
+    tags: str
+    text: str
+    
+# ------------------------------------------------------------------------------
+
+# 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.
+def get_embedding(text: str) -> np.ndarray:
+
+    max_retries = 4
+    max_wait_time = 10
+    attempt = 0
+
+    while True:
+        try:
+            result = openai.Embedding.create(model=EMBEDDING_MODEL, input=text)
+            return result["data"][0]["embedding"]
+
+        except openai.error.RateLimitError as e:
+
+            attempt += 1
+
+            if attempt > max_retries: raise e
+
+            time.sleep(min(max_wait_time, 2 ** attempt))
+
+
+# Get the k blocks most semantically similar to the query using Pinecone.
+def get_top_k_blocks(index, user_query: str, k: int = 20) -> List[Block]:
+
+    # print time
+    t = time.time()
+
+    # Get the embedding for the query.
+    query_embedding = get_embedding(user_query)
+
+    t1 = time.time()
+    print("Time to get embedding: ", t1 - t)
+
+    query_response = index.query(
+        namespace="alignment-search",  # ugly, sorry
+        top_k=k,
+        include_values=False,
+        include_metadata=True,
+        vector=query_embedding
+    )
+    blocks = []
+    for match in query_response['matches']:
+
+        date = match['metadata']['date']
+
+        if type(date) == datetime.date: date = date.strftime("%Y-%m-%d") # iso8601
+
+        blocks.append(Block(
+            title = match['metadata']['title'],
+            author = match['metadata']['author'],
+            date = date,
+            url = match['metadata']['url'],
+            tags = match['metadata']['tags'],
+            text = match['metadata']['text']
+        ))
+
+    t2 = time.time()
+
+    print("Time to get top-k blocks: ", t2 - t1)
+    
+    # for all blocks that are "the same" (same title, author, date, url, tags),
+    # combine their text with "....." in between. Return them in order such
+    # that the combined block has the minimum index of the blocks combined.
+
+    key = lambda bi: (bi[0].title or "", bi[0].author or "", bi[0].date or "", bi[0].url or "", bi[0].tags or "")
+
+    blocks_plus_old_index = [(block, i) for i, block in enumerate(blocks)]
+    blocks_plus_old_index.sort(key=key)
+
+    unified_blocks: List[Tuple[Block, int]] = []
+
+    for key, group in itertools.groupby(blocks_plus_old_index, key=key):
+        group = list(group)
+        if len(group) == 0: continue
+        
+        group = group[:3] # limit to a max of 3 blocks from any one source
+
+        text = "\n.....\n".join([block[0].text for block in group])
+
+        min_index = min([block[1] for block in group])
+
+        unified_blocks.append((Block(key[0], key[1], key[2], key[3], key[4], text), min_index))
+
+    unified_blocks.sort(key=lambda bi: bi[1])
+    return [block for block, _ in unified_blocks]
+
+
+# we add the title and authors inside the contents of the block, so that
+# searches for the title or author will be more likely to pull it up. This
+# strips it back out.
+def strip_block(text: str) -> str:
+    r = re.match(r"^\"(.*)\"\s*-\s*Title:.*$", text, re.DOTALL)
+    if not r:
+        print("Warning: couldn't strip block")
+        print(text)
+    return r.group(1) if r else text

api/main.py

@@ -0,0 +1,70 @@
+from flask import Flask, jsonify, request, Response
+from flask_cors import CORS, cross_origin
+from get_blocks import get_top_k_blocks
+from chat import talk_to_robot
+import dataclasses
+import os
+import openai
+import pinecone
+
+
+# ---------------------------------- env setup ---------------------------------
+
+
+if os.path.exists('.env'):
+    from dotenv import load_dotenv
+    load_dotenv()
+
+OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
+openai.api_key = OPENAI_API_KEY
+
+PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
+PINECONE_ENV = "us-east1-gcp"
+pinecone.init(
+    api_key=PINECONE_API_KEY,
+    environment=PINECONE_ENV
+)
+INDEX_NAME = "alignment-search"
+index = pinecone.Index(index_name=INDEX_NAME)
+
+
+# ---------------------------------- web setup ---------------------------------
+
+app = Flask(__name__)
+cors = CORS(app)
+app.config['CORS_HEADERS'] = 'Content-Type'
+
+# ---------------------------------- sse stuff ---------------------------------
+
+def stream(src):
+    yield from ('data: ' + '\ndata: '.join(message.splitlines()) + '\n\n' for message in src)
+    yield 'event: close\n\n'
+
+# ------------------------------- semantic search ------------------------------
+
+
+@app.route('/semantic', methods=['POST'])
+@cross_origin()
+def semantic():
+    query = request.json['query']
+    return jsonify([dataclasses.asdict(block) for block in get_top_k_blocks(index, query)])
+
+
+# ------------------------------------ chat ------------------------------------
+
+
+@app.route('/chat', methods=['POST'])
+@cross_origin()
+def chat():
+
+    query = request.json['query']
+    history = request.json['history']
+
+    return Response(stream(talk_to_robot(index, query, history)), mimetype='text/event-stream')
+
+
+# ------------------------------------------------------------------------------
+
+
+if __name__ == '__main__':
+    app.run(debug=True, port=3000)

api/nixpacks.toml

@@ -0,0 +1,5 @@
+# [phases.build]
+# cmds = ['python3 dataset_dl.py']
+
+[start]
+cmd = 'gunicorn --timeout 300 "main:app"'

api/requirements.txt

@@ -0,0 +1,22 @@
+# ---- <flask stuff> ----
+
+Flask==1.1.2
+
+click==7.1.2
+gunicorn==20.0.4
+itsdangerous==1.1.0
+Jinja2==2.11.3
+MarkupSafe==1.1.1
+Werkzeug==2.2.3
+
+flask-cors
+
+# ---- </flask stuff> ----
+
+openai==0.27.2
+numpy==1.24.2
+tenacity==8.2.2
+tiktoken
+pinecone-client
+
+python-dotenv

readme.old

@@ -0,0 +1,38 @@
+# AlignmentSearch
+
+This project creates embeddings for every set of a few paragraphs from the source dataset, in order to do real-time semantic search and question answering on them.
+
+The very barebones of the project is currently in src/testing.ipynb. file which contains: 
+
+To try out, add a src/config.py file which contains your OPENAI_API_KEY.
+
+## TODO:
+- Getting data:
+    - Figure out the right format for the dataset
+    - Get entirety of data
+    - Searches for new posts/papers/etc and scrape them, runs once a day
+    - Async API calls for embeddings (otherwise it is going to take years)
+- Semantic search:
+    - Test out other techniques than just vector similarity (e.g. LSH-index, see Dense Retrieval methods (here)[https://medium.com/@aikho/deep-learning-in-information-retrieval-part-ii-dense-retrieval-1f9fecb47de9])
+    - Test other embeddings models ((SimCSE)[https://github.com/princeton-nlp/SimCSE] possibly SOTA?)
+- Question answering:
+    - Test out other models prompts to see which is best
+- Summarization:
+    - Test out other models and prompts to see which is best (Forefront?)
+- Info extraction from PDF:
+    - Specifically mentioned by Anson. Look into methods by Mely.ai to extract tables from PDFs maybe?
+    - Test various techniques to make it more performant
+- Finetuning:
+    - Finetune embeddings model
+    - Finetune q&a model
+    - Finetune summarization model
+    - Finetune info extraction model
+- Other:
+    - Create website/other. Not sure what would be most helpful here (Find someone that can figure this out)
+    - Document search (give descriptions of a post/document/video/article/book/etc related to alignment, and get top semantically related result)
+
+## Ideas
+- The whole dataset with embeddings doesn't fit on the frontend, but requests take a while. Solution: have a useful but small fraction of the dataset on the front end, get similarity score for embeddings, and do real time semantic search. However, every time the user presses space or enter, do the call to do semantic search over the full dataset.
+
+
+Source dataset: Kirchner, J. H., Smith, L., Thibodeau, J., McDonnell, K., and Reynolds, L. "Understanding AI alignment research: A Systematic Analysis." arXiv preprint arXiv:2022.4338861 (2022).

src/dataset/create_dataset.py

@@ -0,0 +1,401 @@
+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"))
+    """
+    

src/dataset/settings.py

@@ -0,0 +1,12 @@
+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 / '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')

src/dataset/text_splitter.py

@@ -0,0 +1,220 @@
+import re
+from typing import List
+import tiktoken
+
+import re
+from typing import List
+import nltk
+
+
+# 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.
+
+    :param text: The input text to be split.
+    :return: A list of sentences.
+    """
+    text = text.replace("\n", " ")    # Replace newline characters with spaces
+    sentences = nltk.sent_tokenize(text)    # Use the Punkt tokenizer from the NLTK library to split the text into sentences
+    sentences = [s.strip() for s in sentences]    # Strip leading and trailing whitespace from each sentence
+    return sentences
+
+
+class TokenSplitter:
+    """Splits text into blocks of tokens according to chatgpt's tokenizer."""
+
+    def __init__(self, min_tokens: int = 200, max_tokens: int = 300):
+        self.encoding = tiktoken.get_encoding("cl100k_base")
+        self.min_tokens = min_tokens
+        self.max_tokens = max_tokens
+        self.default_signature = "{url, title, author} unknown"
+
+    def _text_splitter(self, text: str, signature: str) -> List[str]:
+        """Splits text into blocks of tokens according to chatgpt's tokenizer."""
+        # enc = self.encoding.encode  # takes a string and returns a list of ints (tokens)
+        enc = self.encoding.encode_ordinary  # takes a string and returns a list of ints (tokens)
+        dec = self.encoding.decode  # takes a list of ints (tokens) and returns a string
+        tok_len = lambda x: len(enc(x))  # length of a string in tokens
+
+        max_tokens = self.max_tokens - tok_len(signature) - 10  # 10 to be safe
+        assert max_tokens > 0, "max_tokens is too small for the signature"
+
+        min_tokens = self.min_tokens - tok_len(signature) - 10  # 10 to be safe
+        assert min_tokens > 0, "min_tokens is too small for the signature"
+
+        blocks = []
+        current_block = ""
+        paragraphs = text.split("\n\n")
+        
+        for paragraph in paragraphs:
+            sentences = split_into_sentences(paragraph)
+            if current_block != "":
+                current_block += "\n\n"
+
+            for sentence in sentences:
+                potential_new_block = f"{current_block} {sentence}"
+                
+                if tok_len(potential_new_block) <= max_tokens:
+                    current_block = potential_new_block
+                
+                else:
+                    blocks.append(current_block)
+                    if tok_len(sentence) < max_tokens:
+                        current_block = sentence
+                    else:
+                        blocks.append(dec(enc(sentence)[:max_tokens]))
+                        current_block = ""
+            
+            if tok_len(current_block) > min_tokens:
+                blocks.append(current_block)
+                current_block = ""
+
+        if current_block != "":
+            if len(blocks) == 0:
+                blocks.append(current_block)
+            else:
+                latest_block = blocks[-1]
+                len_cur_block = tok_len(current_block)
+                latest_plus_current = latest_block + current_block
+
+                if len_cur_block > min_tokens:
+                    blocks.append(current_block)
+                
+                else:
+                    # select the last self.max_tokens tokens from the latest block
+                    last_block = dec(enc(latest_plus_current)[-max_tokens:])
+                    blocks.append(last_block)
+
+        return [block.strip() for block in blocks]
+
+    def split(self, text: str, signature: str = None) -> List[str]:
+        if signature is None:
+            signature = self.default_signature
+
+        blocks = self._text_splitter(text, signature)
+
+        # Check all block elements are strings
+        assert all([isinstance(block, str) for block in blocks]), "block elements are not strings"
+
+        output = [f'"{block}"\n- {signature}' for block in blocks]
+        # Check all output elements are strings
+        assert all([isinstance(block, str) for block in output]), "output elements are not strings"
+
+        return output
+
+if __name__ == "__main__":
+    text = """This post has been recorded as part of the LessWrong Curated Podcast, and an be listened to on Spotify, Apple Podcasts, and Libsyn.
+
+Over the last few years, deep-learning-based AI has progressed extremely rapidly in fields like natural language processing and image generation. However, self-driving cars seem stuck in perpetual beta mode, and aggressive predictions there have repeatedly been disappointing. Google's self-driving project started four years before AlexNet kicked off the deep learning revolution, and it still isn't deployed at large scale, thirteen years later. Why are these fields getting such different results?
+
+Right now, I think the biggest answer is that ML benchmarks judge models by average-case performance, while self-driving cars (and many other applications) require matching human worst-case performance. For MNIST, an easy handwriting recognition task, performance tops out at around 99.9% even for top models; it's not very practical to design for or measure higher reliability than that, because the test set is just 10,000 images and a handful are ambiguous. Redwood Research, which is exploring worst-case performance in the context of AI alignment, got reliability rates around 99.997% for their text generation models.
+
+By comparison, human drivers are ridiculously reliable. The US has around one traffic fatality per 100 million miles driven; if a human driver makes 100 decisions per mile, that gets you a worst-case reliability of ~1:10,000,000,000 or ~99.999999999%. That's around five orders of magnitude better than a very good deep learning model, and you get that even in an open environment, where data isn't pre-filtered and there are sometimes random mechanical failures. Matching that bar is hard! I'm sure future AI will get there, but each additional "nine" of reliability is typically another unit of engineering effort. (Note that current self-driving systems use a mix of different models embedded in a larger framework, not one model trained end-to-end like GPT-3.)
+
+(The numbers here are only rough Fermi estimates. I'm sure one could nitpick them by going into pre-pandemic vs. post-pandemic crash rates, laws in the US vs. other countries, what percentage of crashes are drunk drivers, do drunk drivers count, how often would a really bad decision be fatal, etc. But I'm confident that whichever way you do the math, you'll still find that humans are many orders of magnitude more reliable.)
+
+Other types of accidents are similarly rare. Eg. pre-pandemic, there were around 40 million commercial flights per year, but only a handful of fatal crashes. If each flight involves 100 chances for the pilot to crash the plane by screwing up, then that would get you a reliability rate around 1:1,000,000,000, or ~99.99999999%.
+
+Even obviously dangerous activities can have very low critical failure rates. For example, shooting is a popular hobby in the US; the US market buys around 10 billion rounds of ammunition per year. There are around 500 accidental gun deaths per year, so shooting a gun has a reliability rate against accidental death of ~1:20,000,000, or 99.999995%. In a military context, the accidental death rate was around ten per year against ~1 billion rounds fired, for a reliability rate of ~99.9999999%. Deaths by fire are very rare compared to how often humans use candles, stoves, and so on; New York subway deaths are rare compared to several billion annual rides; out of hundreds of millions of hikers, only a tiny percentage fall off of cliffs; and so forth.
+
+The 2016 AI Impacts survey asked hundreds of AI researchers when they thought AI would be capable of doing certain tasks, playing poker, proving theorems and so on. Some tasks have been solved or have a solution "in sight", but right now, we're nowhere close to an AI that can replace human surgeons; robot-assisted surgeries still have manual control by human operators. Cosmetic surgeries on healthy patients have a fatality rate around 1:300,000, even before excluding unpredictable problems like blood clots. If a typical procedure involves two hundred chances to kill the patient by messing up, then an AI surgeon would need a reliability rate of at least 99.999998%.
+
+One concern with GPT-3 has been that it might accidentally be racist or offensive. Humans are, of course, sometimes racist or offensive, but in a tightly controlled Western professional context, it's pretty rare. Eg., one McDonald's employee was fired for yelling racial slurs at a customer. But McDonald's serves 70 million people a day, ~1% of the world's population. Assuming that 10% of such incidents get a news story and there's about one story per year, a similar language model would need a reliability rate of around 1:2,500,000,000, or 99.99999996%, to match McDonald's workers. When I did AI for the McDonald's drive-thru, the language model wasn't allowed to generate text at all. All spoken dialog had to be pre-approved and then manually engineered in. Reliability is hard!
+
+On the one hand, this might seem slightly optimistic for AI alignment research, since commercial AI teams will have to get better worst-case bounds on AI behavior for immediate economic reasons. On the other hand, because so much of the risk of AI is concentrated into a small number of very bad outcomes, it seems like such engineering might get us AIs that appear safe, and almost always are safe, but will still cause catastrophic failure in conditions that weren't anticipated. That seems bad."""
+    text = """Imagine it's late autumn of 332 BC. You're Alexander the Great, and your armies are marching toward Egypt from Gaza. There’s just one little problem: you need to cross the Sinai peninsula - 150 miles of hot, barren desert. How will you carry food and water for the troops?
+
+
+Green triangle on the left is the Nile river delta in Egypt; green chunk in the upper right is Israel. The big desert peninsula between them is the Sinai.
+
+Option 1: carry it
+
+A physically-active human needs about 3 lbs of food per day. (Modern hikers can probably find lighter calorie-dense foodstuffs, but we’re talking ancient history here.) Water requirements vary; 5 lbs is a minimum, but the US Army Quartermaster Corps recommends 20 lbs/day when marching through a hot desert. Alexander’s army crossed the desert in 7 days. Food might be reasonable, but to carry the water would mean 7*20 = 140 lbs per person, plus 50+ lbs of armor, weapons, etc.
+
+When I go hiking, I aim for a 20-30 lb pack. US marines are apparently expected to be able to carry 150 lbs for 9 miles - quite a bit less than the 20+ miles/day Alexander’s army managed, and with no comment on how long the marine in question might need to rest afterwards. (Also, I’m not sure I trust that source - 150 lbs for 9 miles sounds unrealistic to me, and if it’s true then I’m very impressed by marines.)
+
+Suffice to say that carrying that much water across that much desert is not a realistic option, even if we drink it along the way.
+
+Option 2: horses
+
+A horse consumes 20 lbs of food (half of which may be forage) and 80 lbs of water per day. In exchange, it can carry about 200 lbs (surprisingly, my source claims that horses can carry more than they can pull). Of course, that 200 lbs has to include the horse’s own food and water, plus whatever useful load it’s carrying. So, marching through a desert, a horse can only transport (200 lbs)/(80+20 lbs/day) = 2 days of supplies for itself, and that’s before whatever useful things actually need to be transported.
+
+In other words, there’s a hard upper limit on how far goods can be transported by horse without refilling supplies along the way. That limit is around 2 days travel time without any refill, 10 days if there’s plenty of fresh water along the route, or 20 days if there’s both water and forage. At 20 miles/day, that’s 40, 200, or 400 miles. Realistically, if we want the number of horses to be reasonable, the limit is more like half that much - 20 miles, 100 miles, or 200 miles, respectively.
+
+So horses also won’t work.
+
+Option 2.5: camels or other pack animals
+
+Contrary to popular image, camels actually need more water than horses. They can go a couple days without, but then need to fill up all at once. They can also carry a bit more weight, but they eat more food. At the end of the day, the numbers end up quite similar.
+
+Mules also end up with similar numbers, and cattle are generally worse.
+
+Option 3: ships
+
+Assuming the army marches along the coast, a supply fleet can sail alongside. At the time, a single large merchant ship could carry 400 tons - in other words, as much as about 4000 horses. Presumably the ship would cost a lot less than the horses, too.
+
+Well then, there’s our answer. Ships are clearly a vastly superior way to move goods. Range is a non-issue, capacity is far larger, and they’re far cheaper. They’re perfect for crossing the Sinai, which runs right along the coast anyway.
+
+Fast forward a few years to 327 BC, and Alexander is marching his armies back from India. He plans to cross the Gedrosian desert, along the coast of modern-day Pakistan and Iran. The plan is much like the Sinai: a supply fleet will sail alongside the army. Unfortunately, neither Alexander nor his commanders knows about the monsoons: across most of south Asia, the wind blows consistently southwest for half the year, and consistently northeast for the other half. There is nothing like it in the Mediterranean. And so, Alexander marches out expecting the fleet to catch up as soon as the winds turn - not realizing that the winds will not turn for months. Three quarters of his soldiers die in the desert.
+
+Thus end the campaigns of Alexander.
+
+Generalization
+The above numbers are drawn from Donald Engels’ book Alexander the Great and the Logistics of Macedonian Army. But it tells us a lot more about the world than just the logistics of one particular ancient army.
+
+First, this highlights the importance of naval dominance in premodern warfare. A fleet was a far superior supply train, capable of moving a high volume of food and water over long distance at relatively low cost. Without a fleet, transport of food became expensive at best, regular resupply became a strategic necessity, and long routes through arid terrain became altogether impassable. Destroying an enemy’s fleet meant starving the army. Likewise, controlling ports wasn’t just for show - without a port, feeding the army became a serious problem.
+
+Another interesting insight into premodern warfare: away from friendly seas and rivers, the only way to keep an army fed was to either seize grain from enemies, or buy it from allies, either of whom needed to already be nearby. In Alexander’s case, deals were often struck to establish supply caches along the army’s intended route.
+
+An interesting exercise: to what extent was transportation a binding constraint on the size of premodern towns/cities? (One number you may want: Braudel (pg 121) estimates that 5000 square meters of land growing wheat would provide one person-year of food, not accounting for crop rotation.) Leave a comment if you try a calculation here; I'm curious to see how other peoples' models compare to my own.
+
+Modern Day
+Today we have trains and trucks and roads, so the transportation constraint has relaxed somewhat. But here’s an interesting comparison: a modern 18-wheeler in the US is legally limited to haul 40 tons, while a panamax ship could carry about 50k tons through the canal (prior to the opening of the new locks in 2016). That’s a ratio of a bit over 1000 - surprisingly similar to the ship/horse ratio of antiquity, especially considering the much larger new-panamax and super-panamax ships also in use today.
+
+
+Can we get a quick-and-dirty feel for tautness of the transportation constraint today? Here are a few very different angles:
+
+This USDA study shows rates on produce transport, typically about 7-20 cents per pound (see figure 6). The Smart & Final grocery store near me sells the cheaper produce items looked at in that study (bell peppers, cantaloupes, tomatoes, oranges) for 70-100 cents per pound, so transport alone is roughly 10-20% of the cost-to-consumer.
+What about transporting humans? Average commute in the US is ~30 minutes each way; driving is usually in the 20-30 minute range, while public transit is usually 30-50. Assuming 8 hr workdays, that means commutes are typically ~10-20% of our work-hours.
+The bureau of transportation estimates transport at 5.6% of the US economy for a very narrow measure, or 8.9% with a broader measure (though this still excludes non-market transport costs like e.g. commute time).
+My interpretation: the transportation constraint becomes taut when it accounts for 10-20% of cost. If it’s less than that, it usually doesn’t limit production - we see plenty of goods which aren’t transportation-dependent or which are higher-value-per-weight, and the transportation constraint is generally slack for those. But once transportation hits about 10-20%, people start looking for alternatives, i.e. producing the goods somewhere else or using alternative goods. Obviously this is not based on very much data, but I find it intuitively plausible.
+
+Compared to ancient times, transportation constraints have obviously relaxed quite a lot. Yet qualitatively, the world today still does not look like a world of fully slack transportation constraints. To wrap up, let’s discuss what that would look like.
+
+Extreme Slackness
+In Material Goods as an Abundant Resource, we discussed the world of the duplicator - a device capable of copying any item placed on it. In such a world, material scarcity is removed as an economic constraint - all material constraints are completely slack.
+
+What would be a corresponding sci-fi device for transportation constraints, and what would that world look like?
+
+I suggest portals: imagine we can create pairs of devices capable of transporting things from one device to the other, across any distance, at the speed of light. (We could instead imagine teleporters, removing the need for a pre-installed device at either end, but then the entire discussion would be about security.) What does the world of the portal look like?
+
+First, there’s complete geographical decoupling of production from consumption. People have no need to live near where they work; companies can put offices and factories wherever real estate is cheap. We can enjoy miles of wilderness on the back porch and a downtown district on the front porch; a swimming pool can open right into the ocean. Buying direct from the farm or factory is standard for most material goods.
+
+What are now tourist destinations would become options for an evening activity. Disneyworld would sell a park-hopper ticket that includes Disneyland California, Paris, and Shanghai, but the price of that ticket would be high enough to prevent the parks from becoming unpleasantly crowded - probably quite a bit more expensive than today, though possibly cheaper than today’s flights to Orlando.
+
+Obviously roads would cease to exist. Huge amounts of land would revert from asphalt to wilderness, but buildings would also be much more spread out. Buildings would be built close together more for show than for function - e.g. to provide the ambiance of a downtown or a community to those who want it. Physical life, in general, would look more like the structure of the internet rather than the structure of geography; “cities” would be clusters very spread out in space but very tightly connected via the portal network. Filter bubbles would be a much more physically tangible phenomenon.
+
+Geographically-defined governments would likely be replaced by some other form of government - governments based around access to portal hubs/networks are one natural possibility. Security would be a priority, early on - carrying an unauthorized portal into an area would earn a facefull of high explosives. On the other hand, it would be hard to prevent a high degree of mobility between areas controlled by different governments; the implications for government behavior are conceptually similar to seasteading.
+
+The structure of space near portal networks would be different in a big-O sense; the amount of space at a distance of about 
+r
+ would increase exponentially, rather than like 
+r
+2
+. A nuclear warhead could go off five hundred feet away and you’d feel a breeze through a fast-branching portal network. On the other hand, viruses could spread much more rapidly.
+
+Anyway, at this point we’re getting into specifics of portals, so I’ll cut off the speculation. The point is: if transportation continues to get cheaper and more efficient over time, then we will converge to the world of the portal, or at least something like it. The details do matter - portals are different from teleportation or whatever might actually happen - but any method of fully relaxing transportation constraints will have qualitatively similar results, to a large extent."""
+   
+    signature = "Title: Humans are very reliable agents, Author: alyssavance"
+
+    splitting = TokenSplitter(max_tokens=200, min_tokens=300)
+    blocks = splitting.split(text, signature)
+    context = "Context: " + "\n\n---\n\n".join(blocks)
+    print(context)

src/main.py

@@ -0,0 +1,163 @@
+import openai
+"""
+import config
+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,
+)
+
+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
+
+
+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,
+    )
+    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")
+
+    
+    
+
+
+
+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)

src/settings.py

@@ -0,0 +1,12 @@
+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')

web/.env.example

@@ -0,0 +1,14 @@
+# Since the ".env" file is gitignored, you can use the ".env.example" file to
+# build a new ".env" file when you clone the repo. Keep this file up-to-date
+# when you add new variables to `.env`.
+
+# This file will be committed to version control, so make sure not to have any
+# secrets in it. If you are cloning this repo, create a copy of this file named
+# ".env" and populate it with your secrets.
+
+# When adding additional environment variables, the schema in "/src/env.mjs"
+# should be updated accordingly.
+
+# Example:
+# SERVERVAR="foo"
+# NEXT_PUBLIC_CLIENTVAR="bar"

web/.gitignore

@@ -0,0 +1,42 @@
+# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.
+
+# dependencies
+/node_modules
+/.pnp
+.pnp.js
+
+# testing
+/coverage
+
+# database
+/prisma/db.sqlite
+/prisma/db.sqlite-journal
+
+# next.js
+/.next/
+/out/
+next-env.d.ts
+
+# production
+/build
+
+# misc
+.DS_Store
+*.pem
+
+# debug
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+.pnpm-debug.log*
+
+# local env files
+# do not commit any .env files to git, except for the .env.example file. https://create.t3.gg/en/usage/env-variables#using-environment-variables
+.env
+.env*.local
+
+# vercel
+.vercel
+
+# typescript
+*.tsbuildinfo

web/README.md

@@ -0,0 +1,28 @@
+# Create T3 App
+
+This is a [T3 Stack](https://create.t3.gg/) project bootstrapped with `create-t3-app`.
+
+## What's next? How do I make an app with this?
+
+We try to keep this project as simple as possible, so you can start with just the scaffolding we set up for you, and add additional things later when they become necessary.
+
+If you are not familiar with the different technologies used in this project, please refer to the respective docs. If you still are in the wind, please join our [Discord](https://t3.gg/discord) and ask for help.
+
+- [Next.js](https://nextjs.org)
+- [NextAuth.js](https://next-auth.js.org)
+- [Prisma](https://prisma.io)
+- [Tailwind CSS](https://tailwindcss.com)
+- [tRPC](https://trpc.io)
+
+## Learn More
+
+To learn more about the [T3 Stack](https://create.t3.gg/), take a look at the following resources:
+
+- [Documentation](https://create.t3.gg/)
+- [Learn the T3 Stack](https://create.t3.gg/en/faq#what-learning-resources-are-currently-available) — Check out these awesome tutorials
+
+You can check out the [create-t3-app GitHub repository](https://github.com/t3-oss/create-t3-app) — your feedback and contributions are welcome!
+
+## How do I deploy this?
+
+Follow our deployment guides for [Vercel](https://create.t3.gg/en/deployment/vercel), [Netlify](https://create.t3.gg/en/deployment/netlify) and [Docker](https://create.t3.gg/en/deployment/docker) for more information.

web/next.config.mjs

@@ -0,0 +1,22 @@
+/**
+ * Run `build` or `dev` with `SKIP_ENV_VALIDATION` to skip env validation.
+ * This is especially useful for Docker builds.
+ */
+!process.env.SKIP_ENV_VALIDATION && (await import("./src/env.mjs"));
+
+/** @type {import("next").NextConfig} */
+const config = {
+  reactStrictMode: true,
+
+  /**
+   * If you have the "experimental: { appDir: true }" setting enabled, then you
+   * must comment the below `i18n` config out.
+   *
+   * @see https://github.com/vercel/next.js/issues/41980
+   */
+  i18n: {
+    locales: ["en"],
+    defaultLocale: "en",
+  },
+};
+export default config;

web/package.json

@@ -0,0 +1,39 @@
+{
+  "name": "alignment_search",
+  "version": "0.1.0",
+  "private": true,
+  "scripts": {
+    "build": "next build",
+    "dev": "next dev",
+    "lint": "next lint",
+    "start": "next start"
+  },
+  "dependencies": {
+    "autosize": "^6.0.1",
+    "next": "^13.2.4",
+    "react": "18.2.0",
+    "react-dom": "18.2.0",
+    "react-textarea-autosize": "^8.4.1",
+    "zod": "^3.20.6"
+  },
+  "devDependencies": {
+    "@types/eslint": "^8.21.1",
+    "@types/node": "^18.14.0",
+    "@types/prettier": "^2.7.2",
+    "@types/react": "^18.0.28",
+    "@types/react-dom": "^18.0.11",
+    "@typescript-eslint/eslint-plugin": "^5.53.0",
+    "@typescript-eslint/parser": "^5.53.0",
+    "autoprefixer": "^10.4.7",
+    "eslint": "^8.34.0",
+    "eslint-config-next": "^13.2.1",
+    "postcss": "^8.4.14",
+    "prettier": "^2.8.1",
+    "prettier-plugin-tailwindcss": "^0.2.1",
+    "tailwindcss": "^3.2.0",
+    "typescript": "^4.9.5"
+  },
+  "ct3aMetadata": {
+    "initVersion": "7.8.0"
+  }
+}

web/postcss.config.cjs

@@ -0,0 +1,8 @@
+const config = {
+  plugins: {
+    tailwindcss: {},
+    autoprefixer: {},
+  },
+};
+
+module.exports = config;

web/src/env.mjs

@@ -0,0 +1,74 @@
+import { z } from "zod";
+
+/**
+ * Specify your server-side environment variables schema here. This way you can ensure the app isn't
+ * built with invalid env vars.
+ */
+const server = z.object({
+  NODE_ENV: z.enum(["development", "test", "production"]),
+});
+
+/**
+ * Specify your client-side environment variables schema here. This way you can ensure the app isn't
+ * built with invalid env vars. To expose them to the client, prefix them with `NEXT_PUBLIC_`.
+ */
+const client = z.object({
+  // NEXT_PUBLIC_CLIENTVAR: z.string().min(1),
+});
+
+/**
+ * You can't destruct `process.env` as a regular object in the Next.js edge runtimes (e.g.
+ * middlewares) or client-side so we need to destruct manually.
+ *
+ * @type {Record<keyof z.infer<typeof server> | keyof z.infer<typeof client>, string | undefined>}
+ */
+const processEnv = {
+  NODE_ENV: process.env.NODE_ENV,
+  // NEXT_PUBLIC_CLIENTVAR: process.env.NEXT_PUBLIC_CLIENTVAR,
+};
+
+// Don't touch the part below
+// --------------------------
+
+const merged = server.merge(client);
+
+/** @typedef {z.input<typeof merged>} MergedInput */
+/** @typedef {z.infer<typeof merged>} MergedOutput */
+/** @typedef {z.SafeParseReturnType<MergedInput, MergedOutput>} MergedSafeParseReturn */
+
+let env = /** @type {MergedOutput} */ (process.env);
+
+if (!!process.env.SKIP_ENV_VALIDATION == false) {
+  const isServer = typeof window === "undefined";
+
+  const parsed = /** @type {MergedSafeParseReturn} */ (
+    isServer
+      ? merged.safeParse(processEnv) // on server we can validate all env vars
+      : client.safeParse(processEnv) // on client we can only validate the ones that are exposed
+  );
+
+  if (parsed.success === false) {
+    console.error(
+      "❌ Invalid environment variables:",
+      parsed.error.flatten().fieldErrors,
+    );
+    throw new Error("Invalid environment variables");
+  }
+
+  env = new Proxy(parsed.data, {
+    get(target, prop) {
+      if (typeof prop !== "string") return undefined;
+      // Throw a descriptive error if a server-side env var is accessed on the client
+      // Otherwise it would just be returning `undefined` and be annoying to debug
+      if (!isServer && !prop.startsWith("NEXT_PUBLIC_"))
+        throw new Error(
+          process.env.NODE_ENV === "production"
+            ? "❌ Attempted to access a server-side environment variable on the client"
+            : `❌ Attempted to access server-side environment variable '${prop}' on the client`,
+        );
+      return target[/** @type {keyof typeof target} */ (prop)];
+    },
+  });
+}
+
+export { env };

web/src/header.tsx

@@ -0,0 +1,43 @@
+import React from "react";
+import Link from "next/link";
+
+const Header: React.FC<{page: "index" | "semantic"}> = ({page}) => {
+    const sidebar = page === "index" ? (
+        <span className="flex flex-col font-semibold flex-1 justify-start text-right">
+            <p className="my-0">Conversational Agent</p>
+            <Link href="/semantic">Semantic Search</Link>
+        </span>
+    ) : (
+        <span className="flex flex-col font-semibold flex-1 justify-start text-right">
+            <Link href="/">Conversational Agent</Link>
+            <p className="my-0">Semantic Search</p>
+        </span>
+    );
+        
+    return (<>
+                <div className="flex my-4">
+                    <h1 className="flex-1 my-0">AlignmentSearch</h1>
+                    {sidebar}
+                </div>
+                <p>
+                    This site has been created by McGill students Henri Lemoine, Fraser Lee, and Thomas Lemoine
+                    as an attempt to create a "conversational FAQ" that can answer questions about AI alignment.
+                    When asked a question, we
+                </p>
+                <ol>
+                    <li>Embed the question into a low dimensional semantic space</li>
+                    <li>Pull the semantically closest passages out of a massive alignment dataset</li>
+                    <li>Instruct an LLM to construct a response while citing these passages</li>
+                    <li>Display this response in conversational flow with inline citations</li>
+                </ol>
+                <p>
+                    We've created this as an attempt on the <a href="https://www.lesswrong.com/posts/SLRLuiuDykfTdmesK/speed-running-everyone-through-the-bad-alignement-bingo">$5k bounty for a LW conversational agent</a>.
+                </p>
+                <p>
+                    We hope that it can be a useful tool for the community, and we're eager to hear feedback and suggestions.
+                    For a technical report on our implemention, see our <a href="https://www.lesswrong.com/posts/bGn9ZjeuJCg7HkKBj/introducing-alignmentsearch-an-ai-alignment-informed">LessWrong post</a>.
+                </p>
+            </>);
+};
+
+export default Header;

web/src/pages/_app.tsx

@@ -0,0 +1,9 @@
+import { type AppType } from "next/dist/shared/lib/utils";
+
+import "~/styles/globals.css";
+
+const MyApp: AppType = ({ Component, pageProps }) => {
+  return <Component {...pageProps} />;
+};
+
+export default MyApp;

web/src/pages/index.tsx

@@ -0,0 +1,417 @@
+const API_URL = process.env.NEXT_PUBLIC_API_URL || "http://127.0.0.1:3000";
+
+import Head from "next/head";
+import React from "react";
+import { type NextPage } from "next";
+import { useState } from "react";
+
+import Header from "../header";
+import SearchBox from "../searchbox";
+
+type Citation = {
+    title: string;
+    author: string;
+    date: string;
+    url: string;
+}
+
+type UserEntry = {
+    role: "user";
+    content: string;
+}
+
+type AssistantEntry = {
+    role: "assistant";
+    content: string;
+    citations: Citation[];
+    base_count: number; // the number to start counting citations at
+}
+
+type ErrorMessage = {
+    role: "error";
+    content: string;
+}
+
+type Entry = UserEntry | AssistantEntry | ErrorMessage;
+
+// const Colours = ["blue", "cyan", "teal", "green", "amber"].map(colour => `bg-${colour}-100 border-${colour}-300 text-${colour}-800`);
+// this would be nice, but Tailwind needs te actual string of the class to be in
+// the source file for it to be included in the build
+
+const Colours = [
+    "bg-red-100    border-red-300    text-red-800",
+    "bg-amber-100  border-amber-300  text-amber-800",
+    "bg-orange-100 border-orange-300 text-orange-800",
+    "bg-lime-100   border-lime-300   text-lime-800",
+    "bg-green-100  border-green-300  text-green-800",
+    "bg-cyan-100   border-cyan-300   text-cyan-800",
+    "bg-blue-100   border-blue-300   text-blue-800",
+    "bg-violet-100 border-violet-300 text-violet-800",
+    "bg-pink-100   border-pink-300   text-pink-800",
+];
+
+const ShowCitation: React.FC<{citation: Citation, i: number}> = ({citation, i}) => {
+
+    var c_str = citation.title;
+
+    if (citation.author && citation.author !== "")
+        c_str += " - " + citation.author;
+    if (citation.date && citation.date !== "")
+        c_str += " - " + citation.date;
+
+    return (
+        <A className={Colours[i % Colours.length] + " border-2 flex items-center rounded my-2 text-sm no-underline w-fit"}
+            href={citation.url}>
+            <span className="mx-1"> [{i + 1}] </span>
+            <p className="mx-1 my-0"> {c_str} </p>
+        </A>
+    );
+};
+
+const ShowInTextCitation: React.FC<{citation: Citation, i: number}> = ({citation, i}) => {
+    return (
+        <A className={Colours[i % Colours.length] + " border-2 rounded text-sm no-underline w-min px-0.5 pb-0.5 ml-1 mr-0.5"}
+            href={citation.url}>
+            [{i + 1}]
+        </A>
+    );
+};
+
+const A: React.FC<{href: string, className?: string, children: React.ReactNode}> = ({href, className, children}) => {
+    return href && href !== "" ? (
+        <a className={className} href={href} target="_blank" rel="noreferrer">
+            {children}
+        </a>
+    ) : (
+        <a className={className}>
+            {children}
+        </a>
+    );
+}
+
+
+
+
+
+// todo: memoize this if too slow.
+const ProcessText: (text: string, base_count: number) => [string, Map<string, number>] = (text, base_count) => {
+
+    // ---------------------- normalize citation form ----------------------
+
+    // transform all things that look like [a, b, c] into [a][b][c]
+    let response = text.replace(
+
+                /\[((?:[a-z]+,\s*)*[a-z]+)\]/g, // identify groups of this form
+
+                (block: string) => block.split(',')
+                                        .map((x) => x.trim())
+                                        .join("][")
+    )
+
+    // transform all things that look like [(a), (b), (c)] into [(a)][(b)][(c)]
+    response = response.replace(
+
+            /\[((?:\([a-z]+\),\s*)*\([a-z]+\))\]/g, // identify groups of this form
+
+            (block: string) => block.split(',')
+                                    .map((x) => x.trim())
+                                    .join("][")
+    )
+
+    // transform all things that look like [(a)] into [a]
+    response = response.replace(
+        /\[\(([a-z]+)\)\]/g,
+        (_match: string, x: string) => `[${x}]`
+    )
+
+    // transform all things that look like [ a ] into [a]
+    response = response.replace(
+        /\[\s*([a-z]+)\s*\]/g,
+        (_match: string, x: string) => `[${x}]`
+    )
+
+    // -------------- map citations from strings into numbers --------------
+
+    // figure out what citations are in the response, and map them appropriately
+    const cite_map = new Map<string, number>();
+    let cite_count = 0;
+
+    // scan a regex for [x] over the response. If x isn't in the map, add it.
+    const regex = /\[([a-z]+)\]/g;
+    let match;
+    let response_copy = ""
+    while ((match = regex.exec(response)) !== null) {
+        if (!cite_map.has(match[1]!)) {
+            cite_map.set(match[1]!, base_count + cite_count++);
+        }
+        // replace [x] with [i]
+        response_copy += response.slice(response_copy.length, match.index) + `[${cite_map.get(match[1]!)! + 1}]`;
+    }
+
+    response = response_copy + response.slice(response_copy.length);
+
+    return [response, cite_map]
+}
+
+
+const ShowAssistantEntry: React.FC<{entry: AssistantEntry}> = ({entry}) => {
+    const in_text_citation_regex = /\[([0-9]+)\]/g;
+
+    let [response, cite_map] = ProcessText(entry.content, entry.base_count);
+
+    // ----------------- create the ordered citation array -----------------
+
+    const citations = new Map<number, Citation>();
+    cite_map.forEach((value, key) => {
+        const index = key.charCodeAt(0) - 'a'.charCodeAt(0);
+        if (index >= entry.citations.length) {
+            console.log("invalid citation index: " + index);
+        } else {
+            citations.set(value, entry.citations[index]!);
+        }
+    });
+
+    return (
+        <div className="mt-3 mb-8">
+            {   // split into paragraphs
+                response.split("\n").map(paragraph => ( <p> {
+                    paragraph.split(in_text_citation_regex).map((text, i) => {
+                        if (i % 2 === 0) {
+                            return text.trim();
+                        }
+                        i = parseInt(text) - 1;
+                        if (!citations.has(i)) return `[${text}]`;
+                        const citation = citations.get(i)!;
+                        return (
+                            <ShowInTextCitation citation={citation} i={i} />
+                        );
+                    })
+                } </p>))
+            }
+            <ul className="mt-5">
+                {   // show citations
+                    Array.from(citations.entries()).map(([i, citation]) => (
+                        <li key={i}>
+                            <ShowCitation citation={citation} i={i} />
+                        </li>
+                    ))
+                }
+            </ul>
+        </div>
+    );
+};
+
+
+
+
+
+type State = {
+    state: "idle";
+} | {
+    state: "loading";
+    phase: "semantic" | "prompt" | "llm";
+    citations: Citation[];
+} | {
+    state: "streaming";
+    response: AssistantEntry;
+};
+
+
+const Home: NextPage = () => {
+
+    const [ entries, setEntries ] = useState<Entry[]>([]);
+    const [ runningIndex, setRunningIndex ] = useState(0);
+    const [ loadState, setLoadState ] = useState<State>({state: "idle"});
+
+    const search = async (
+        query: string,
+        setQuery: (query: string) => void,
+        setLoading: (loading: boolean) => void
+    ) => {
+
+        // clear the query box, append to entries
+
+        const old_entries = entries;
+        const new_entries: Entry[] = [...old_entries, {role: "user", content: query}];
+        setEntries(new_entries);
+        setQuery("");
+        setLoading(true);
+
+        // do SSE on a POST request.
+
+        const res = await fetch(API_URL + "/chat", {
+            method: "POST",
+            cache: "no-cache",
+            keepalive: true,
+            headers: {
+                "Content-Type": "application/json",
+                "Accept": "text/event-stream",
+                "Allow-Control-Allow-Origin": "*"
+            },
+
+            body: JSON.stringify({query: query, history:
+                old_entries.filter((entry) => entry.role !== "error")
+                           .map((entry) => {
+                               return {
+                                   "role" : entry.role,
+                                   "content" : entry.content.trim(),
+                               }
+                           })
+            }),
+
+        });
+
+        if (!res.ok) {
+            setLoading(false);
+            setLoadState({state: "idle"});
+            setEntries([...new_entries, {role: "error", content: "POST Error: " + res.status}]);
+            return;
+        }
+
+        // read back the SSE stream
+
+        const reader = res.body!.getReader();
+        var message = "";
+        read: while (true) {
+
+            const {done, value} = await reader.read();
+
+            if (done) break;
+            const chunk = new TextDecoder("utf-8").decode(value);
+            if (chunk.startsWith("event: close\n")) break;
+
+            // note: this form isn't even remotely close to optimal in terms of network usage.
+
+            for (const line of chunk.split('\n')) {
+
+                // Most times, it seems that a single read() call will be one SSE "message",
+                // but I'll do the proper aggregation spec thing in case that's not always true.
+
+                if (line.startsWith("data: ")) message += line.slice(6);
+                if (line === "") {
+                    if (message !== "") {
+                        const data = JSON.parse(message);
+
+                        switch (data.state) {
+
+                            case "loading":
+
+                                // display loading phases, once citations are available toss them
+                                // into the loading state.
+
+                                setLoadState((s) => {
+                                    var citations = s.state === "loading" ? s.citations : [];
+                                    if (data.citations !== undefined) {
+                                        citations = data.citations;
+                                    }
+                                    return {state: "loading", phase: data.phase, citations: citations};
+                                });
+
+                                break;
+
+                            case "streaming":
+
+                                // incrementally build up the response
+
+                                setLoadState((s) => {
+                                    const response = s.state === "streaming" ? s.response :
+                                                {role: "assistant",
+                                                 content: "",
+                                                 citations: s.state === "loading" ? s.citations : [],
+                                                 base_count: runningIndex
+                                                };
+
+                                    return {state: "streaming", response: {
+                                        role: "assistant",
+                                        content: response.content + data.content,
+                                        citations: response.citations,
+                                        base_count: response.base_count
+                                    }};
+                                });
+
+                                // smooth-scroll to the bottom of the window if we're already less than 30% a screen away
+                                // note: finicky interaction with "smooth" - maybe fix later.
+                                if (document.documentElement.scrollHeight - window.scrollY < window.innerHeight * 1.3)
+                                    window.scrollTo({top: document.body.scrollHeight, behavior: "smooth"});
+                                break;
+
+                            case "done":
+
+                                // append the response to the entries, reset to normal
+
+                                setLoadState((s) => {
+                                    if (s.state === "streaming") {
+                                        setEntries([...new_entries, s.response]);
+                                        setRunningIndex((i) => (i + ProcessText(s.response.content, 0)[1].size));
+                                    }
+                                    return {state: "idle"};
+                                });
+                                break read;
+
+                            case "error":
+                                setEntries([...new_entries, {role: "error", content: data.error}]);
+                                break read;
+
+                        }
+                    }
+                    message = "";
+                }
+            }
+        }
+
+        setLoading(false);
+        setLoadState({state: "idle"});
+        if (document.documentElement.scrollHeight - window.scrollY < window.innerHeight * 1.3)
+            window.scrollTo({top: document.body.scrollHeight, behavior: "smooth"});
+
+    };
+
+    return (
+        <>
+            <Head>
+                <title>Alignment Search</title>
+            </Head>
+            <main>
+                <Header page="index" />
+                <ul>
+                    {entries.map((entry, i) => {
+                        if (entry.role === "user") {
+                            return <li key={i}>
+                                <p className="border border-gray-300 px-1 text-right"> {entry.content} </p>
+                            </li>
+                        }
+                        if (entry.role === "error") {
+                            return <li key={i}>
+                                <p className="border bg-red-100 border-red-500 text-red-800 px-1"> {entry.content} </p>
+                            </li>
+                        }
+                        if (entry.role === "assistant") {
+                            return <li key={i}>
+                                <ShowAssistantEntry entry={entry}/>
+                            </li>
+                        }
+                        return <></>
+                    })}
+
+                    <SearchBox search={search} />
+
+                    {(() => {
+                        if (loadState.state === "loading") {
+                            switch (loadState.phase) {
+                                case "semantic": return <p>Loading: Performing semantic search...</p>;
+                                case "prompt": return <p>Loading: Creating prompt...</p>;
+                                case "llm": return <p>Loading: Waiting for LLM...</p>;
+                            }
+                        } else if (loadState.state === "streaming") {
+                            return <ShowAssistantEntry entry={loadState.response}/>;
+                        }
+                        return <></>;
+                    })()}
+
+                </ul>
+            </main>
+        </>
+    );
+};
+
+export default Home;

web/src/pages/semantic.tsx

@@ -0,0 +1,100 @@
+const API_URL = process.env.NEXT_PUBLIC_API_URL || "http://127.0.0.1:3000";
+
+import { type NextPage } from "next";
+import React from "react";
+import Head from "next/head";
+import Header from "../header";
+import SearchBox from "../searchbox";
+import { useState } from "react";
+
+const Semantic: NextPage = () => {
+
+    const [results, setResults] = useState<SemanticEntry[]>([]);
+
+    const semantic_search = async (
+        query: string,
+        setQuery: (query: string) => void,
+        setLoading: (loading: boolean) => void
+    ) => {
+        
+        setLoading(true);
+        setQuery("");
+
+        const res = await fetch(API_URL + "/semantic", {
+            method: "POST",
+            headers: { "Content-Type": "application/json", "Access-Control-Allow-Origin": "*", },
+            body: JSON.stringify({query: query}),
+        })
+
+        if (!res.ok) {
+            setLoading(false);
+            console.log("load failure: " + res.status);
+        }
+
+        const data = await res.json();
+
+        setResults(data);
+        setLoading(false);
+
+    };
+
+    return (
+        <>
+            <Head>
+                <title>Alignment Search</title>
+            </Head>
+            <main>
+                <Header page="semantic" />
+                <h2>See the raw results of a semantic search</h2>
+                <SearchBox search={semantic_search} />
+                <ul>
+                    {results.map((entry, i) => (
+                        <li key={"entry" + i}>
+                            <ShowSemanticEntry entry={entry} />
+                        </li>
+                    ))}
+                </ul>
+            </main>
+        </>
+    );
+};
+
+// Round trip test. If this works, our heavier usecase probably will (famous last words)
+// The one real difference is we'll want to send back a series of results as we get
+// them back from OpenAI - I think we can just do this with a websocket, which
+// shouldn't be too much harder.
+
+type SemanticEntry = {
+    title: string;
+    author: string;
+    date: string;
+    url: string;
+    tags: string;
+    text: string;
+};
+
+const ShowSemanticEntry: React.FC<{entry: SemanticEntry}> = ({entry}) => {
+
+    return (
+        <div className="my-3">
+
+            {/* horizontally split first row, title on left, author on right */}
+            <div className="flex">
+                <h3 className="text-xl flex-1">{entry.title}</h3>
+                <p className="flex-1 text-right my-0">{entry.author} - {entry.date}</p>
+            </div>
+            { entry.text.split("\n").map((paragraph, i) => {
+                const p = paragraph.trim();
+                if (p === "") return <></>;
+                if (p === ".....") return <hr key={"b" + i} />;
+                return <p className="text-sm" key={"p" + i}> {paragraph} </p>
+              })
+            }
+
+            <a href={entry.url}>Read more</a>
+        </div>
+    );
+};
+
+
+export default Semantic;

web/src/searchbox.tsx

@@ -0,0 +1,52 @@
+import React from "react";
+import { useState, useEffect } from "react";
+import TextareaAutosize from 'react-textarea-autosize';
+
+const SearchBox: React.FC<{search: (
+    query: string,
+    setQuery: (query: string) => void,
+    setLoading: (loading: boolean) => void
+) => void}> = ({search}) => {
+
+    const [ query, setQuery ] = useState("");
+
+    const [ loading, setLoading ] = useState(false);
+
+    const inputRef = React.useRef<HTMLTextAreaElement>(null);
+
+
+    useEffect(() => {
+        // set focus on the input box
+        if (!loading) inputRef.current?.focus();
+    }, [loading]);
+
+    if (loading) return <></>;
+    return (<>
+        <form className="flex mb-2" onSubmit={async (e) => {
+            e.preventDefault();
+            search(query, setQuery, setLoading);
+        }}>
+
+            <TextareaAutosize
+                className="border border-gray-300 px-1 flex-1 resize-none"
+                ref={inputRef}
+                value={query}
+                onChange={(e) => setQuery(e.target.value)}
+                onKeyDown={(e) => {
+                    // if <esc>, blur the input box
+                    if (e.key === "Escape") e.currentTarget.blur();
+                    // if <enter> without <shift>, submit the form (if it's not empty)
+                    if (e.key === "Enter" && !e.shiftKey) {
+                        e.preventDefault();
+                        if (query.trim() !== "") search(query, setQuery, setLoading);
+                    }
+                }}
+            />
+            <button className="ml-2" type="submit" disabled={loading}>
+                {loading ? "Loading..." : "Search"}
+            </button>
+        </form>
+    </>);
+};
+
+export default SearchBox;

web/src/styles/globals.css

@@ -0,0 +1,39 @@
+@tailwind base;
+@tailwind components;
+@tailwind utilities;
+
+h1 {
+  @apply text-4xl font-bold my-4;
+}
+
+h2 {
+  @apply text-xl font-semibold my-4;
+}
+
+main {
+  @apply flex flex-col justify-center;
+  max-width: 800px;
+  margin: 0 auto;
+  padding: 0 2rem;
+  margin-top: 4rem;
+  margin-bottom: 4rem;
+}
+
+a {
+  @apply underline;
+}
+
+p {
+  @apply my-2;
+}
+
+button {
+  @apply bg-white hover:bg-gray-300 px-0.5 py-0 w-fit h-fit;
+  @apply border border-gray-300;
+  @apply text-gray-700;
+}
+
+ol {
+  @apply list-decimal;
+}
+

web/tailwind.config.cjs

@@ -0,0 +1,10 @@
+/** @type {import('tailwindcss').Config} */
+const config = {
+  content: ["./src/**/*.{js,ts,jsx,tsx}"],
+  theme: {
+    extend: {},
+  },
+  plugins: [],
+};
+
+module.exports = config;

web/tsconfig.json

@@ -0,0 +1,33 @@
+{
+  "compilerOptions": {
+    "target": "es2017",
+    "lib": ["dom", "dom.iterable", "esnext"],
+    "allowJs": true,
+    "checkJs": true,
+    "skipLibCheck": true,
+    "strict": true,
+    "forceConsistentCasingInFileNames": true,
+    "noEmit": true,
+    "esModuleInterop": true,
+    "module": "esnext",
+    "moduleResolution": "node",
+    "resolveJsonModule": true,
+    "isolatedModules": true,
+    "jsx": "preserve",
+    "incremental": true,
+    "noUncheckedIndexedAccess": true,
+    "baseUrl": ".",
+    "paths": {
+      "~/*": ["./src/*"]
+    }
+  },
+  "include": [
+    ".eslintrc.cjs",
+    "next-env.d.ts",
+    "**/*.ts",
+    "**/*.tsx",
+    "**/*.cjs",
+    "**/*.mjs"
+  ],
+  "exclude": ["node_modules"]
+}