Update app.py

app.py

@@ -1,10 +1,10 @@
 from fastapi import FastAPI
-from fastapi.responses import JSONResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 import pandas as pd
 import os
 import requests
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer
 from io import StringIO
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import HfFolder
@@ -41,25 +41,18 @@ def preprocess_user_prompt(user_prompt):
 # Define prompt template
 prompt_template = """\
 You are an expert in generating synthetic data for machine learning models.
-
 Your task is to generate a synthetic tabular dataset based on the description provided below.
-
 Description: {description}
-
 The dataset should include the following columns: {columns}
-
 Please provide the data in CSV format.
-
 Example Description:
 Generate a dataset for predicting house prices with columns: 'Size', 'Location', 'Number of Bedrooms', 'Price'
-
 Example Output:
 Size,Location,Number of Bedrooms,Price
 1200,Suburban,3,250000
 900,Urban,2,200000
 1500,Rural,4,300000
 ...
-
 Description:
 {description}
 Columns:
@@ -120,7 +113,7 @@ def process_generated_data(csv_data, expected_columns):
         return f"Failed to parse CSV data: {e}"
 
 def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100):
-    data_frames = []
+    csv_data_all = ""
     
     for _ in tqdm(range(num_rows // rows_per_generation), desc="Generating Data"):
         generated_data = generate_synthetic_data(description, columns)
@@ -129,12 +122,12 @@ def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_
 
         df_synthetic = process_generated_data(generated_data, columns)
         if isinstance(df_synthetic, pd.DataFrame) and not df_synthetic.empty:
-            data_frames.append(df_synthetic)
+            csv_data_all += df_synthetic.to_csv(index=False, header=False)
         else:
             print("Skipping invalid generation.")
 
-    if data_frames:
-        return pd.concat(data_frames, ignore_index=True)
+    if csv_data_all:
+        return csv_data_all
     else:
         return "No valid data frames to concatenate."
 
@@ -147,12 +140,13 @@ def generate_data(request: DataGenerationRequest):
     if isinstance(generated_data, str) and "Error" in generated_data:
         return JSONResponse(content={"error": generated_data}, status_code=500)
     
-    # Process the generated CSV data into a DataFrame
-    df_synthetic = process_generated_data(generated_data, columns)
-    if isinstance(df_synthetic, pd.DataFrame):
-        return JSONResponse(content={"data": df_synthetic.to_dict(orient="records")})
-    else:
-        return JSONResponse(content={"error": "Failed to generate valid synthetic data"}, status_code=500)
+    # Create a streaming response to return the CSV data
+    csv_buffer = StringIO(generated_data)
+    return StreamingResponse(
+        csv_buffer,
+        media_type="text/csv",
+        headers={"Content-Disposition": "attachment; filename=generated_data.csv"}
+    )
 
 @app.get("/")
 def greet_json():