Update custom_llm.py

custom_llm.py

@@ -145,7 +145,9 @@ def format_df(df):
     return out
 
 
-out_prompt = PromptTemplate.from_template("""<s><INST>Fix the following code:
+def out_format(text, df):
+
+    prompt = PromptTemplate.from_template("""<s><INST>Fix the following code:
     {code}
     
     Error Message : {err}
@@ -158,63 +160,54 @@ out_prompt = PromptTemplate.from_template("""<s><INST>Fix the following code:
     
     </INST></s>""")
 
+    err_chain = prompt | llm
 
+    e_ = None
 
-df_prompt = PromptTemplate.from_template("""<s><INST>You have access to a pandas 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 
+    for _ in range(6):
 
-
-Do not import pandas and Do not create or re-assign "df" variable
-
-
-The output must follow the following example format:
-```python
-# Generated Code
-```
-
-</INST></s>""")
+        try :
+            # text = text.split("```python")[-1].split("```")[0]
+            text = text.split("# Generated Code")[-1].split("`")[0]
+            print(text)
+            exec(text)
+            
+            return text
+        except Exception as e:
+            print(e)
+            text = err_chain.invoke({"code":text, "err":str(e)})
+            e_ = e
+    # exec(text)
+    
+    return "Bad Python Code, Error Message : " + str(e_)
 
 
 
 def custom_dataframe_chain(llm, df):
-    prompt = df_prompt
-
-    def out_format(text, df):
+    prompt = PromptTemplate.from_template("""<s><INST>You have access to a pandas dataframe variable named "df". Below are the examples of the dataframe:
 
-        prompt = out_prompt
+    {df_example}
     
-        err_chain = prompt | llm
+    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:
     
-        e_ = None
+    {question}
     
-        for _ in range(6):
+    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)
     
-            try :
-                # text = text.split("```python")[-1].split("```")[0]
-                text = text.split("# Generated Code")[-1].split("`")[0]
-                print(text)
-                exec(text)
-                
-                return text
-            except Exception as e:
-                print(e)
-                text = err_chain.invoke({"code":text, "err":str(e)})
-                e_ = e
-        # exec(text)
-        
-        return "Bad Python Code, Error Message : " + str(e_)
-
+    
+    Always use idxmin or idxmax instead of array indicies whenever it is possible 
+    
+    
+    Do not import pandas and Do not create or re-assign "df" variable
+    
+    
+    The output must follow the following example format:
+    ```python
+    # Generated Code
+    ```
+    
+    </INST></s>""")
 
     return prompt | llm | RunnableLambda(lambda x:out_format(x, df))