Gladiator/gradient_dissent_bot

1import os2from dataclasses import asdict3 4import pandas as pd5from langchain.callbacks import get_openai_callback6from langchain.document_loaders import DataFrameLoader7from langchain.embeddings.openai import OpenAIEmbeddings8from langchain.text_splitter import TokenTextSplitter9from langchain.vectorstores import Chroma10from tqdm import tqdm11from wandb.integration.langchain import WandbTracer12 13import wandb14from config import config15 16 17def get_data(artifact_name: str, total_episodes=None):18    podcast_artifact = wandb.use_artifact(artifact_name, type="dataset")19    podcast_artifact_dir = podcast_artifact.download(config.root_artifact_dir)20    filename = artifact_name.split(":")[0].split("/")[-1]21    df = pd.read_csv(os.path.join(podcast_artifact_dir, f"{filename}.csv"))22    if total_episodes is not None:23        df = df.iloc[:total_episodes]24    return df25 26 27def create_embeddings(episode_df: pd.DataFrame, index: int):28    # load docs into langchain format29    loader = DataFrameLoader(episode_df, page_content_column="transcript")30    data = loader.load()31 32    # split the documents33    text_splitter = TokenTextSplitter.from_tiktoken_encoder(chunk_size=1000, chunk_overlap=0)34    docs = text_splitter.split_documents(data)35 36    title = data[0].metadata["title"]37    print(f"Number of documents for podcast {title}: {len(docs)}")38 39    # initialize embedding engine40    embeddings = OpenAIEmbeddings()41 42    db = Chroma.from_documents(43        docs,44        embeddings,45        persist_directory=os.path.join(config.root_data_dir / "chromadb", str(index)),46    )47    db.persist()48 49 50if __name__ == "__main__":51    # initialize wandb tracer52    WandbTracer.init(53        {54            "project": config.project_name,55            "job_type": "embed_transcripts",56            "config": asdict(config),57        }58    )59 60    # get data61    df = get_data(artifact_name=config.summarized_que_data_artifact)62 63    # create embeddings64    with get_openai_callback() as cb:65        for episode in tqdm(df.iterrows(), total=len(df), desc="Embedding transcripts"):66            episode_data = episode[1].to_frame().T67 68            create_embeddings(episode_data, index=episode[0])69 70        print("*" * 25)71        print(cb)72        print("*" * 25)73 74        wandb.log(75            {76                "total_prompt_tokens": cb.prompt_tokens,77                "total_completion_tokens": cb.completion_tokens,78                "total_tokens": cb.total_tokens,79                "total_cost": cb.total_cost,80            }81        )82 83    # log embeddings to wandb artifact84    artifact = wandb.Artifact("transcript_embeddings", type="dataset")85    artifact.add_dir(config.root_data_dir / "chromadb")86    wandb.log_artifact(artifact)87 88    WandbTracer.finish()89