Update app.py

app.py

@@ -14,53 +14,34 @@ app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # You can specify domains here
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Access the Hugging Face API token from environment variables
 hf_token = os.getenv('HF_API_TOKEN')
-
 if not hf_token:
     raise ValueError("Hugging Face API token is not set. Please set the HF_API_TOKEN environment variable.")
 
-# Load GPT-2 model and tokenizer
 tokenizer_gpt2 = GPT2Tokenizer.from_pretrained('gpt2')
 model_gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')
-
-# Create a pipeline for text generation using GPT-2
 text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokenizer_gpt2)
 
-def preprocess_user_prompt(user_prompt):
-    # Generate a structured prompt based on the user input
-    generated_text = text_generator(user_prompt, max_length=50, num_return_sequences=1)[0]["generated_text"]
-    return generated_text
-
-# 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:
-{columns}
-Output: """
+"""
 
 tokenizer_mixtral = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1", token=hf_token)
 
+def preprocess_user_prompt(user_prompt):
+    generated_text = text_generator(user_prompt, max_length=50, num_return_sequences=1)[0]["generated_text"]
+    return generated_text
+
 def format_prompt(description, columns):
     processed_description = preprocess_user_prompt(description)
     prompt = prompt_template.format(description=processed_description, columns=",".join(columns))
@@ -93,14 +74,10 @@ def generate_synthetic_data(description, columns):
 
 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
@@ -115,11 +92,12 @@ def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_
     
     for _ in tqdm(range(num_rows // rows_per_generation), desc="Generating Data"):
         generated_data = generate_synthetic_data(description, columns)
+        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.")
+        if df_synthetic is not None and not df_synthetic.empty:
+            data_frames.append(df_synthetic)
+        else:
+            print("Skipping invalid generation.")
        
     if data_frames:
         return pd.concat(data_frames, ignore_index=True)
@@ -140,12 +118,10 @@ def generate_data(request: DataGenerationRequest):
     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,
         media_type="text/csv",
@@ -155,3 +131,4 @@ def generate_data(request: DataGenerationRequest):
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
+