import gradio as gr
from pandas import DataFrame
from pandasai import PandasAI
from pandasai.llm.starcoder import Starcoder
from datasets import load_dataset

def send_prompt(token: str, df: DataFrame, prompt: str) -> str:
    llm = Starcoder(api_token=token)
    pandas_ai = PandasAI(llm, conversational=True)
    return pandas_ai(df, prompt=prompt) 

def get_result(token, dataset, config, split, prompt) -> str:
    try:
        dataset = load_dataset(dataset, config, split=split)
        df = dataset.to_pandas()
        return send_prompt(token, df, prompt=prompt)
    except Exception as e:
        return str(e)

with gr.Blocks() as demo:
    gr.Markdown(" ## PandasAI demo using datasets library")
    gr.Markdown(" pandasai library https://github.com/gventuri/pandas-ai")
    gr.Markdown(" datasets library https://huggingface.co/docs/datasets")

    hf_token = gr.Textbox(label="hugging face token", placeholder="hf_xxx", type="password")
    dataset = gr.Textbox(label="dataset", placeholder="mstz/iris", value="mstz/iris")
    config = gr.Textbox(label="config", placeholder="iris", value="iris")
    split = gr.Textbox(label="split", placeholder="train", value="train")
    prompt = gr.Textbox(label="prompt (str)", placeholder="How many records do I have?. Give me the median of sepal_width. Show me the first 5 rows.")
    output = gr.Textbox(label="Output Box")
    get_result_button = gr.Button("Submit")
    get_result_button.click(get_result, inputs=[hf_token, dataset, config, split, prompt], outputs=output)

if __name__ == "__main__":
    demo.launch()