Update app.py

app.py

@@ -1,43 +1,17 @@
 import gradio as gr
 import pandas as pd
 import requests
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, pipeline, AutoModelForCausalLM
-from huggingface_hub import HfFolder
 from io import StringIO
 import os
-import torch
 
 # Access the Hugging Face API token from environment variables
 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)
 
-# 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')
-
-# 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)
+# Set the inference endpoint URL
+inference_endpoint = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3.1-8B"
 
 # Define your prompt template
 prompt_template = """\
@@ -62,53 +36,45 @@ Columns:
 Output: """
 
 def preprocess_user_prompt(user_prompt):
-    generated_text = text_generator(user_prompt, max_length=60, num_return_sequences=1)[0]["generated_text"]
-    return generated_text
+    return user_prompt
 
 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()
-        
-        # Prepare the input for the Llama model
+        # Format the prompt
         formatted_prompt = format_prompt(description, columns)
         
-        # Tokenize the prompt
-        inputs = tokenizer_llama(formatted_prompt, return_tensors="pt").to(model_llama.device)
+        # Send a POST request to the Hugging Face Inference API
+        headers = {
+            "Authorization": f"Bearer {hf_token}",
+            "Content-Type": "application/json"
+        }
+        data = {
+            "inputs": formatted_prompt,
+            "parameters": {
+                "max_new_tokens": 512,
+                "top_p": 0.90,
+                "temperature": 0.8
+            }
+        }
         
-        # 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)
+        response = requests.post(inference_endpoint, json=data, headers=headers)
+
+        if response.status_code != 200:
+            return f"Error: {response.status_code}, {response.text}"
         
-        # Return the generated synthetic data
+        # Extract the generated text from the response
+        generated_text = response.json()[0]['generated_text']
         return generated_text
+
     except Exception as e:
         print(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
@@ -147,7 +113,6 @@ iface = gr.Interface(
     api_name="generate"  # Set the API name directly here
 )
 
-
 iface.api_name = "generate"
 
 # Run the Gradio app