Update app.py

app.py

@@ -129,6 +129,81 @@ else:
                 # Get raw answer again from TAPAS
                 raw_answer = tqa(table=df, query=question, truncation=True)
 
+# Extract the raw answer from TAPAS and the aggregator it returned
+raw_answer = tqa(table=df, query=question, truncation=True)
+
+# Get the answer, coordinates, cells, and aggregator from the raw TAPAS output
+answer = raw_answer['answer']
+aggregator = raw_answer.get('aggregator', '')
+coordinates = raw_answer.get('coordinates', [])
+cells = raw_answer.get('cells', [])
+
+# Manually fix the aggregator if it returns an incorrect one
+if 'MEDIAN' in question.upper() and 'AVERAGE' in aggregator.upper():
+    aggregator = 'MEDIAN'
+elif 'MIN' in question.upper() and 'AVERAGE' in aggregator.upper():
+    aggregator = 'MIN'
+elif 'MAX' in question.upper() and 'AVERAGE' in aggregator.upper():
+    aggregator = 'MAX'
+elif 'TOTAL' in question.upper() and 'SUM' in aggregator.upper():
+    aggregator = 'SUM'
+
+# Use the corrected aggregator for further processing
+summary_type = aggregator.lower()
+
+# Now, calculate the correct value using pandas based on the corrected aggregator
+if summary_type == 'sum':
+    numeric_value = df_numeric[column_name].sum()
+elif summary_type == 'max':
+    numeric_value = df_numeric[column_name].max()
+elif summary_type == 'min':
+    numeric_value = df_numeric[column_name].min()
+elif summary_type == 'average':
+    numeric_value = df_numeric[column_name].mean()
+elif summary_type == 'count':
+    numeric_value = df_numeric[column_name].count()
+elif summary_type == 'median':
+    numeric_value = df_numeric[column_name].median()
+elif summary_type == 'std_dev':
+    numeric_value = df_numeric[column_name].std()
+else:
+    numeric_value = processed_answer  # Fallback if something went wrong
+
+# Construct a natural language response
+if summary_type == 'sum':
+    natural_language_answer = f"The total {column_name} is {numeric_value}."
+elif summary_type == 'maximum':
+    natural_language_answer = f"The highest {column_name} is {numeric_value}."
+elif summary_type == 'minimum':
+    natural_language_answer = f"The lowest {column_name} is {numeric_value}."
+elif summary_type == 'average':
+    natural_language_answer = f"The average {column_name} is {numeric_value}."
+elif summary_type == 'count':
+    natural_language_answer = f"The number of entries in {column_name} is {numeric_value}."
+elif summary_type == 'median':
+    natural_language_answer = f"The median {column_name} is {numeric_value}."
+elif summary_type == 'std_dev':
+    natural_language_answer = f"The standard deviation of {column_name} is {numeric_value}."
+else:
+    natural_language_answer = f"The value for {column_name} is {numeric_value}."
+
+# Display the result to the user
+st.markdown("<p style='font-family:sans-serif;font-size: 0.9rem;'> Analysis Results: </p>", unsafe_allow_html=True)
+st.success(f"""
+    • Answer: {natural_language_answer}
+
+    Data Location:
+    • Column: {column_name}
+
+    Additional Context:
+    • Query Asked: "{question}"
+""")
+
+
+
+
+                
+
                 # Display raw result for debugging purposes
                 st.markdown("<p style='font-family:sans-serif;font-size: 0.9rem;'> Raw Result: </p>", unsafe_allow_html=True)
                 st.success(raw_answer)