import streamlit as st
import pandas as pd
import torch
from transformers import pipeline
from pandasai import SmartDataframe

#llm =  pipeline("text-generation", model="microsoft/Phi-3-mini-128k-instruct", trust_remote_code=True)
#llm = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", torch_dtype=torch.bfloat16)
llm = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")

st.title("Data Analysis with Pandas AI")

uploader_file = st.file_uploader("Upload a excel file", type = ["xlsx"] )

if uploader_file is not None:
    data = pd.read_excel(uploader_file)
    st.write(data.head(3))
    df = SmartDataframe(data, config={"llm": llm})
    prompt = st.text_area("Enter your Prompt:")

    if st.button("Generate"):
        if prompt:
            with st.spinner("Generative response ...."):
                st.write(df.chat(prompt))
        else:
            st.warning("Please enter a Prompt")