Intro

This notebook will show you an example of how to use DocstringWalker from Llama Hub, combined with Llama Index and LLM of your choice.

Lib install for Collab

python

!pip install llama_index

python

!pip install llama_hub

For this exercise we will use PyTorch Geometric (PyG) module for inspecting multi-module doctstrings.

python

!pip install torch_geometric

Lib imports

python

import os

from pprint import pprint

from llama_index import (
    ServiceContext,
    VectorStoreIndex,
    SummaryIndex,
)

import llama_hub.docstring_walker as docstring_walker

Example 1 - reading Docstring Walker's own docstrings

Let's start by using it.... on itself :) We will see what information gets extracted from the module.

python

# Step 1 - create DocstringWalker object
walker = docstring_walker.DocstringWalker()

# Step 2 - prepare path to module
path_to_docstring_walker = os.path.dirname(docstring_walker.__file__)

# Step 3 - load documents from docstrings
example1_docs = walker.load_data(path_to_docstring_walker)

python

print(example1_docs[0].text)

Now we can use the doc to generate Llama index and use it with LLM.

python

# Step 1 - create vector store index
example1_index = VectorStoreIndex(example1_docs)

# Step 2 - turn vector store into the query engine
example1_query_engine = example1_index.as_query_engine()

python

pprint(
    example1_query_engine.query("What is the main purpose of DocstringWalker?").response
)

python

print(
    example1_query_engine.query(
        "What are the main functions used in DocstringWalker. Use numbered list, briefly describe each function."
    ).response
)

Example 2 - checking multi-module project

Now we can use the same approach to check a multi-module project. Let's use PyTorch Geometric (PyG) Knowledge Graph (KG) module for this exercise.

python

import torch_geometric.nn.kge as kge

path_to_module = os.path.dirname(kge.__file__)
example2_docs = walker.load_data(path_to_module)

python

example2_index = SummaryIndex(example2_docs)
example2_docs = example2_index.as_query_engine()

python

print(
    example2_docs.query(
        "What classes are available and what is their main purpose? Use nested numbered list to describe: the class name, short summary of purpose, papers or literature review for each one of them"
    ).response
)

python

print(example2_docs.query("What are the parameters required by TransE class?").response)