Update app.py

app.py

@@ -30,12 +30,11 @@ text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokeniz
 # Load the Llama-3 model and tokenizer once during startup
 tokenizer_llama = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B", token=hf_token)
 model_llama = AutoModelForCausalLM.from_pretrained(
-        "meta-llama/Meta-Llama-3.1-8B",
-        torch_dtype='auto',
-        device_map='auto',
-        token=hf_token
-    )
-
+    "meta-llama/Meta-Llama-3.1-8B",
+    torch_dtype='auto',
+    device_map='auto',
+    token=hf_token
+)
 
 # Define your prompt template
 prompt_template = """\
@@ -80,11 +79,11 @@ def generate_synthetic_data(description, columns):
     try:
         # Prepare the input for the Llama model
         formatted_prompt = format_prompt(description, columns)
-        
+
         # Tokenize the prompt with truncation enabled
         inputs = tokenizer_llama(formatted_prompt, return_tensors="pt", truncation=True, max_length=512).to(model_llama.device)
         print(f"Input Tensor Size: {inputs['input_ids'].size()}")
-        
+
         # Generate synthetic data
         with torch.no_grad():
             outputs = model_llama.generate(
@@ -94,10 +93,10 @@ def generate_synthetic_data(description, columns):
                 temperature=generation_params["temperature"],
                 num_return_sequences=1,
             )
-        
+
         # Decode the generated output
         generated_text = tokenizer_llama.decode(outputs[0], skip_special_tokens=True)
-        
+
         # Return the generated synthetic data
         return generated_text
     except Exception as e:
@@ -110,20 +109,20 @@ def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_
 
     # Create a progress bar
     progress_bar = st.progress(0)
-    
+
     for i in tqdm(range(num_iterations)):
         generated_data = generate_synthetic_data(description, columns)
-
-print("Generated Data:\n", generated_data)  # Add this line to debug
+        print("Generated Data:\n", generated_data)  # Move the print statement here
 
         if "Error" in generated_data:
             return generated_data
+
         df_synthetic = process_generated_data(generated_data)
         data_frames.append(df_synthetic)
-        
+
         # Update the progress bar
         progress_bar.progress((i + 1) / num_iterations)
-    
+
     return pd.concat(data_frames, ignore_index=True)
 
 def process_generated_data(csv_data):
@@ -137,7 +136,6 @@ def process_generated_data(csv_data):
         print("DataFrame Shape:", df.shape)
         print("DataFrame Head:\n", df.head())
 
-
         # Check if the DataFrame is empty
         if df.empty:
             raise ValueError("Generated DataFrame is empty.")
@@ -157,7 +155,7 @@ if st.button("Generate"):
     description = description.strip()
     columns = [col.strip() for col in columns.split(',')]
     df_synthetic = generate_large_synthetic_data(description, columns)
-    
+
     if isinstance(df_synthetic, str) and "Error" in df_synthetic:
         st.error(df_synthetic)  # Display error message if any
     else: