File size: 2,889 Bytes
a2134eb 09e2bc4 a2134eb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
#!/usr/bin/env python3
"""
Visualization script using PandasAI.
This script creates a sample dataframe and uses PandasAI to generate
and save visualizations based on user queries.
Usage:
python visualize.py "Create a bar chart of sales by region"
Requirements:
- pandas
- pandasai
- matplotlib
"""
import os
import sys
import pandas as pd
import matplotlib.pyplot as plt
import pandasai as pai
from dotenv import load_dotenv
def create_sample_dataframe():
"""Create a sample dataframe with sales data."""
# 'Region': ['North', 'South', 'East', 'West', 'North', 'South', 'East', 'West'],
# 'Product': ['Widget', 'Widget', 'Widget', 'Widget', 'Gadget', 'Gadget', 'Gadget', 'Gadget'],
# 'Sales': [150, 200, 120, 180, 90, 110, 95, 130],
# 'Quarter': ['Q1', 'Q1', 'Q1', 'Q1', 'Q2', 'Q2', 'Q2', 'Q2'],
data = {
'Year': [2023, 2023, 2023, 2023, 2023, 2023, 2023, 2023]
}
return pai.DataFrame(data)
def visualize_data(df, query):
"""
Generate visualization based on user query using PandasAI.
Args:
df: Pandas DataFrame containing the data
query: User query string describing the desired visualization
Returns:
Path to the saved visualization file
"""
# Initialize PandasAI with an LLM
# Note: In a real application, you would need to set up your OpenAI API key
# Either set OPENAI_API_KEY environment variable or pass it directly
try:
# llm = OpenAI(api_token=api_key)
# pandas_ai = PandasAI(llm)
load_dotenv()
pai.api_key.set(os.environ["PANDAS_KEY"])
df.chat(query)
# Generate the visualization
print(f"Generating visualization for query: '{query}'")
# Save the current figure
output_file = "visualization_output.png"
plt.savefig(output_file)
plt.close()
print(f"Visualization saved to {output_file}")
return output_file
except Exception as e:
print(f"Error generating visualization: {str(e)}")
return None
def main():
"""Main function to run the visualization script."""
# Get query from command line argument
# if len(sys.argv) < 2:
# print("Usage: python visualize.py \"Your visualization query here\"")
# print("Example: python visualize.py \"Create a bar chart of sales by region\"")
# return
# query = sys.argv[1]
query = "Plot a bar chart of sales by region"
# Create sample dataframe
df = create_sample_dataframe()
print("Sample DataFrame created:")
print(df.head())
# Generate and save visualization
output_file = visualize_data(df, query)
if output_file:
print(f"Visualization process completed. Output saved to: {output_file}")
else:
print("Visualization process failed.")
if __name__ == "__main__":
main() |