Update app.py

app.py

@@ -1,12 +1,12 @@
-
-import transformers
 import streamlit as st
 import pandas as pd
 import os
 import torch
+import transformers
 from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, AutoModelForCausalLM, pipeline
 from huggingface_hub import HfFolder
 from io import StringIO
+from tqdm import tqdm  # To display progress bar in Streamlit
 
 # Access the Hugging Face API token from environment variables
 hf_token = os.getenv('HF_API_TOKEN')
@@ -18,26 +18,25 @@ HfFolder.save_token(hf_token)
 # Set environment variable to avoid floating-point errors
 os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
 
-# Load the tokenizer and model
-tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+# Load the GPT-2 tokenizer and model
+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)
+text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokenizer_gpt2)
 
-# Load the Llama model and tokenizer at the beginning
-model_name = "meta-llama/Meta-Llama-3.1-8B"
-tokenizer_llama = AutoTokenizer.from_pretrained(model_name, token=hf_token)
+# 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(
-    model_name,
-    torch_dtype="auto",
-    device_map="auto",
+    "meta-llama/Meta-Llama-3.1-8B",
+    torch_dtype= 'auto',
+    device_map= 'auto',
     token=hf_token
 )
 
 
 
-
+# Define your 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.
@@ -100,19 +99,25 @@ def generate_synthetic_data(description, columns):
         # Return the generated synthetic data
         return generated_text
     except Exception as e:
-        print(f"Error in generate_synthetic_data: {e}")
+        st.error(f"Error in generate_synthetic_data: {e}")
         return f"Error: {e}"
 
 def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100):
     data_frames = []
     num_iterations = num_rows // rows_per_generation
 
-    for _ in range(num_iterations):
+    # Create a progress bar
+    progress_bar = st.progress(0)
+    
+    for i in tqdm(range(num_iterations)):
         generated_data = generate_synthetic_data(description, columns)
         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)