Update app.py

app.py

@@ -3,10 +3,12 @@ from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 import pandas as pd
 import os
-import torch
+import requests
 from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, AutoModelForCausalLM, pipeline
 from io import StringIO
 from fastapi.middleware.cors import CORSMiddleware
+from huggingface_hub import HfFolder
+from tqdm import tqdm
 
 app = FastAPI()
 
@@ -24,39 +26,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 the GPT-2 tokenizer and model
+# Load GPT-2 model and tokenizer
 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)
 
-# Load the Llama-3 model and tokenizer once during startup
-device = "cuda" if torch.cuda.is_available() else "cpu"
-tokenizer_llama = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B", token=hf_token)
-model_llama = AutoModelForCausalLM.from_pretrained(
-    "meta-llama/Meta-Llama-3-8B",
-    torch_dtype='auto',
-    device_map='balanced',
-    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
 
-# Define your prompt template
+# Define 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 with a minimum of 100 rows per generation.
-Ensure that the data is realistic, does not contain any duplicate rows, and follows any specific conditions mentioned.
+
+Please provide the data in CSV format.
+
 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:
@@ -67,66 +70,89 @@ class DataGenerationRequest(BaseModel):
     description: str
     columns: list
 
-def preprocess_user_prompt(user_prompt):
-    generated_text = text_generator(user_prompt, max_length=60, num_return_sequences=1, truncation=True)[0]["generated_text"]
-    return generated_text
+# Set up the Mixtral model and tokenizer
+token = hf_token  # Use environment variable for the 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 = {
     "top_p": 0.90,
     "temperature": 0.8,
     "max_new_tokens": 512,
+    "return_full_text": False,
+    "use_cache": False
 }
 
 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 {token}"}, json=payload)
+    response_data = response.json()
+    
+    if 'error' in response_data:
+        return f"Error: {response_data['error']}"
+    
+    return response_data[0]["generated_text"]
+
+def process_generated_data(csv_data, expected_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)
-
-        # Generate synthetic data
-        with torch.no_grad():
-            outputs = model_llama.generate(
-                **inputs,
-                max_length=512,
-                top_p=generation_params["top_p"],
-                temperature=generation_params["temperature"],
-                num_return_sequences=1,
-            )
-
-        # Decode the generated output
-        generated_text = tokenizer_llama.decode(outputs[0], skip_special_tokens=True)
+        # 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=',')
         
-        # Return the generated synthetic data
-        return generated_text
-    except Exception as e:
-        return f"Error: {e}"
-
-def process_generated_data(csv_data):
-    # Convert to DataFrame
-    data = StringIO(csv_data)
-    df = pd.read_csv(data)
+        # Check if the DataFrame has the expected columns
+        if set(df.columns) != set(expected_columns):
+            return f"Unexpected columns in the generated data: {df.columns}"
+        
+        return df
+    except pd.errors.ParserError as e:
+        return f"Failed to parse CSV data: {e}"
+
+def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100):
+    data_frames = []
     
-    return df
+    for _ in tqdm(range(num_rows // rows_per_generation), desc="Generating Data"):
+        generated_data = generate_synthetic_data(description, columns)
+        if "Error" in generated_data:
+            return generated_data  # Return the error message
+
+        df_synthetic = process_generated_data(generated_data, columns)
+        if isinstance(df_synthetic, pd.DataFrame) 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:
+        return "No valid data frames to concatenate."
 
 @app.post("/generate/")
 def generate_data(request: DataGenerationRequest):
     description = request.description.strip()
     columns = [col.strip() for col in request.columns]
-    generated_data = generate_synthetic_data(description, columns)
+    generated_data = generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100)
     
-    if "Error" in generated_data:
+    if isinstance(generated_data, str) and "Error" in generated_data:
         return JSONResponse(content={"error": generated_data}, status_code=500)
     
     # Process the generated CSV data into a DataFrame
-    df_synthetic = process_generated_data(generated_data)
-    return JSONResponse(content={"data": df_synthetic.to_dict(orient="records")})
+    df_synthetic = process_generated_data(generated_data, columns)
+    if isinstance(df_synthetic, pd.DataFrame):
+        return JSONResponse(content={"data": df_synthetic.to_dict(orient="records")})
+    else:
+        return JSONResponse(content={"error": "Failed to generate valid synthetic data"}, status_code=500)
 
 @app.get("/")
 def greet_json():