Joint Tabular/Semantic QA over Tesla 10Q

In this example, we show how to ask questions over 10Q with understanding of both the unstructured text as well as embedded tables.

We use Unstructured to parse out the tables, and use LlamaIndex recursive retrieval to index/retrieve tables if necessary given the user question.

from pydantic import BaseModel
from unstructured.partition.html import partition_html
import pandas as pd

pd.set_option("display.max_rows", None)
pd.set_option("display.max_columns", None)
pd.set_option("display.width", None)
pd.set_option("display.max_colwidth", None)

Extract Elements

We use Unstructured to extract table and non-table elements from the 10-K filing.

from typing import Any, Optional
import pandas as pd


class Element(BaseModel):
    id: str
    type: str
    element: Any
    summary: Optional[str] = None
    table: Optional[pd.DataFrame] = None

    class Config:
        arbitrary_types_allowed = True
from lxml import html
import pandas as pd


def html_to_df(html_str):
    # print(html_str)
    tree = html.fromstring(html_str)
    # print(tree.xpath('//table'))
    table_element = tree.xpath("//table")[0]
    rows = table_element.xpath(".//tr")

    data = []
    for row in rows:
        cols = row.xpath(".//td")
        cols = [c.text.strip() if c.text is not None else "" for c in cols]
        data.append(cols)

    df = pd.DataFrame(data[1:], columns=data[0])
    return df
# simple heuristic to filter the table (if there's only one row or one column)
def filter_table(table_element):
    table_df = html_to_df(table_element.metadata.text_as_html)
    if len(table_df) <= 1 or len(table_df.columns) <= 1:
        return False
    else:
        return True
def extract_elements(filename, table_filters=[]):
    elements = partition_html(filename=filename)
    output_els = []
    for idx, element in enumerate(elements):
        if "unstructured.documents.html.HTMLTable" in str(type(element)):
            should_keep = all([tf(element) for tf in table_filters])
            if should_keep:
                table_df = html_to_df(str(element.metadata.text_as_html))
                output_els.append(
                    Element(
                        id=f"id_{idx}", type="table", element=element, table=table_df
                    )
                )
            else:
                pass
        else:
            output_els.append(Element(id=f"id_{idx}", type="text", element=element))
    return output_els
def get_table_elements(elements):
    return [e for e in elements if e.type == "table"]


def get_text_elements(elements):
    return [e for e in elements if e.type == "text"]
elements = extract_elements("tsla-20211231.htm", table_filters=[filter_table])
table_elements = get_table_elements(elements)
text_elements = get_text_elements(elements)
len(table_elements)
105

Summarize Tables

We specifically go through tables and use LlamaIndex to help extract a summary.

from llama_index import SummaryIndex, Document
from llama_index import ServiceContext
from llama_index.llms import OpenAI
from tqdm.notebook import tqdm

llm = OpenAI(model="gpt-4")

system_prompt = """\
You are an assistant designed to extract insights from messy tables in a financial report.

You are also designed to filter out "tables" that are not useful to keep. For instance, if the table \
is a wrongfully extracted piece of text, or does not contain any useful information.
"""

service_context = ServiceContext.from_defaults(system_prompt=system_prompt, llm=llm)
class TableOutput(BaseModel):
    """Output from analyzing a table."""

    summary: str
    should_keep: bool


def extract_table_summaries(elements):
    """Go through elements, extract out summaries that are tables."""
    for element in tqdm(elements):
        if element.type != "table":
            continue
        index = SummaryIndex.from_documents([Document(text=str(element.element))])
        query_engine = index.as_query_engine(output_cls=TableOutput)
        query_str = """\
What is this table about? Give a very concise summary (imagine you are adding a caption), \
and also output whether or not the table should be kept.
"""
        response = query_engine.query(query_str)
        # print(str(response))
        element.summary = response.response.summary
        # print(element.summary)
        # raise Exception
extract_table_summaries(table_elements)
table_elements[0].summary
'Delaware 91-2197729'
# [optional] save
import pickle
pickle.dump(elements, open("elements.pkl", "wb"))
# [optional] load

elements = pickle.load(open("elements.pkl", "rb"))

Setup Recursive Retriever

Now that we’ve extracted tables and their summaries, we can setup a recursive retriever in LlamaIndex to query these tables.

Create Nodes

from llama_index.schema import TextNode, IndexNode
from llama_index.node_parser import SimpleNodeParser
# join all non-table elements into
def _get_nodes_from_buffer(buffer, node_parser):
    doc = Document(text="\n\n".join([t for t in buffer]))
    nodes = node_parser.get_nodes_from_documents([doc])
    return nodes


