Update app.py

app.py

@@ -46,7 +46,7 @@ model_llama = AutoModelForCausalLM.from_pretrained(
     token=hf_token
 )
 
-# Define your prompt template
+# Define your prompt template with a check for the purpose of dataset generation
 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.
@@ -54,6 +54,7 @@ Description: {description}
 The dataset should include the following columns: {columns}
 Please provide the data in CSV format with a minimum of 100 rows per generation.
 Ensure that the data is realistic, does not contain any duplicate rows, and follows any specific conditions mentioned.
+If the description is not related to generating synthetic datasets, please return the message: "This request is not for generating synthetic datasets."
 Example Description:
 Generate a dataset for predicting house prices with columns: 'Size', 'Location', 'Number of Bedrooms', 'Price'
 Example Output:
@@ -143,12 +144,14 @@ def generate_data(request: DataGenerationRequest):
     if "Error" in generated_data:
         return JSONResponse(content={"error": generated_data}, status_code=500)
     
+    # Check if the model generated a message about incorrect purpose
+    if "This request is not for generating synthetic datasets." in generated_data:
+        return JSONResponse(content={"message": "This request is not for generating synthetic datasets."}, status_code=400)
+    
     # Process the generated CSV data into a DataFrame
     df_synthetic = process_generated_data(generated_data)
     return JSONResponse(content={"data": df_synthetic.to_dict(orient="records")})
 
-
-
 @app.get("/")
 def greet_json():
-    return {"Hello": "World!"}
+    return {"Hello": "World!"}