Update app.py

app.py

@@ -24,103 +24,94 @@ 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.")
 
-tokenizer_gpt2 = GPT2Tokenizer.from_pretrained('gpt2')
+
+# Load GPT-2 model and tokenizer
+tokenizer_gpt2 = AutoTokenizer.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)
 
+# Define prompt template for generating the dataset
 prompt_template = """\
-You are an expert in generating synthetic data for machine learning models.
+You are an AI specialized in generating synthetic tabular data specifically for machine learning purposes.
 
-Your task is to generate a synthetic tabular dataset based on the description provided below.
+Task: Generate a synthetic dataset based on the provided description and column names.
 
 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'
+Columns: {columns}
 
-Example Output:
-Size,Location,Number of Bedrooms,Price
-1200,Suburban,3,250000
-900,Urban,2,200000
-1500,Rural,4,300000
-...
+Instructions:
 
-Description:
-{description}
-Columns:
-{columns}
-Output: """
+Output only the tabular data in valid CSV format.
+Include the header row followed by the data rows.
+Do not generate any additional text, explanations, comments, or code.
+Ensure that the values for each column are contextually appropriate.
 
-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))
-    return prompt
+Format Example (do not include this line or the following example in your output):
+Column1,Column2,Column3
+Value1,Value2,Value3
+Value4,Value5,Value6
+"""
 
-API_URL = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"
 
+# Define generation parameters
 generation_params = {
     "top_p": 0.90,
     "temperature": 0.8,
-    "max_new_tokens": 512,
+    "max_new_tokens": 1024,
     "return_full_text": False,
     "use_cache": False
 }
 
+def format_prompt(description, columns):
+    prompt = prompt_template.format(description=description, columns=",".join(columns))
+    return prompt
+
 def generate_synthetic_data(description, columns):
     formatted_prompt = format_prompt(description, columns)
     payload = {"inputs": formatted_prompt, "parameters": generation_params}
+
+    # Call Mixtral model to generate data
+    response = requests.post("https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1", 
+                             headers={"Authorization": f"Bearer {token}"}, json=payload)
+    
+    if response.status_code == 200:
+        return response.json()[0]["generated_text"]
+    else:
+        print(f"Error generating data: {response.status_code}, {response.text}")
+        return None
+
+def process_generated_data(csv_data):
     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:
-        # Replace inconsistent line endings
+        # Ensure the data is cleaned and correctly formatted
         cleaned_data = csv_data.replace('\r\n', '\n').replace('\r', '\n')
-        
-        # Check for common CSV formatting issues and apply corrections
-        cleaned_data = cleaned_data.strip().replace('|', ',').replace('  ', ' ').replace(' ,', ',')
-        
-        # Load the cleaned data into a DataFrame
         data = StringIO(cleaned_data)
-        df = pd.read_csv(data, delimiter=',')
-        
+
+        # Read the CSV data with specific parameters to handle irregularities
+        df = pd.read_csv(data)
+
         return df
-    
     except pd.errors.ParserError as e:
         print(f"Failed to parse CSV data: {e}")
         return None
 
 def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100):
     data_frames = []
-    
+
     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 generated_data:
+           
+            df_synthetic = process_generated_data(generated_data)
+
+            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)
     else: