Update app.py

app.py

@@ -1,26 +1,37 @@
+import os
 import streamlit as st
+from st_aggrid import AgGrid
 import pandas as pd
+from transformers import pipeline, T5ForConditionalGeneration, T5Tokenizer
 import plotly.express as px
-from transformers import pipeline
 
 # Set the page layout for Streamlit
 st.set_page_config(layout="wide")
 
-# Initialize TAPAS pipeline for table-based question answering (multilingual)
+# Initialize TAPAS pipeline
 tqa = pipeline(task="table-question-answering", 
               model="google/tapas-large-finetuned-wtq",
-              device=0)  # Assuming GPU is available, otherwise set device="cpu"
+              device="cpu")
+
+# Initialize T5 tokenizer and model for text generation
+t5_tokenizer = T5Tokenizer.from_pretrained("t5-small")
+t5_model = T5ForConditionalGeneration.from_pretrained("t5-small")
 
 # Title and Introduction
-st.title("Data Table with TAPAS NLP Integration")
+st.title("Table Question Answering and Data Analysis App")
 st.markdown(""" 
-    This app allows you to upload a table (CSV or Excel) and ask questions to extract information from the data.
-    Using **TAPAS**, the app can interpret your questions and provide the corresponding answers.
+    This app allows you to upload a table (CSV or Excel) and ask questions about the data.
+    Based on your question, it will provide the corresponding answer using the **TAPAS** model and additional data processing.
 
     ### Available Features:
-    - **Table Question Answering**: Ask questions related to the uploaded table.
+    - **mean()**: For "average", it computes the mean of the entire numeric DataFrame.
+    - **sum()**: For "sum", it calculates the sum of all numeric values in the DataFrame.
+    - **max()**: For "max", it computes the maximum value in the DataFrame.
+    - **min()**: For "min", it computes the minimum value in the DataFrame.
+    - **count()**: For "count", it counts the non-null values in the entire DataFrame.
+    - **Graph Generation**: You can ask questions like "make a graph of column sales?" or "make a graph between sales and expenses?". The app will generate interactive graphs for you.
     
-    Upload your data and ask questions to extract answers.
+    Upload your data and ask questions to get both answers and visualizations.
 """)
 
 # File uploader in the sidebar
@@ -41,61 +52,107 @@ else:
             df = None
 
         if df is not None:
-            # Convert object columns to numeric where possible
-            df = df.apply(pd.to_numeric, errors='ignore')
+            numeric_columns = df.select_dtypes(include=['object']).columns
+            for col in numeric_columns:
+                df[col] = pd.to_numeric(df[col], errors='ignore')
 
             st.write("Original Data:")
             st.write(df)
 
-            # Display a sample of data for user reference
-            st.write("Sample data:")
-            st.write(df.head())
-
+            df_numeric = df.copy()
+            df = df.astype(str)
+
+            # Display the first 5 rows of the dataframe in an editable grid
+            grid_response = AgGrid(
+                df.head(5),
+                fit_columns_on_grid_load=True,  # Correct parameter to fit columns on grid load
+                editable=True, 
+                height=300, 
+                width='100%',
+            )
+            
     except Exception as e:
         st.error(f"Error reading file: {str(e)}")
 
     # User input for the question
-    question = st.text_input(f'Ask your question related to the table')
+    question = st.text_input('Type your question')
+
+    # Check if the question is about generating a graph
+    is_graph_query = False
+    is_count_query = False
 
+    # Check if the question contains "count"
+    if 'count' in question.lower():
+        is_count_query = True
+    elif 'graph' in question.lower():
+        is_graph_query = True
+
+    # Process the answer using TAPAS and T5
     with st.spinner():
-        if st.button('Get Answer'):
+        if st.button('Answer'):
             try:
-                # Ensure the question is a valid string
-                if not question or not isinstance(question, str):
-                    st.error("Please enter a valid question.")
-                else:
-                    # Use TAPAS model to process the question
-                    result = tqa(table=df, query=question)
+                if not is_graph_query:
+                    # Process TAPAS-related questions if it's not a graph query
+                    raw_answer = tqa(table=df, query=question, truncation=True)
+
+                    # Display raw answer from TAPAS on the screen
+                    st.markdown("<p style='font-family:sans-serif;font-size: 1rem;'>Raw TAPAS Answer: </p>", unsafe_allow_html=True)
+                    st.write(raw_answer)  # Display the raw TAPAS output
+
+                    # Extract relevant values for Plotly
+                    answer = raw_answer.get('answer', '')
+                    coordinates = raw_answer.get('coordinates', [])
+                    cells = raw_answer.get('cells', [])
+
+                    st.markdown("<p style='font-family:sans-serif;font-size: 1rem;'>Relevant Data for Plotly: </p>", unsafe_allow_html=True)
+                    st.write(f"Answer: {answer}")
+                    st.write(f"Coordinates: {coordinates}")
+                    st.write(f"Cells: {cells}")
+
+                    # If TAPAS is returning a list of numbers for "average" like you mentioned
+                    if "average" in question.lower() and cells:
+                        # Assuming cells are numeric values that can be plotted in a graph
+                        plot_data = [float(cell) for cell in cells]  # Convert cells to numeric data
 
-                    # Display the raw output from TAPAS
-                    st.write("TAPAS Raw Output (Response):")
-                    st.write(result)  # This will display the raw output from TAPAS
+                        # Create a DataFrame for Plotly
+                        plot_df = pd.DataFrame({ 'Index': list(range(1, len(plot_data) + 1)), 'Value': plot_data })
+
+                        # Generate a graph using Plotly
+                        fig = px.line(plot_df, x='Index', y='Value', title=f"Graph for '{question}'")
+                        st.plotly_chart(fig, use_container_width=True)
+
+                    else:
+                        st.write(f"No data to plot for the question: '{question}'")
+
+                else:
+                    # Handle graph-related questions
+                    if is_count_query:
+                        # Extract the column name to count
+                        column_name = question.split('count')[-1].strip()
 
-                    # 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:
+                            # Ask TAPAS to count the rows for this specific column
                             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")
+                    elif 'between' in question.lower() and 'and' in question.lower():
+                        columns = question.split('between')[-1].split('and')
+                        columns = [col.strip() for col in columns]
+                        if len(columns) == 2 and all(col in df.columns for col in columns):
+                            fig = px.scatter(df, x=columns[0], y=columns[1], title=f"Graph between {columns[0]} and {columns[1]}")
                             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']}")
+                            st.success(f"Here is the graph between '{columns[0]}' and '{columns[1]}'.")
+                        else:
+                            st.warning("Columns not found in the dataset.")
+                    elif 'column' in question.lower():
+                        column = question.split('of')[-1].strip()
+                        if column in df.columns:
+                            fig = px.line(df, x=df.index, y=column, title=f"Graph of column '{column}'")
+                            st.plotly_chart(fig, use_container_width=True)
+                    
+                    st.stop()  # This halts further execution
 
             except Exception as e:
                 st.warning(f"Error processing question or generating answer: {str(e)}")
+                st.warning("Please retype your question and make sure to use the column name and cell value correctly.")