Update app.py

app.py

@@ -425,10 +425,15 @@ default_attributes = [
 def dataframe_to_code(df):
     df = df.head(4)
     columns = df.columns.tolist()
-    dict_string = ", ".join([f"'{col}': {df[col].tolist()}" for col in columns])
-    code_string = f"df = pd.DataFrame(\n    {{{dict_string}}}\n)"
+    dict_entries = []
+    for col in columns:
+        values_list = repr(df[col].tolist())
+        dict_entries.append(f"        '{col}': {values_list}")
+    dict_string = ",\n".join(dict_entries)
+    code_string = f"df = pd.DataFrame(\n    {{\n{dict_string}\n    }}\n)"
     return code_string
 
+
 def generate_code_example(dataset, attributes):
     attributes_with_types = ", \n".join(
         [f"\t\tfantix.type.{attr.upper().replace(' ', '_')}" for attr in attributes]
@@ -439,6 +444,7 @@ def generate_code_example(dataset, attributes):
     code = f"""
 import fantix
 import pandas as pd
+
 client = fantix.Client(api_key="YOUR_API_KEY")
 
 {dataframe_code}
@@ -1106,27 +1112,27 @@ def predict(dataset, attributes, access_token):
     - attributes (list): The attributes to predict.
     - access_token (str): The access token for API authentication.
     Returns:
-    - tuple: A message about the prediction, prediction results as a DataFrame, 
+    - tuple: A message about the prediction, prediction results as a DataFrame,
                 and the number of predictions made in the given time frame.
     """
     api_url = "https://rb3mw988lz88cvpz.us-east-1.aws.endpoints.huggingface.cloud"
     headers = {
         "Accept": "application/json",
         "Authorization": f"Bearer {access_token}",
-        "Content-Type": "application/json" 
+        "Content-Type": "application/json",
     }
-    
+
     payload = {
         "inputs": [
             {
                 "input_data": dataset.to_dict(orient="records"),
                 "attributes_to_predict": [
-                   attribute.lower().replace(" ", "_") for attribute in attributes
+                    attribute.lower().replace(" ", "_") for attribute in attributes
                 ],
             }
         ],
     }
-    
+
     start_time = time.time()
     response = requests.post(api_url, headers=headers, json=payload)
     end_time = time.time()
@@ -1144,9 +1150,10 @@ def predict(dataset, attributes, access_token):
 
     return prediction_message, prediction_results
 
+
 def load_dataset_and_predict(dataset, attributes, access_token):
     loaded_data = load_dataset(dataset)
-    
+
     code_example = generate_code_example(dataset, attributes)
 
     if access_token: