W&B 문서를 위한 질의응답 봇 만들기

PROJECT = "wandb_docs_bot"
run = wandb.init(project=PROJECT)

def download_raw_dataset():
    dataset_artifact_path = 'parambharat/wandb_docs_bot/docs_dataset:latest'
    artifact = run.use_artifact(dataset_artifact_path, type='dataset')
    artifact_path = artifact.get_path("wandb_docs.json")
    file = artifact_path.download()
    return file

import json
from langchain.docstore.document import Document
from langchain.text_splitter import CharacterTextSplitter

def load_documents(fname):
    source_chunks = []
    splitter = CharacterTextSplitter(separator=" ", chunk_size=1024, chunk_overlap=0)
    for line in open(fname, "r"):
        line = json.loads(line)
        for chunk in splitter.split_text(line["reference"]):
            source_chunks.append(Document(page_content=chunk, metadata={"source": line["source"]}))
    return source_chunks

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores.faiss import FAISS

def create_and_save_index(documents):
    store = FAISS.from_documents(documents,OpenAIEmbeddings())
    artifact = wandb.Artifact("faiss_store", type="search_index")
    faiss.write_index(store.index, "docs.index")
    artifact.add_file("docs.index")
    store.index = None
    with artifact.new_file("faiss_store.pkl", "wb") as f:
        pickle.dump(store, f)
    wandb.log_artifact(artifact, "docs_index", type="embeddings_index")
    return store

def load_prompt():
  dataset_artifact_path = 'parambharat/wandb_docs_bot/docs_dataset:latest'
  artifact = run.use_artifact(dataset_artifact_path, type='dataset')
  artifact_path = artifact.get_path("combine_prompt.txt")
  file = artifact_path.download()
  prompt_template = (open(file, "r").read())
  prompt = PromptTemplate(input_variables=["question", "summaries"],
                          template=prompt_template)
  return prompt

def load_chain(openai_api_key):
  if validate_openai_key(openai_api_key):
    vectorstore = load_vectostore()
    prompt = load_prompt()
    chain = VectorDBQAWithSourcesChain.from_chain_type(
      llm=OpenAI(temperature=0, openai_api_key=openai_api_key),
      chain_type="map_reduce",
      vectorstore=vectorstore,
      combine_prompt=prompt,)
    return chain

def get_answer(question, chain):
  if chain is not None:
    result = chain(
      {
        "question": question,
      },
      return_only_outputs=True,
    )
    response = f"Answer:\t{result['answer']}\n\nSources:\t{result['sources']}\n"
    return response

class Chat:
  def __init__(self):
    self.chain = None

  def __call__(self, message, history, openai_api_key):
    if self.chain is None:
      self.chain = load_chain(openai_api_key)

    history = history or []
    message = message.lower()
    response = get_answer(message, self.chain)
    if response is None:
      response = "Please enter a valid Openai API Key and try again. "
    history.append((message, response))
    return history, history

with gr.Blocks() as demo:
  with gr.Row():
    question = gr.Textbox(
      label='Type in your questions about wandb here and press Enter!',
      placeholder='How do i log images with wandb ?')
    openai_api_key = gr.Textbox(type='password',
                                label="Enter your OpenAI API key here",)
  state = gr.State()
  chatbot = gr.Chatbot()
  question.submit(Chat(), [question, state, openai_api_key], [chatbot, state])