Update app.py

app.py

@@ -87,63 +87,59 @@ generation_params = {
 def generate_synthetic_data(description, columns):
     formatted_prompt = format_prompt(description, columns)
     payload = {"inputs": formatted_prompt, "parameters": generation_params}
-    response = requests.post(API_URL, headers={"Authorization": f"Bearer {token}"}, json=payload)
-    response_data = response.json()
+    response = requests.post(API_URL, headers={"Authorization": f"Bearer {hf_token}"}, json=payload)
+    
+    try:
+        response_data = response.json()
+    except ValueError:
+        raise HTTPException(status_code=500, detail="Failed to parse response from the API.")
     
     if 'error' in response_data:
-        return f"Error: {response_data['error']}"
+        raise HTTPException(status_code=500, detail=f"API Error: {response_data['error']}")
+    
+    if 'generated_text' not in response_data[0]:
+        raise HTTPException(status_code=500, detail="Unexpected API response format.")
     
     return response_data[0]["generated_text"]
 
 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):
-            return f"Unexpected columns in the generated data: {df.columns}"
+            raise ValueError("Unexpected columns in the generated data.")
         
         return df
     except pd.errors.ParserError as e:
-        return f"Failed to parse CSV data: {e}"
+        raise HTTPException(status_code=500, detail=f"Failed to parse CSV data: {e}")
 
 def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100):
-    csv_data_all = ""
+    csv_data_all = StringIO()
     
     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 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:
-            print("Skipping invalid generation.")
-
-    if csv_data_all:
+            df_synthetic.to_csv(csv_data_all, index=False, header=False)
+    
+    if csv_data_all.tell() > 0:  # Check if there's any data in the buffer
+        csv_data_all.seek(0)  # Rewind the buffer to the beginning
         return csv_data_all
     else:
-        return "No valid data frames to concatenate."
+        raise HTTPException(status_code=500, detail="No valid data frames generated.")
 
 @app.post("/generate/")
 def generate_data(request: DataGenerationRequest):
     description = request.description.strip()
     columns = [col.strip() for col in request.columns]
-    generated_data = generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100)
-    
-    if isinstance(generated_data, str) and "Error" in generated_data:
-        return JSONResponse(content={"error": generated_data}, status_code=500)
+    csv_data = generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100)
     
-    # Create a streaming response to return the CSV data
-    csv_buffer = StringIO(generated_data)
+    # Return the CSV data as a downloadable file
     return StreamingResponse(
-        csv_buffer,
+        csv_data,
         media_type="text/csv",
         headers={"Content-Disposition": "attachment; filename=generated_data.csv"}
     )