Update app.py

app.py

@@ -1,27 +1,27 @@
 import gradio as gr
 import json
 from transformers import pipeline
+from huggingface_hub import InferenceApi
 
-# Load a text generation model from Hugging Face without an API key
-# You can use a smaller model if gpt-3.5-turbo isn't available
-generator = pipeline("text-generation", model="gpt-3.5-turbo")
+# Set up the Inference API to use a GPU in the Hugging Face Space
+inference = InferenceApi(repo_id="gpt-3.5-turbo", task="text-generation", use_gpu=True)
 
-# Function to attempt to fix JSON using a Hugging Face model
+# Function to attempt to fix JSON using a Hugging Face model with GPU
 def ai_fix_json(json_data):
     # Prepare a prompt for the AI model to fix the JSON
     prompt = f"Fix the following JSON data and make it valid:\n\n{json_data}\n\nFixed JSON:"
 
-    # Generate a response from the model
-    response = generator(prompt, max_length=1024, num_return_sequences=1)[0]['generated_text']
+    # Generate a response from the model using GPU
+    response = inference(inputs=prompt, max_length=1024, num_return_sequences=1)
 
     # Extract the fixed JSON from the response (you may need to fine-tune this depending on the model's output format)
-    fixed_json = response.split("Fixed JSON:")[-1].strip()
+    fixed_json = response[0]['generated_text'].split("Fixed JSON:")[-1].strip()
     
     # Try to load the fixed JSON to ensure it's valid
     try:
         parsed_data = json.loads(fixed_json)
         pretty_json = json.dumps(parsed_data, indent=4)
-        return pretty_json, "JSON fixed using AI."
+        return pretty_json, "JSON fixed using AI with GPU."
     
     except json.JSONDecodeError as e:
         return None, f"Failed to fix JSON: {str(e)}"
@@ -46,8 +46,8 @@ iface = gr.Interface(
     fn=process_file,
     inputs=gr.File(label="Upload your JSON file"),
     outputs=[gr.JSON(label="Fixed JSON"), "text", gr.File(label="Download cleaned JSON file")],
-    title="AI-Powered JSON Cleaner",
-    description="Upload a JSON file to automatically fix, remove duplicates, and download the cleaned version using AI."
+    title="AI-Powered JSON Cleaner with GPU",
+    description="Upload a JSON file to automatically fix, remove duplicates, and download the cleaned version using AI on a GPU."
 )
 
 iface.launch()