Update app.py

app.py

@@ -95,11 +95,30 @@ def generate_synthetic_data(description, columns):
     
     return response_data[0]["generated_text"]
 
+def extract_valid_csv(csv_data, expected_columns):
+    # Attempt to extract valid CSV data from the response
+    lines = csv_data.split('\n')
+    valid_lines = []
+    
+    # Find the start of the valid CSV data by looking for the column headers
+    header_found = False
+    for line in lines:
+        if header_found:
+            if line.strip() == '':
+                continue
+            valid_lines.append(line)
+        elif set(line.split(',')) == set(expected_columns):
+            header_found = True
+            valid_lines.append(line)
+    
+    # Combine valid lines into a single string
+    valid_csv_data = '\n'.join(valid_lines)
+    return valid_csv_data
+
 def process_generated_data(csv_data, expected_columns):
     try:
-        # Ensure the data is cleaned and correctly formatted
-        cleaned_data = csv_data.replace('\r\n', '\n').replace('\r', '\n')
-        data = StringIO(cleaned_data)
+        valid_csv_data = extract_valid_csv(csv_data, expected_columns)
+        data = StringIO(valid_csv_data)
         
         # Read the CSV data
         df = pd.read_csv(data, delimiter=',')
@@ -148,6 +167,7 @@ def generate_data(request: DataGenerationRequest):
         headers={"Content-Disposition": "attachment; filename=generated_data.csv"}
     )
 
+
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}