Update app.py

app.py

@@ -27,13 +27,17 @@ 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 = GPT2LMHeadModel.from_pretrained('gpt2')
+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.
@@ -62,12 +66,14 @@ Columns:
 {columns}
 Output: """
 
-# Set up the Mixtral model and tokenizer
-token = os.getenv("HF_TOKEN")
-HfFolder.save_token(token)
 
 tokenizer_mixtral = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1", token=token)
 
+def format_prompt(description, columns):
+    processed_description = preprocess_user_prompt(description)
+    prompt = prompt_template.format(description=processed_description, columns=",".join(columns))
+    return prompt
+
 API_URL = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"
 
 generation_params = {
@@ -78,15 +84,6 @@ generation_params = {
     "use_cache": False
 }
 
-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
-
 def generate_synthetic_data(description, columns):
     formatted_prompt = format_prompt(description, columns)
     payload = {"inputs": formatted_prompt, "parameters": generation_params}
@@ -95,12 +92,18 @@ 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
+        
         return df
     except pd.errors.ParserError as e:
         print(f"Failed to parse CSV data: {e}")
@@ -108,19 +111,24 @@ def process_generated_data(csv_data, expected_columns):
 
 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 data_frames:
         return pd.concat(data_frames, ignore_index=True)
     else:
         print("No valid data frames to concatenate.")
         return pd.DataFrame(columns=columns)
 
+
+
 @app.route('/generate', methods=['POST'])
 def generate():
     data = request.json