Update app.py

app.py

@@ -1,14 +1,15 @@
 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
-from tqdm import tqdm
+import logging
+import random
+import csv
 
 app = FastAPI()
 
@@ -88,47 +89,36 @@ def generate_synthetic_data(description, columns):
         else:
             raise ValueError("Invalid response format from Hugging Face API.")
     except (requests.RequestException, ValueError) as e:
-        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
+        logging.error(f"Error during API request or response processing: {e}")
+        return "name,age,course,grade\nSampleName,20,Course,0"
 
 def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100):
-    data_frames = []
+    csv_buffer = StringIO()
+    writer = csv.writer(csv_buffer)
+    
+    # Write header
+    writer.writerow(columns)
     
-    for _ in tqdm(range(num_rows // rows_per_generation), desc="Generating Data"):
+    rows_generated = 0
+    while rows_generated < num_rows:
         generated_data = generate_synthetic_data(description, columns)
-        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.")
+        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 data_frames:
-        return pd.concat(data_frames, ignore_index=True)
-    else:
-        print("No valid data frames to concatenate.")
-        return pd.DataFrame(columns=columns)
+    csv_buffer.seek(0)
+    return csv_buffer
 
 class DataGenerationRequest(BaseModel):
     description: str
@@ -138,16 +128,8 @@ class DataGenerationRequest(BaseModel):
 def generate_data(request: DataGenerationRequest):
     description = request.description.strip()
     columns = [col.strip() for col in request.columns]
-    csv_data = generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100)
+    csv_buffer = generate_large_synthetic_data(description, columns, num_rows=10000)
     
-    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,