stampy-chat/src/dataset.py

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 tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential,
)  # for exponential backoff

from text_splitter import TokenSplitter, split_into_sentences
from settings import PATH_TO_DATA, PATH_TO_EMBEDDINGS, PATH_TO_DATASET, EMBEDDING_MODEL, LEN_EMBEDDINGS

import openai
openai.api_key = os.environ.get('OPENAI_API_KEY')

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 the dataset .jsonl file.
            custom_sources: List[str] = None,  # List of sources to include, like "alignment forum", "lesswrong", "arxiv",etc.
            rate_limit_per_minute: int = 3_500,  # Rate limit for the OpenAI API.
            min_tokens_per_block: int = 400, # Minimum number of tokens per block.
            max_tokens_per_block: int = 600, # Maximum number of tokens per block.
            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 of an article, its URL, and text. E.g.: [('title', 'url', 'text'), ...]
        self.embedding_strings: List[str] = []  # List of strings, each being a few paragraphs from a single article (not exceeding 1000 words).
        self.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
            
        # 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 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.")
                    
                    random_number = random.random()
                    if random_number > self.fraction_of_articles_to_use:
                        continue
                    
                    # 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
                    
                    self.articles_count[entry['source']] += 1
                    self.total_articles_count += 1
                    
                    # Get title, author, date, URL, tags, and text
                    title, author, date_published, url, tags, text = self.extract_info_from_article(entry)
                                                            
                    # 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] * 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):
        # Get an embedding for each text, with retries if necessary
        #TODO: check batch size stuff at https://github.com/openai/openai-cookbook/blob/main/examples/vector_databases/pinecone/Gen_QA.ipynb
        # to speed up the process 

        @retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(5))
        def get_embedding_at_index(text: str, i: int, delay_in_seconds: float = 0) -> np.ndarray:
            time.sleep(delay_in_seconds)
            embedding = openai.Embedding.create(
                model=EMBEDDING_MODEL, 
                input=text
            )
            return i, embedding["data"][0]["embedding"]
        
        start = time.time()
        self.embeddings = np.zeros((len(self.embedding_strings), LEN_EMBEDDINGS))
        
        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = [executor.submit(get_embedding_at_index, text, i) for i, text in enumerate(self.embedding_strings)]
            num_completed = 0
            for future in concurrent.futures.as_completed(futures):
                i, embedding = future.result()
                self.embeddings[i] = embedding
                num_completed += 1
                if num_completed % 50 == 0:
                    print(f"Completed {num_completed}/{len(self.embedding_strings)} embeddings in {time.time() - start:.2f} seconds.")
        print(f"Completed {num_completed}/{len(self.embedding_strings)} embeddings in {time.time() - start:.2f} seconds.")

    #TODO: complete this to speed up embeddings
    """ def get_embeddings_in_batches(self):
        # Get an embedding for each text, with retries if necessary
        
        @retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(5))
        def get_embedding_in_batches(batch: List[str], i: int, delay_in_seconds: float = 0) -> np.ndarray:
                try:
                    res = openai.Embedding.create(input=batch, engine=EMBEDDING_MODEL)
                except:
                    done = False
                    while not done:
                        time.sleep(5)
                        try:
                            res = openai.Embedding.create(input=batch, engine=EMBEDDING_MODEL)
                            done = True
                        except:
                            pass
                """

            

    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):
        with open(path, 'wb') as f:
            pickle.dump(self, 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





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_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("embeddings.npy")
    
    # dataset.save_class("dataset.pkl")
    # dataset = pickle.load(open("dataset.pkl", "rb"))