Update app.py

app.py

@@ -1,15 +1,14 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
+import pandas as pd
 import os
 import requests
 from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, pipeline
 from io import StringIO
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import HfFolder
-import logging
-import random
-import csv
+from tqdm import tqdm
 
 app = FastAPI()
 
@@ -89,36 +88,47 @@ def generate_synthetic_data(description, columns):
         else:
             raise ValueError("Invalid response format from Hugging Face API.")
     except (requests.RequestException, ValueError) as e:
-        logging.error(f"Error during API request or response processing: {e}")
-        return "name,age,course,grade\nSampleName,20,Course,0"
+        print(f"Error during API request or response processing: {e}")
+        return ""
+
+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)
+        
+        # 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):
+            print(f"Unexpected columns in the generated data: {df.columns}")
+            return None
+        
+        return df
+    except pd.errors.ParserError as e:
+        print(f"Failed to parse CSV data: {e}")
+        return None
 
 def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100):
-    csv_buffer = StringIO()
-    writer = csv.writer(csv_buffer)
-    
-    # Write header
-    writer.writerow(columns)
+    data_frames = []
     
-    rows_generated = 0
-    while rows_generated < num_rows:
+    for _ in tqdm(range(num_rows // rows_per_generation), desc="Generating Data"):
         generated_data = generate_synthetic_data(description, columns)
-        cleaned_data = generated_data.replace('\r\n', '\n').replace('\r', '\n')
-        data = StringIO(cleaned_data)
-        
-        # Append rows to CSV buffer
-        reader = csv.reader(data)
-        header_written = False
-        for row in reader:
-            if not header_written:
-                header_written = True
-                continue  # Skip the header of the generated data
-            writer.writerow(row)
-            rows_generated += 1
-            if rows_generated >= num_rows:
-                break
+        if generated_data:
+            df_synthetic = process_generated_data(generated_data, columns)
+            if df_synthetic is not None and not df_synthetic.empty:
+                data_frames.append(df_synthetic)
+            else:
+                print("Skipping invalid generation.")
+        else:
+            print("Skipping empty or invalid generation.")
 
-    csv_buffer.seek(0)
-    return csv_buffer
+    if data_frames:
+        return pd.concat(data_frames, ignore_index=True)
+    else:
+        print("No valid data frames to concatenate.")
+        return pd.DataFrame(columns=columns)
 
 class DataGenerationRequest(BaseModel):
     description: str
@@ -128,8 +138,16 @@ class DataGenerationRequest(BaseModel):
 def generate_data(request: DataGenerationRequest):
     description = request.description.strip()
     columns = [col.strip() for col in request.columns]
-    csv_buffer = generate_large_synthetic_data(description, columns, num_rows=10000)
+    csv_data = generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100)
     
+    if csv_data.empty:
+        return JSONResponse(content={"error": "No valid data generated"}, status_code=500)
+    
+    # Convert the DataFrame to CSV format
+    csv_buffer = StringIO()
+    csv_data.to_csv(csv_buffer, index=False)
+    csv_buffer.seek(0)
+
     # Return the CSV data as a downloadable file
     return StreamingResponse(
         csv_buffer,