Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import pandas as pd
+import plotly.express as px
 from transformers import pipeline
 
 # Set the page layout for Streamlit
@@ -70,16 +71,31 @@ else:
                     st.write("TAPAS Raw Output (Response):")
                     st.write(result)  # This will display the raw output from TAPAS
 
-                    # Optionally, you can output the raw output as plain text:
-                    st.text("Raw TAPAS Output (Plain Text):")
-                    st.text(str(result))  # This will display raw output as plain text
-
-                    # Check if TAPAS is returning the expected answer
-                    answer = result.get('answer', None)
-                    if answer:
-                        st.write(f"TAPAS Answer: {answer}")
-                    else:
-                        st.warning("TAPAS did not return a valid answer.")
+                    # If the user asked for a count of a column or specific data:
+                    if "count" in question.lower():
+                        # Ask TAPAS to count rows of a specific column
+                        column_name = question.split("count")[-1].strip()  # Extract column name
+                        if column_name in df.columns:
+                            count_result = tqa(table=df, query=f"count of {column_name}")
+                            st.write(f"Count for column '{column_name}': {count_result['answer']}")
+                        else:
+                            st.warning(f"Column '{column_name}' not found in the dataset.")
+                    
+                    elif isinstance(result.get("answer"), list):
+                        # Handle structured data for graphing (e.g., scatter plot or other visualizations)
+                        answer_data = result["answer"]
+                        if answer_data and isinstance(answer_data, list) and isinstance(answer_data[0], dict):
+                            # Extract column data for x and y axes for Plotly
+                            x_data = [item.get("column1") for item in answer_data]  # Replace column1 with actual column name
+                            y_data = [item.get("column2") for item in answer_data]  # Replace column2 with actual column name
+                            
+                            # Create a scatter plot using Plotly
+                            fig = px.scatter(x=x_data, y=y_data, title="Scatter Plot based on TAPAS Answer")
+                            st.plotly_chart(fig, use_container_width=True)
+
+                    elif isinstance(result.get("answer"), str):
+                        # Handle simple answers (e.g., sums, counts, etc.)
+                        st.write(f"TAPAS Answer: {result['answer']}")
 
             except Exception as e:
                 st.warning(f"Error processing question or generating answer: {str(e)}")