Update custom_llm.py

custom_llm.py

@@ -126,6 +126,94 @@ def custom_chain_with_history(llm, memory):
     # return {"chat_history":prompt_memory, "context":asyncio.run(create_vectorstore()).as_retriever(search_type="similarity", search_kwargs={"k": 12}) | format_docs, "question": RunnablePassthrough()} | prompt | llm
     return {"chat_history":lambda x:prompt_memory(x['memory']), "context":itemgetter("question") | asyncio.run(create_vectorstore()).as_retriever(search_type="similarity", search_kwargs={"k": 100000}) | format_docs, "question": lambda x:x['question']} | prompt | llm
 
+
+
+
+def format_df(df):
+    out = ""
+
+    for x in df.columns:
+        out+= x + "|"
+    out = out[:-1]
+
+    for _,row in df.iterrows():
+        for x in row.values:
+            out += str(x) + "|"
+    
+        out = out[:-1]
+    
+    return out
+
+
+def custom_dataframe_chain(llm, dataframe):
+    prompt = PromptTemplate.from_template("""<s><INST>You have access to a dataframe variable named df. Below are the examples of the dataframe:
+
+    {df_example}
+    
+    Given the following user input, create relevant python code to get the relevant information in the dataframe and store the response string result in a variable named "response". Do not explain, just create the python code:
+    
+    {question}
+    
+    Always change the corresponding columns into datetime format with parameter day_first=True, example:
+    df['column_name'] = pd.to_datetime(df['column_name'], day_first=True)
+    
+    
+    Always use idxmin or idxmax instead of array indicies whenever it is possible 
+    
+    
+    The output must follow the following example format:
+    ```python
+    # Generated Code
+    ```
+    
+    </INST></s>""")
+
+    def out_format(text:str):
+
+        prompt = PromptTemplate.from_template("""<s><INST>Fix the following code:
+        {code}
+    
+        Error Message : {err}
+    
+    
+        The output must follow the following example format:
+        ```python
+        # Generated Code
+        ```
+        
+        </INST></s>""")
+    
+        err_chain = prompt | llm
+    
+        e_ = None
+    
+        for _ in range(6):
+    
+            try :
+                text = text.split("```python")[-1].split("```")[0]
+                # print(text)
+                exec(text)
+                break
+                # return response
+            except Exception as e:
+                # print(e)
+                text = err_chain.invoke({"code":text, "err":str(e)})
+                e_ = e
+        exec(text)
+    
+        return text
+        
+        return "Bad Python Code, Error Message : " + str(e_)
+
+
+    return RunnablePassthrough.assign(df_example=format_df(dataframe.head(4))) | prompt | llm | out_format
+
+    
+
+
+
+
+
 class CustomLLM(LLM):
     repo_id : str
     api_token : str