Question Answering Over Your Own Data With LlamaIndex and W&B

pip install llama-index

!pip install wandb

!pip install llama-index
!pip install wandb

import os
import json
import pandas as pd
from llama_index import Document, GPTVectorStoreIndex
from llama_index.node_parser import SimpleNodeParser
from llama_index.callbacks import LlamaDebugHandler, WandbCallbackHandler
from llama_index import (GPTVectorStoreIndex, ServiceContext, LLMPredictor, StorageContext)
from llama_index.callbacks import CallbackManager
from langchain.chat_models import ChatOpenAI
from llama_index import load_index_from_storage
import wandb
import openai

os.environ['OPENAI_API_KEY'] = '<your_open_ai_key>'

# Initialize a new wandb run
run_args = dict(
    project="llamaindex",
)
wandb_callback = WandbCallbackHandler(run_args=run_args)
llama_debug = LlamaDebugHandler(print_trace_on_end=True)

# Configure LLM Predictor
llm_predictor = LLMPredictor(llm=ChatOpenAI(model_name='gpt-4', temperature=0))
callback_manager = CallbackManager([llama_debug, wandb_callback])
service_context = ServiceContext.from_defaults(callback_manager=callback_manager, llm_predictor=llm_predictor)

# Load JSON file
with open('/kaggle/input/stanford-question-answering-dataset/train-v1.1.json', 'r') as f:
    squad_dict = json.load(f)

# Parse JSON file and Create DataFrame
id_list, context_list, title_list = [], [], []
for topic in squad_dict['data']:
    for para in topic['paragraphs']:
        context = para['context']
        for qa in para['qas']:
            id_list.append(qa['id'])
            context_list.append(context)
            title_list.append(topic['title'])
data = pd.DataFrame({'id': id_list, 'context': context_list, 'title': title_list})

documents = []
for i, row in data.iterrows():
    # Log the progress
    wandb.log({"Document Index": i})

    document = Document(
        text=row['context'],  # Use 'text' instead of 'content'
        doc_id=row['id'],
        extra_info={"title": row['title'],}
    )
    documents.append(document)

parser = SimpleNodeParser()
nodes = parser.get_nodes_from_documents(documents)

index = GPTVectorStoreIndex(nodes, service_context=service_context)
wandb_callback.persist_index(index, index_name="simple_vector_store")

index.storage_context.persist(persist_dir="/kaggle/working/index")
storage_context = wandb_callback.load_storage_context(artifact_url=[entity]/[project]/[artifact-name]:[version])

index = load_index_from_storage(storage_context, service_context=service_context)

query_engine = index.as_query_engine()
response = query_engine.query("in what year was the college of engineering established?")
print(response, sep="\n")

response = query_engine.query("What is the name of the manifestation of the Virgin Mary that appeared in 1531?")
print(response, sep="\n")