Update app.py

app.py

@@ -4,7 +4,7 @@ from pydantic import BaseModel
 import pandas as pd
 import os
 import requests
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, pipeline
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer
 from io import StringIO
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import HfFolder
@@ -95,51 +95,36 @@ def generate_synthetic_data(description, columns):
     
     return response_data[0]["generated_text"]
 
-def extract_valid_csv(csv_data, expected_columns):
-    lines = csv_data.split('\n')
-    valid_lines = []
-    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)
-    
-    valid_csv_data = '\n'.join(valid_lines)
-    return valid_csv_data
-
 def process_generated_data(csv_data, expected_columns):
     try:
-        valid_csv_data = extract_valid_csv(csv_data, expected_columns)
-        data = StringIO(valid_csv_data)
+        # Ensure the data is cleaned and correctly formatted
+        cleaned_data = csv_data.replace('\r\n', '\n').replace('\r', '\n')
+        data = StringIO(cleaned_data)
         
+        # Read the CSV data
         df = pd.read_csv(data, delimiter=',')
         
+        # Check if the DataFrame has the expected columns
         if set(df.columns) != set(expected_columns):
             return f"Unexpected columns in the generated data: {df.columns}"
         
         return df
     except pd.errors.ParserError as e:
-        logging.error(f"Failed to parse CSV data: {e}")
         return f"Failed to parse CSV data: {e}"
 
 def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100):
     csv_data_all = ""
     
-    for _ in range(num_rows // rows_per_generation):
+    for _ in tqdm(range(num_rows // rows_per_generation), desc="Generating Data"):
         generated_data = generate_synthetic_data(description, columns)
         if "Error" in generated_data:
-            return generated_data
+            return generated_data  # Return the error message
 
         df_synthetic = process_generated_data(generated_data, columns)
         if isinstance(df_synthetic, pd.DataFrame) and not df_synthetic.empty:
             csv_data_all += df_synthetic.to_csv(index=False, header=False)
         else:
-            logging.info("Skipping invalid generation.")
+            print("Skipping invalid generation.")
 
     if csv_data_all:
         return csv_data_all
@@ -155,6 +140,7 @@ def generate_data(request: DataGenerationRequest):
     if isinstance(generated_data, str) and "Error" in generated_data:
         return JSONResponse(content={"error": generated_data}, status_code=500)
     
+    # Create a streaming response to return the CSV data
     csv_buffer = StringIO(generated_data)
     return StreamingResponse(
         csv_buffer,
@@ -162,7 +148,6 @@ def generate_data(request: DataGenerationRequest):
         headers={"Content-Disposition": "attachment; filename=generated_data.csv"}
     )
 
-
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}