Update app.py

app.py

@@ -6,9 +6,7 @@ import os
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, AutoModelForCausalLM, pipeline
 from io import StringIO
-from tqdm import tqdm
-import accelerate
-from accelerate import init_empty_weights, disk_offload
+from accelerate import Accelerator
 from fastapi.middleware.cors import CORSMiddleware
 import re
 
@@ -17,7 +15,7 @@ app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # You can specify domains here
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -36,36 +34,22 @@ 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
+# Initialize accelerator with disk offload
+accelerator = Accelerator(cpu=False, disk_offload=True)
+
+# Load the Llama-3 model and tokenizer with disk offload
 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='auto',
+    offload_folder="offload",  # Folder to offload weights to disk
+    offload_state_dict=True,   # Offload state_dict to disk
     token=hf_token
-)
+).to(accelerator.device)
 
 # 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.
-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: """
+prompt_template = """..."""  # Your existing prompt template here
 
 class DataGenerationRequest(BaseModel):
     description: str
@@ -113,7 +97,6 @@ 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:
@@ -124,10 +107,8 @@ 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)
     
@@ -142,12 +123,9 @@ 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!"}