Update Demo.py

Demo.py

@@ -53,22 +53,24 @@ def create_pipeline(model):
         .setInputCols(["document_table"])\
         .setOutputCol("table")
 
-    tapas = TapasForQuestionAnswering\
-        .pretrained(model, "en")\
+    tapas_wtq = TapasForQuestionAnswering\
+        .pretrained("table_qa_tapas_base_finetuned_wtq", "en")\
         .setInputCols(["questions", "table"])\
         .setOutputCol("answers_wtq")
+    
+    tapas_sqa = TapasForQuestionAnswering\
+        .pretrained("table_qa_tapas_base_finetuned_sqa", "en")\
+        .setInputCols(["questions", "table"])\
+        .setOutputCol("answers_sqa")
 
-    pipeline = Pipeline(stages=[document_assembler, sentence_detector, table_assembler, tapas])
+    pipeline = Pipeline(stages=[document_assembler, sentence_detector, table_assembler, tapas_wtq, tapas_sqa])
     return pipeline
 
 def fit_data(pipeline, json_data, question):
     spark_df = spark.createDataFrame([[json_data, question]]).toDF("table_json", "questions")
     model = pipeline.fit(spark_df)
-    lightPipelineModel = LightPipeline(model)
-    table, data = f""" table_json:{json_data} """ , """ questions:{question} """
-    result = lightPipelineModel.fullAnnotate([table, data])
-    st.write(result)
-    return result
+    result = model.transform(spark_df)
+    return str(result.select("answers_wtq.result", "answers_sqa.result"))
 
 # Sidebar content
 model = st.sidebar.selectbox(
@@ -170,6 +172,7 @@ output = fit_data(pipeline, table_json_str, text_to_analyze)
 st.markdown("---")
 st.subheader("Processed Output")
 
+output
 # # Check if output is available
 # if output:
 #     results_wtq = output[0][0] if output[0][0] else "No results found."