Gladiator/gradient_dissent_bot

1import os2from dataclasses import asdict3 4import pandas as pd5from langchain.callbacks import get_openai_callback6from langchain.chains.summarize import load_summarize_chain7from langchain.chat_models import ChatOpenAI8from langchain.document_loaders import DataFrameLoader9from langchain.prompts import PromptTemplate10from langchain.text_splitter import TokenTextSplitter11from tqdm import tqdm12from wandb.integration.langchain import WandbTracer13 14import wandb15from config import config16 17 18def get_data(artifact_name: str, total_episodes: int = None):19    podcast_artifact = wandb.use_artifact(artifact_name, type="dataset")20    podcast_artifact_dir = podcast_artifact.download(config.root_artifact_dir)21    filename = artifact_name.split(":")[0].split("/")[-1]22    df = pd.read_csv(os.path.join(podcast_artifact_dir, f"{filename}.csv"))23    if total_episodes is not None:24        df = df.iloc[:total_episodes]25    return df26 27 28def summarize_episode(episode_df: pd.DataFrame):29    # load docs into langchain format30    loader = DataFrameLoader(episode_df, page_content_column="transcript")31    data = loader.load()32 33    # split the documents34    text_splitter = TokenTextSplitter.from_tiktoken_encoder(chunk_size=1000, chunk_overlap=0)35    docs = text_splitter.split_documents(data)36    print(f"Number of documents for podcast {data[0].metadata['title']}: {len(docs)}")37 38    # initialize LLM39    llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)40 41    # define map prompt42    map_prompt = """Write a concise summary of the following short transcript from a podcast.43    Don't add your opinions or interpretations.44 45    {text}46 47    CONCISE SUMMARY:"""48 49    # define combine prompt50    combine_prompt = """You have been provided with summaries of chunks of transcripts from a podcast.51    Your task is to merge these intermediate summaries to create a brief and comprehensive summary of the entire podcast.52    The summary should encompass all the crucial points of the podcast.53    Ensure that the summary is atleast 2 paragraph long and effectively captures the essence of the podcast.54    {text}55 56    SUMMARY:"""57 58    map_prompt_template = PromptTemplate(template=map_prompt, input_variables=["text"])59    combine_prompt_template = PromptTemplate(template=combine_prompt, input_variables=["text"])60 61    # initialize the summarizer chain62    chain = load_summarize_chain(63        llm,64        chain_type="map_reduce",65        return_intermediate_steps=True,66        map_prompt=map_prompt_template,67        combine_prompt=combine_prompt_template,68    )69 70    summary = chain({"input_documents": docs})71    return summary72 73 74if __name__ == "__main__":75    # initialize wandb tracer76    WandbTracer.init(77        {78            "project": config.project_name,79            "job_type": "summarize",80            "config": asdict(config),81        }82    )83 84    # get scraped data85    df = get_data(artifact_name=config.yt_podcast_data_artifact)86 87    summaries = []88    with get_openai_callback() as cb:89        for episode in tqdm(df.iterrows(), total=len(df), desc="Summarizing episodes"):90            episode_data = episode[1].to_frame().T91 92            summary = summarize_episode(episode_data)93            summaries.append(summary["output_text"])94 95        print("*" * 25)96        print(cb)97        print("*" * 25)98 99        wandb.log(100            {101                "total_prompt_tokens": cb.prompt_tokens,102                "total_completion_tokens": cb.completion_tokens,103                "total_tokens": cb.total_tokens,104                "total_cost": cb.total_cost,105            }106        )107 108    df["summary"] = summaries109 110    # save data111    path_to_save = os.path.join(config.root_data_dir, "summarized_podcasts.csv")112    df.to_csv(path_to_save, index=False)113 114    # log to wandb artifact115    artifact = wandb.Artifact("summarized_podcasts", type="dataset")116    artifact.add_file(path_to_save)117    wandb.log_artifact(artifact)118 119    # create wandb table120    table = wandb.Table(dataframe=df)121    wandb.log({"summarized_podcasts": table})122 123    WandbTracer.finish()124