Update app.py

app.py

@@ -1,10 +1,11 @@
 import gradio as gr
 import pandas as pd
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoModelForCausalLM
+import requests
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, pipeline, AutoModelForCausalLM
+from huggingface_hub import HfFolder
 from io import StringIO
 import os
 import torch
-from huggingface_hub import HfFolder
 
 # Access the Hugging Face API token from environment variables
 hf_token = os.getenv('HF_API_TOKEN')
@@ -12,17 +13,31 @@ 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.")
 HfFolder.save_token(hf_token)
-# Load the GPT-2 tokenizer and model
-tokenizer_gpt2 = GPT2Tokenizer.from_pretrained('gpt2')
+
+# 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')
 model_gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')
 
-# Load the Llama3 model in sharded mode
-model_name = "meta-llama/Meta-Llama-3.1-8B"
-try:
-    model_llama = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16,
-        load_in_8bit=True, token = hf_token)  # use device_map for automatic sharding
-except OSError as e:
-    print(f"Error loading model: {e}")
+# Create a pipeline for text generation using GPT-2
+text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokenizer)
+
+# Lazy loading function for Llama-3 model
+model_llama = None
+tokenizer_llama = None
+
+def load_llama_model():
+    global model_llama, tokenizer_llama
+    if model_llama is None:
+        model_name = "meta-llama/Meta-Llama-3.1-8B"
+        model_llama = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,  # Use FP16 for reduced memory
+            use_auth_token=hf_token
+        )
+        tokenizer_llama = AutoTokenizer.from_pretrained(model_name, token=hf_token)
 
 # Define your prompt template
 prompt_template = """\
@@ -47,21 +62,42 @@ Columns:
 Output: """
 
 def preprocess_user_prompt(user_prompt):
-    generated_text = model_gpt2.generate(tokenizer_gpt2.encode(user_prompt, return_tensors='pt'), max_length=60)[0]
-    return tokenizer_gpt2.decode(generated_text, skip_special_tokens=True)
+    generated_text = text_generator(user_prompt, max_length=60, num_return_sequences=1)[0]["generated_text"]
+    return generated_text
 
 def format_prompt(description, columns):
     processed_description = preprocess_user_prompt(description)
     prompt = prompt_template.format(description=processed_description, columns=",".join(columns))
     return prompt
 
+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):
     try:
+        # Load the Llama model only when generating data
+        load_llama_model()
+        
         formatted_prompt = format_prompt(description, columns)
-        inputs = tokenizer_llama(formatted_prompt, return_tensors="pt")
-        generated_output = model_llama.generate(**inputs, max_new_tokens=512)
-        generated_text = tokenizer_llama.decode(generated_output[0], skip_special_tokens=True)
-        return generated_text
+        payload = {"inputs": formatted_prompt, "parameters": generation_params}
+        headers = {"Authorization": f"Bearer {hf_token}"}
+        
+        response = requests.post(API_URL, headers=headers, json=payload)
+        
+        if response.status_code == 200:
+            response_json = response.json()
+            if isinstance(response_json, list) and len(response_json) > 0 and "generated_text" in response_json[0]:
+                return response_json[0]["generated_text"]
+            else:
+                raise ValueError("Unexpected response format or missing 'generated_text' key")
+        else:
+            print(f"Error details: {response.text}") 
+            raise ValueError(f"API request failed with status code {response.status_code}: {response.text}")
     except Exception as e:
         print(f"Error in generate_synthetic_data: {e}")
         return f"Error: {e}"