Update app.py

app.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 import pandas as pd
 import os
@@ -87,7 +87,7 @@ 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 = requests.post(API_URL, headers={"Authorization": f"Bearer {hf_token}"}, json=payload)
     return response.json()[0]["generated_text"]
 
 def process_generated_data(csv_data, expected_columns):
@@ -127,26 +127,31 @@ def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_
         print("No valid data frames to concatenate.")
         return pd.DataFrame(columns=columns)
 
+class DataGenerationRequest(BaseModel):
+    description: str
+    columns: list[str]
 
-
-@app.route('/generate', methods=['POST'])
-def generate():
-    data = request.json
-    description = data.get('description')
-    columns = data.get('columns')
-    num_rows = data.get('num_rows', 1000)
-    
-    if not description or not columns:
-        return jsonify({"error": "Please provide 'description' and 'columns' in the request."}), 400
+@app.post("/generate/")
+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)
     
-    df_synthetic = generate_large_synthetic_data(description, columns, num_rows=num_rows)
+    if csv_data.empty:
+        return JSONResponse(content={"error": "No valid data generated"}, status_code=500)
     
-    if df_synthetic is not None and not df_synthetic.empty:
-        file_path = 'synthetic_data.csv'
-        df_synthetic.to_csv(file_path, index=False)
-        return send_file(file_path, as_attachment=True)
-    else:
-        return jsonify({"error": "Failed to generate a valid synthetic dataset."}), 500
-
-if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=8000)
+    # 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,
+        media_type="text/csv",
+        headers={"Content-Disposition": "attachment; filename=generated_data.csv"}
+    )
+
+@app.get("/")
+def greet_json():
+    return {"Hello": "World!"}