Update app.py

app.py

@@ -6,7 +6,9 @@ import os
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, AutoModelForCausalLM, pipeline
 from io import StringIO
-from accelerate import Accelerator
+from tqdm import tqdm
+import accelerate
+from accelerate import init_empty_weights, disk_offload
 from fastapi.middleware.cors import CORSMiddleware
 import re
 
@@ -15,7 +17,7 @@ app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # You can specify domains here
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -39,13 +41,32 @@ 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='float16',
-    device_map='auto',
+    torch_dtype='auto',
+    device_map='balanced',
     token=hf_token
-).to(device) 
+).to(device)
 
 # Define your prompt template
-prompt_template = """..."""  # Your existing prompt template here
+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.
+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:
+{columns}
+Output: """
 
 class DataGenerationRequest(BaseModel):
     description: str
@@ -93,6 +114,7 @@ def generate_synthetic_data(description, columns):
         return f"Error: {e}"
         
 def clean_generated_text(generated_text):
+    # Extract CSV part using a regular expression
     csv_match = re.search(r'(\n?([A-Za-z0-9_]+,)*[A-Za-z0-9_]+\n([^\n,]*,)*[^\n,]*\n*)+', generated_text)
     
     if csv_match:
@@ -103,8 +125,10 @@ def clean_generated_text(generated_text):
     return csv_text
 
 def process_generated_data(csv_data):
+    # Clean the generated data
     cleaned_data = clean_generated_text(csv_data)
     
+    # Convert to DataFrame
     data = StringIO(cleaned_data)
     df = pd.read_csv(data)
     
@@ -119,9 +143,12 @@ def generate_data(request: DataGenerationRequest):
     if "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")})
 
+
+
 @app.get("/")
 def greet_json():
-    return {"Hello": "World!"}
+    return {"Hello": "World!"}