|
|
|
import pandas as pd |
|
import pandasai as pdai |
|
from pandasai import PandasAI |
|
|
|
from pandasai.llm.starcoder import Starcoder |
|
|
|
|
|
models = {"Starcoder": Starcoder} |
|
|
|
|
|
model_to_run = 'Starcoder' |
|
|
|
|
|
|
|
API_KEY = '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' |
|
pd.DataFrame({ |
|
"country": ["United States", "United Kingdom", "France", "Germany", "Italy", "Spain", "Canada", "Australia", "Japan", "China"], |
|
"gdp": [21400000, 2940000, 2830000, 3870000, 2160000, 1350000, 1780000, 1320000, 516000, 14000000], |
|
"happiness_index": [7.3, 7.2, 6.5, 7.0, 6.0, 6.3, 7.3, 7.3, 5.9, 5.0] |
|
|
|
llm = models[model_to_run](api_token=API_KEY) |
|
pandas_ai = PandasAI(llm, conversational=False, verbose=True) |
|
|
|
prompt = 'Plot the histogram of countries showing for each the gdp, using different colors for each bar' |
|
response = pandas_ai.run(df, prompt=prompt, |
|
is_conversational_answer=False) |