def get_nodes_and_mappings(elements):
    pd.options.display.max_columns = None
    node_parser = SimpleNodeParser.from_defaults()

    nodes = []
    node_mappings = {}
    other_mappings = {}
    cur_text_el_buffer = []
    for element in elements:
        if element.type == "table":
            # flush text buffer
            if len(cur_text_el_buffer) > 0:
                cur_text_nodes = _get_nodes_from_buffer(cur_text_el_buffer, node_parser)
                nodes.extend(cur_text_nodes)
                cur_text_el_buffer = []

            index_node = IndexNode(
                text=str(element.summary), index_id=(element.id + "_table")
            )
            # print(str(element.table))
            table_df = element.table
            table_str = table_df.to_string()
            # node_mappings[(element.id + "_table")] = (TextNode(text=table_str), element.table, str(element.summary))
            node_mappings[(element.id + "_table")] = TextNode(text=table_str)
            other_mappings[(element.id + "_table")] = (
                element.table,
                str(element.summary),
            )
            nodes.append(index_node)
        else:
            cur_text_el_buffer.append(str(element.element))

    # flush text buffer
    if len(cur_text_el_buffer) > 0:
        cur_text_nodes = _get_nodes_from_buffer(cur_text_el_buffer, node_parser)
        nodes.extend(cur_text_nodes)
        cur_text_el_buffer = []

    return nodes, node_mappings, other_mappings
nodes, node_mappings, other_mappings = get_nodes_and_mappings(elements)
# # print(nodes[11].get_content())
# # # print(nodes[9].index_id)
# # print([n.get_content() if n
# tmp = str(node_mappings["id_1715_table"][0].get_content())
# print(tmp)
print(other_mappings["id_1715_table"])
(                                                                               \
0                               Year Ended December 31,                         
1                         2021                                         2020     
2    United States           $                           23,973                 
3            China                                       13,844                 
4            Other                                       16,006                 
5            Total           $                           53,823                 

                                                 
0                                                
1                 2019                           
2  $    15,207                $    12,653        
3        6,662                      2,979        
4        9,667                      8,946        
5  $    31,536                $    24,578        , 'Revenue by country for the years 2019, 2020, and 2021')

Construct Retrievers

from llama_index.retrievers import RecursiveRetriever
from llama_index.query_engine import RetrieverQueryEngine
from llama_index import VectorStoreIndex
# construct top-level vector index + query engine
vector_index = VectorStoreIndex(nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=1)
vector_query_engine = vector_index.as_query_engine(similarity_top_k=1)
from llama_index.retrievers import RecursiveRetriever

recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever},
    node_dict=node_mappings,
    verbose=True,
)
query_engine = RetrieverQueryEngine.from_args(recursive_retriever)

Run some Queries

response = query_engine.query("What was the revenue in 2020?")
print(str(response))
Retrieving with query id None: What was the revenue in 2020?
Retrieved node with id, entering: id_1715_table
Retrieving with query id id_1715_table: What was the revenue in 2020?
The revenue in 2020 was $31,536.
# compare against the baseline retriever
response = vector_query_engine.query("What was the revenue in 2020?")
print(str(response))
The revenue in 2020 was not provided in the context information.
response = query_engine.query("What were the total cash flows in 2021?")
Retrieving with query id None: What were the total cash flows in 2021?
Retrieved node with id, entering: id_558_table
Retrieving with query id id_558_table: What were the total cash flows in 2021?

print(str(response))
The total cash flows in 2021 were $11,497 million.
response = vector_query_engine.query("What were the total cash flows in 2021?")
print(str(response))
The total cash flows in 2021 were not provided in the given context information.
response = query_engine.query("What are the risk factors for Tesla?")
print(str(response))
Retrieving with query id None: What are the risk factors for Tesla?
Retrieving text node: Employees may leave Tesla or choose other employers over Tesla due to various factors, such as a very competitive labor market for talented individuals with automotive or technology experience, or any negative publicity related to us. In regions where we

19

have or will have operations, particularly significant engineering and manufacturing centers, there is strong competition for individuals with skillsets needed for our business, including specialized knowledge of electric vehicles, engineering and electrical and building construction expertise. Moreover, we may be impacted by perceptions relating to reductions in force that we have conducted in the past in order to optimize our organizational structure and reduce costs and the departure of certain senior personnel for various reasons. Likewise, as a result of our temporary suspension of various U.S. manufacturing operations in the first half of 2020, in April 2020, we temporarily furloughed certain hourly employees and reduced most salaried employees’ base salaries. We also compete with both mature and prosperous companies that have far greater financial resources than we do and start-ups and emerging companies that promise short-term growth opportunities.

