spark-nlp/TAPAS

1import streamlit as st2import sparknlp3import pandas as pd4import json5 6from sparknlp.base import *7from sparknlp.annotator import *8from pyspark.ml import Pipeline9from sparknlp.pretrained import PretrainedPipeline10 11# Page configuration12st.set_page_config(13    layout="wide", 14    initial_sidebar_state="auto"15)16 17# CSS for styling18st.markdown("""19    <style>20        .main-title {21            font-size: 36px;22            color: #4A90E2;23            font-weight: bold;24            text-align: center;25        }26        .section {27            background-color: #f9f9f9;28            padding: 10px;29            border-radius: 10px;30            margin-top: 10px;31        }32        .section p, .section ul {33            color: #666666;34        }35    </style>36""", unsafe_allow_html=True)37 38@st.cache_resource39def init_spark():40    return sparknlp.start()41 42@st.cache_resource43def create_pipeline(model):44    document_assembler = MultiDocumentAssembler() \45        .setInputCols("table_json", "questions") \46        .setOutputCols("document_table", "document_questions")47 48    sentence_detector = SentenceDetector() \49        .setInputCols(["document_questions"]) \50        .setOutputCol("questions")51 52    table_assembler = TableAssembler()\53        .setInputCols(["document_table"])\54        .setOutputCol("table")55 56    tapas_wtq = TapasForQuestionAnswering\57        .pretrained("table_qa_tapas_base_finetuned_wtq", "en")\58        .setInputCols(["questions", "table"])\59        .setOutputCol("answers_wtq")60    61    tapas_sqa = TapasForQuestionAnswering\62        .pretrained("table_qa_tapas_base_finetuned_sqa", "en")\63        .setInputCols(["questions", "table"])\64        .setOutputCol("answers_sqa")65 66    pipeline = Pipeline(stages=[document_assembler, sentence_detector, table_assembler, tapas_wtq, tapas_sqa])67    return pipeline68 69def fit_data(pipeline, json_data, question):70    spark_df = spark.createDataFrame([[json_data, question]]).toDF("table_json", "questions")71    model = pipeline.fit(spark_df)72    res = model.transform(spark_df)73    return res.select("answers_wtq.result", "answers_sqa.result").collect()74 75# Sidebar content76model = st.sidebar.selectbox(77    "Choose the pretrained model",78    ["table_qa_tapas_base_finetuned_wtq", "table_qa_tapas_base_finetuned_sqa"],79    help="For more info about the models visit: https://sparknlp.org/models"80)81 82# Set up the page layout83title = 'TAPAS for Table-Based Question Answering with Spark NLP'84sub_title = ("""85TAPAS (Table Parsing Supervised via Pre-trained Language Models) enhances the BERT architecture to effectively process tabular data, allowing it to answer complex questions about tables without needing to convert them into text.<br>86<br>87<strong>table_qa_tapas_base_finetuned_wtq:</strong> This model excels at answering questions that require aggregating data across the entire table, such as calculating sums or averages.<br>88<strong>table_qa_tapas_base_finetuned_sqa:</strong> This model is designed for sequential question-answering tasks where the answer to each question may depend on the context provided by previous answers.89""")90 91st.markdown(f'<div class="main-title">{title}</div>', unsafe_allow_html=True)92st.markdown(f'<div class="section"><p>{sub_title}</p></div>', unsafe_allow_html=True)93 94# Reference notebook link in sidebar95link = """96<a href="https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/streamlit_notebooks/NER_HINDI_ENGLISH.ipynb">97    <img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/>98</a>99"""100st.sidebar.markdown('Reference notebook:')101st.sidebar.markdown(link, unsafe_allow_html=True)102 103# Define the JSON data for the table104# New JSON data105json_data = '''106{107  "header": ["name", "net_worth", "age", "nationality", "company", "industry"],108  "rows": [109    ["Elon Musk", "$200,000,000,000", "52", "American", "Tesla, SpaceX", "Automotive, Aerospace"],110    ["Jeff Bezos", "$150,000,000,000", "60", "American", "Amazon", "E-commerce"],111    ["Bernard Arnault", "$210,000,000,000", "74", "French", "LVMH", "Luxury Goods"],112    ["Bill Gates", "$120,000,000,000", "68", "American", "Microsoft", "Technology"],113    ["Warren Buffett", "$110,000,000,000", "93", "American", "Berkshire Hathaway", "Conglomerate"],114    ["Larry Page", "$100,000,000,000", "51", "American", "Google", "Technology"],115    ["Mark Zuckerberg", "$85,000,000,000", "40", "American", "Meta", "Social Media"],116    ["Mukesh Ambani", "$80,000,000,000", "67", "Indian", "Reliance Industries", "Conglomerate"],117    ["Alice Walton", "$65,000,000,000", "74", "American", "Walmart", "Retail"],118    ["Francoise Bettencourt Meyers", "$70,000,000,000", "70", "French", "L'Oreal", "Cosmetics"],119    ["Amancio Ortega", "$75,000,000,000", "88", "Spanish", "Inditex (Zara)", "Retail"],120    ["Carlos Slim", "$55,000,000,000", "84", "Mexican", "America Movil", "Telecom"]121  ]122}123'''124 125# Define queries for selection126queries = [127    "Who has a higher net worth, Bernard Arnault or Jeff Bezos?",128    "List the top three individuals by net worth.",129    "Who is the richest person in the technology industry?",130    "Which company in the e-commerce industry has the highest net worth?",131    "Who is the oldest billionaire on the list?",132    "Which individual under the age of 60 has the highest net worth?",133    "Who is the wealthiest American, and which company do they own?",134    "Find all French billionaires and list their companies.",135    "How many women are on the list, and what are their total net worths?",136    "Who is the wealthiest non-American on the list?",137    "Find the person who is the youngest and has a net worth over $100 billion.",138    "Who owns companies in more than one industry, and what are those industries?",139    "What is the total net worth of all individuals over 70?",140    "How many billionaires are in the conglomerate industry?"141]142 143# Load the JSON data into a DataFrame and display it144table_data = json.loads(json_data)145df_table = pd.DataFrame(table_data["rows"], columns=table_data["header"])146df_table.index += 1147 148st.write("")149st.write("Context DataFrame (Click To Edit)")150edited_df = st.data_editor(df_table)151 152# Convert edited DataFrame back to JSON format153table_json_data = {154    "header": edited_df.columns.tolist(),155    "rows": edited_df.values.tolist()156}157table_json_str = json.dumps(table_json_data)158 159# User input for questions160selected_text = st.selectbox("Question Query", queries)161custom_input = st.text_input("Try it with your own Question!")162text_to_analyze = custom_input if custom_input else selected_text163 164# Initialize Spark and create the pipeline165spark = init_spark()166pipeline = create_pipeline(model)167 168# Run the pipeline with the selected query and the converted table data169output = fit_data(pipeline, table_json_str, text_to_analyze)170 171# Display the output172st.markdown("---")173st.subheader("Processed Output")174 175# Check if output is available176if output:177    # Extract and Display results178    results_wtq = output[0][0] if output[0][0] else "No results found."179    results_sqa = output[0][1] if output[0][1] else "No results found."180    st.markdown(f"**Answers from WTQ model:** {', '.join(results_wtq)}")181    st.markdown(f"**Answers from SQA model:** {', '.join(results_sqa)}")