File size: 4,293 Bytes
01289c8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import os
import pandas as pd
from pandasai import Agent, SmartDataframe
from typing import Tuple
from PIL import Image
from pandasai.llm import HuggingFaceTextGen
from dotenv import load_dotenv
from langchain_groq.chat_models import ChatGroq

load_dotenv("Groq.txt")
Groq_Token = os.environ["GROQ_API_KEY"]
models = {"mixtral": "mixtral-8x7b-32768", "llama": "llama2-70b-4096", "gemma": "gemma-7b-it"}

hf_token = os.getenv("HF_READ")

def preprocess_and_load_df(path: str) -> pd.DataFrame:
    df = pd.read_csv(path)
    df["Timestamp"] = pd.to_datetime(df["Timestamp"])
    return df

def load_agent(df: pd.DataFrame, context: str, inference_server: str, name="mixtral") -> Agent:
    # llm = HuggingFaceTextGen(
    # inference_server_url=inference_server,
    # max_new_tokens=250,
    # temperature=0.1,
    # repetition_penalty=1.2,
    # top_k=5,
    # )
    # llm.client.headers = {"Authorization": f"Bearer {hf_token}"}
    llm = ChatGroq(model=models[name], api_key=os.getenv("GROQ_API"), temperature=0.1)
    
    agent = Agent(df, config={"llm": llm, "enable_cache": False, "options": {"wait_for_model": True}})
    agent.add_message(context)
    return agent

def load_smart_df(df: pd.DataFrame, inference_server: str, name="mixtral") -> SmartDataframe:
    # llm = HuggingFaceTextGen(
    # inference_server_url=inference_server,
    # )
    # llm.client.headers = {"Authorization": f"Bearer {hf_token}"}
    llm = ChatGroq(model=models[name], api_key=os.getenv("GROQ_API"), temperature=0.1)
    df = SmartDataframe(df, config={"llm": llm, "max_retries": 5, "enable_cache": False})
    return df

def get_from_user(prompt):
    return {"role": "user", "content": prompt}

def ask_agent(agent: Agent, prompt: str) -> Tuple[str, str, str]:
    response = agent.chat(prompt)
    gen_code = agent.last_code_generated
    ex_code = agent.last_code_executed
    last_prompt = agent.last_prompt
    return {"role": "assistant", "content": response, "gen_code": gen_code, "ex_code": ex_code, "last_prompt": last_prompt}

def decorate_with_code(response: dict) -> str:
    return f"""<details>
<summary>Generated Code</summary>
    
```python
{response["gen_code"]}
```
</details>

<details>
<summary>Prompt</summary>

{response["last_prompt"]}
"""

def show_response(st, response):
    with st.chat_message(response["role"]):
        try:
            image = Image.open(response["content"])
            if "gen_code" in response:
                st.markdown(decorate_with_code(response), unsafe_allow_html=True)
            st.image(image)
        except Exception as e:
            if "gen_code" in response:
                display_content = decorate_with_code(response) + f"""</details>

{response["content"]}"""
            else:
                display_content = response["content"]
            st.markdown(display_content, unsafe_allow_html=True)

def ask_question(model_name, question):
    llm = ChatGroq(model=models[model_name], api_key=os.getenv("GROQ_API"), temperature=0.1)

    df_check = pd.read_csv("Data.csv")
    df_check["Timestamp"] = pd.to_datetime(df_check["Timestamp"])
    df_check = df_check.head(5)

    new_line = "\n"

    template = f"""```python
import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv("Data.csv")
df["Timestamp"] = pd.to_datetime(df["Timestamp"])

    # df.dtypes
    {new_line.join(map(lambda x: '# '+x, str(df_check.dtypes).split(new_line)))}

    # {question.strip()}
    # <your code here>
    ```
    """

    query = f"""I have a pandas dataframe data of PM2.5 and PM10.
    * Frequency of data is daily. 
    * `pollution` generally means `PM2.5`.
    * Save result in a variable `answer` and make it global.
    * If result is a plot, save it and save path in `answer`. Example: `answer='plot.png'`
    * If result is not a plot, save it as a string in `answer`. Example: `answer='The city is Mumbai'`

    Complete the following code.

    {template}

    """

    answer = llm.invoke(query)
    code = f"""
    {template.split("```python")[1].split("```")[0]}
    {answer.content.split("```python")[1].split("```")[0]}
    """
    # update variable `answer` when code is executed
    exec(code)

    return {"role": "assistant", "content": answer.content, "gen_code": code, "ex_code": code, "last_prompt": question}