Update app.py

app.py

@@ -24,25 +24,40 @@ 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')
 text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokenizer_gpt2)
 
-# 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.
-"""
 
-# Set up the Mixtral model and tokenizer
+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):
-    # 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
 
@@ -64,28 +79,32 @@ 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 {hf_token}"}, json=payload)
-    
-    if response.status_code != 200:
-        raise HTTPException(status_code=response.status_code, detail="Error from Hugging Face API")
-        
-    return response.json()[0]["generated_text"]
+    try:
+        response = requests.post(API_URL, headers={"Authorization": f"Bearer {hf_token}"}, json=payload)
+        response.raise_for_status()
+        data = response.json()
+        if 'generated_text' in data[0]:
+            return data[0]['generated_text']
+        else:
+            raise ValueError("Invalid response format from Hugging Face API.")
+    except (requests.RequestException, ValueError) as e:
+        print(f"Error during API request or response processing: {e}")
+        return ""
 
 def process_generated_data(csv_data, expected_columns):
     try:
-        # Ensure the data is cleaned and correctly formatted
+        # Replace inconsistent line endings
         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 for common CSV formatting issues and apply corrections
+        cleaned_data = cleaned_data.strip().replace('|', ',').replace('  ', ' ').replace(' ,', ',')
         
-        # 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
+        # Load the cleaned data into a DataFrame
+        data = StringIO(cleaned_data)
+        df = pd.read_csv(data, delimiter=',')
         
         return df
+    
     except pd.errors.ParserError as e:
         print(f"Failed to parse CSV data: {e}")
         return None
@@ -101,7 +120,7 @@ def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_
             data_frames.append(df_synthetic)
         else:
             print("Skipping invalid generation.")
-
+       
     if data_frames:
         return pd.concat(data_frames, ignore_index=True)
     else:
@@ -133,6 +152,4 @@ def generate_data(request: DataGenerationRequest):
 
 @app.get("/")
 def greet_json():
-    return {"Hello": "World!"}
-
-
+    return {"Hello": "World!"}