RedTachyon/hot-ones-trivia

1import pandas as pd2from tqdm.auto import tqdm3import requests4import tiktoken5 6from typarse import BaseParser7from openai import OpenAI8import dotenv9 10import pickle11 12from core import get_batch_embeddings, Chunk, Dataset13 14 15class Parser(BaseParser):16    chunk_size: int = 400017    save_path: str = "dataset.pkl"18 19    _abbrev = {20        "chunk_size": "c",21        "save_path": "s",22    }23 24    _help = {25        "chunk_size": "The maximum number of tokens per chunk",26        "save_path": "The path to save the dataset",27    }28 29 30def get_youtube_title(url: str) -> str | None:31    """32    Get the title of a youtube video from the url33    """34    video_id = url.split("v=")[-1]35    api_url = f"https://www.youtube.com/oembed?url=http://www.youtube.com/watch?v={video_id}&format=json"36    response = requests.get(api_url)37    if response.status_code == 200:38        data = response.json()39        return data["title"]40    else:41        return None42 43 44def num_tokens_from_string(string: str, encoding_name: str) -> int:45    """46    Calculate the number of tokens in a string47    """48    encoding = tiktoken.get_encoding(encoding_name)49    num_tokens = len(encoding.encode(string))50    return num_tokens51 52 53def required_chunks(54    text: str, max_tokens: int = 8191, encoding_name: str = "cl100k_base"55) -> int:56    """57    Calculate the number of chunks required to split a text into chunks of a maximum number of tokens.58    """59    num_tokens = num_tokens_from_string(text, encoding_name)60    num_chunks = num_tokens // max_tokens61    if num_tokens % max_tokens != 0:62        num_chunks += 163    return num_chunks64 65 66def split_in_chunks(67    text: str, max_tokens: int = 8191, encoding_name: str = "cl100k_base"68) -> list[str]:69    """70    Split a long text into chunks of a maximum number of tokens71    """72    encoding = tiktoken.get_encoding(encoding_name)73    tokens = encoding.encode(text)74 75    chunks: list[str] = []76    current_chunk: list[int] = []77    current_chunk_size = 078 79    for token in tokens:80        if current_chunk_size + 1 > max_tokens:81            chunks.append(encoding.decode(current_chunk))82            current_chunk = []83            current_chunk_size = 084        current_chunk.append(token)85        current_chunk_size += 186 87    if current_chunk:88        chunks.append(encoding.decode(current_chunk))89 90    return chunks91 92 93if __name__ == "__main__":94    dotenv.load_dotenv()95 96    client = OpenAI()97    args = Parser()98 99    chunk_size = args.chunk_size100 101    links = pd.read_csv("links.csv").URL.tolist()102    titles = [get_youtube_title(link) for link in tqdm(links)]103 104    # Get all transcripts105    episodes = []106 107    for i in range(17):108        filename = f"transcripts/{i}.vtt"109        with open(filename, "r") as file:110            data = file.read()111            episodes.append(data)112 113    episode_chunks = [114        split_in_chunks(episode, max_tokens=chunk_size) for episode in episodes115    ]116 117    chunk_metadata = [118        Chunk(119            title=titles[i],120            video_idx=i,121            text=episode_chunks[i][j],122            link=links[i],123        )124        for i in range(17)125        for j in range(len(episode_chunks[i]))126    ]127 128    chunk_texts = [chunk.text for chunk in chunk_metadata]129 130    embeddings = get_batch_embeddings(client, chunk_texts)131 132    dataset = Dataset(chunks=chunk_metadata, embeddings=embeddings)133 134    with open(args.save_path, "wb") as file:135        pickle.dump(dataset, file)136