Finally, our compensation philosophy for all of our personnel reflects our startup origins, with an emphasis on equity-based awards and benefits in order to closely align their incentives with the long-term interests of our stockholders. We periodically seek and obtain approval from our stockholders for future increases to the number of awards available under our equity incentive and employee stock purchase plans. If we are unable to obtain the requisite stockholder approvals for such future increases, we may have to expend additional cash to compensate our employees and our ability to retain and hire qualified personnel may be harmed.

We are highly dependent on the services of Elon Musk, Technoking of Tesla and our Chief Executive Officer.

We are highly dependent on the services of Elon Musk, Technoking of Tesla and our Chief Executive Officer. Although Mr. Musk spends significant time with Tesla and is highly active in our management, he does not devote his full time and attention to Tesla. Mr. Musk also currently serves as Chief Executive Officer and Chief Technical Officer of Space Exploration Technologies Corp., a developer and manufacturer of space launch vehicles, and is involved in other emerging technology ventures.

Our information technology systems or data, or those of our service providers or customers or users could be subject to cyber-attacks or other security incidents, which could result in data breaches, intellectual property theft, claims, litigation, regulatory investigations, significant liability, reputational damage and other adverse consequences.

We continue to expand our information technology systems as our operations grow, such as product data management, procurement, inventory management, production planning and execution, sales, service and logistics, dealer management, financial, tax and regulatory compliance systems. This includes the implementation of new internally developed systems and the deployment of such systems in the U.S. and abroad. While, we maintain information technology measures designed to protect us against intellectual property theft, data breaches, sabotage and other external or internal cyber-attacks or misappropriation, our systems and those of our service providers are potentially vulnerable to malware, ransomware, viruses, denial-of-service attacks, phishing attacks, social engineering, computer hacking, unauthorized access, exploitation of bugs, defects and vulnerabilities, breakdowns, damage, interruptions, system malfunctions, power outages, terrorism, acts of vandalism, security breaches, security incidents, inadvertent or intentional actions by employees or other third parties, and other cyber-attacks.

To the extent any security incident results in unauthorized access or damage to or acquisition, use, corruption, loss, destruction, alteration or dissemination of our data, including intellectual property and personal information, or our products or vehicles, or for it to be believed or reported that any of these occurred, it could disrupt our business, harm our reputation, compel us to comply with applicable data breach notification laws, subject us to time consuming, distracting and expensive litigation, regulatory investigation and oversight, mandatory corrective action, require us to verify the correctness of database contents, or otherwise subject us to liability under laws, regulations and contractual obligations, including those that protect the privacy and security of personal information. This could result in increased costs to us and result in significant legal and financial exposure and/or reputational harm.

We also rely on service providers, and similar incidents relating to their information technology systems could also have a material adverse effect on our business. There have been and may continue to be significant supply chain attacks. Our service providers, including our workforce management software provider, have been subject to ransomware and other security incidents, and we cannot guarantee that our or our service providers’ systems have not been breached or that they do not contain exploitable defects, bugs, or vulnerabilities that could result in a security incident, or other disruption to, our or our service providers’ systems. Our ability to monitor our service providers’ security measures is limited, and, in any event, malicious third parties may be able to circumvent those security measures.
The risk factors for Tesla include strong competition for skilled individuals in the labor market, negative publicity, potential impacts from reductions in force and departure of senior personnel, competition from companies with greater financial resources, dependence on the services of Elon Musk, potential cyber-attacks or security incidents, and reliance on service providers who may be vulnerable to security breaches. These factors could harm Tesla's ability to retain and hire qualified personnel, disrupt its business, harm its reputation, result in legal and financial exposure, and cause other adverse consequences.
response = vector_query_engine.query("What are the risk factors for Tesla?")
print(str(response))
The risk factors for Tesla include strong competition for skilled individuals in the labor market, negative publicity, potential impacts from reductions in force and departure of senior personnel, competition from companies with greater financial resources, dependence on the services of Elon Musk, potential cyber-attacks or security incidents, and reliance on service providers who may be vulnerable to security breaches. These factors could harm Tesla's ability to retain and hire qualified personnel, disrupt its business, harm its reputation, result in legal and financial exposure, and cause other adverse consequences.