Gladiator/gradient_dissent_bot

1import os2import re3from dataclasses import asdict4 5import pandas as pd6from langchain.callbacks import get_openai_callback7from langchain.chains import LLMChain8from langchain.chat_models import ChatOpenAI9from langchain.document_loaders import DataFrameLoader10from langchain.prompts import PromptTemplate11from langchain.text_splitter import TokenTextSplitter12from tqdm import tqdm13from wandb.integration.langchain import WandbTracer14 15import wandb16from config import config17 18 19def get_data(artifact_name: str, total_episodes: int = None):20    podcast_artifact = wandb.use_artifact(artifact_name, type="dataset")21    podcast_artifact_dir = podcast_artifact.download(config.root_artifact_dir)22    filename = artifact_name.split(":")[0].split("/")[-1]23    df = pd.read_csv(os.path.join(podcast_artifact_dir, f"{filename}.csv"))24    if total_episodes is not None:25        df = df.iloc[:total_episodes]26    return df27 28 29def extract_questions(episode_df: pd.DataFrame):30    # load docs into langchain format31    loader = DataFrameLoader(episode_df, page_content_column="transcript")32    data = loader.load()33 34    # split the documents35    text_splitter = TokenTextSplitter.from_tiktoken_encoder(chunk_size=1000, chunk_overlap=0)36    docs = text_splitter.split_documents(data)37    print(f"Number of documents for podcast {data[0].metadata['title']}: {len(docs)}")38 39    # initialize LLM40    llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)41 42    # define prompt43    prompt = """You are provided with a short transcript from a podcast episode.44    Your task is to extract the relevant and most important questions one might ask from the transcript and present them in a bullet-point list.45    Ensure that the total number of questions is no more than 3.46 47    TRANSCRIPT:48 49    {text}50 51    QUESTIONS:"""52 53    prompt_template = PromptTemplate(template=prompt, input_variables=["text"])54 55    pattern = r"\d+\.\s"56    que_by_llm = []57    for doc in docs:58        llm_chain = LLMChain(llm=llm, prompt=prompt_template)59        out = llm_chain.run(doc)60        cleaned_ques = re.sub(pattern, "", out).split("\n")61        que_by_llm.extend(cleaned_ques)62 63    return que_by_llm64 65 66if __name__ == "__main__":67    # initialize wandb tracer68    WandbTracer.init(69        {70            "project": config.project_name,71            "job_type": "extract_questions",72            "config": asdict(config),73        }74    )75 76    # get data77    df = get_data(artifact_name=config.summarized_data_artifact)78 79    questions = []80    with get_openai_callback() as cb:81        for episode in tqdm(82            df.iterrows(), total=len(df), desc="Extracting questions from episodes"83        ):84            episode_data = episode[1].to_frame().T85 86            episode_questions = extract_questions(episode_data)87            questions.append(episode_questions)88 89        print("*" * 25)90        print(cb)91        print("*" * 25)92 93        wandb.log(94            {95                "total_prompt_tokens": cb.prompt_tokens,96                "total_completion_tokens": cb.completion_tokens,97                "total_tokens": cb.total_tokens,98                "total_cost": cb.total_cost,99            }100        )101 102    df["questions"] = questions103 104    # log to wandb artifact105    path_to_save = os.path.join(config.root_data_dir, "summarized_que_podcasts.csv")106    df.to_csv(path_to_save, index=False)107    artifact = wandb.Artifact("summarized_que_podcasts", type="dataset")108    artifact.add_file(path_to_save)109    wandb.log_artifact(artifact)110 111    # create wandb table112    df["questions"] = df["questions"].apply(lambda x: "\n".join(x))113    table = wandb.Table(dataframe=df)114    wandb.log({"summarized_que_podcasts": table})115 116    WandbTracer.finish()117