Update app.py

app.py

@@ -1,39 +1,48 @@
 import gradio as gr
 import json
+import torch
 from huggingface_hub import InferenceClient
+from huggingface_hub import spaces
 
-# Initialize the InferenceClient
-client = InferenceClient(model="gpt-3.5-turbo")
+# Initialize the InferenceClient for GPU and CPU
+client_gpu = InferenceClient(model="gpt-3.5-turbo")  # Assuming this will use GPU by default if available
+client_cpu = InferenceClient(model="gpt-3.5-turbo")  # Duplicate client for CPU fallback
 
-# Function to attempt to fix JSON using a Hugging Face model
+# Function to attempt to fix JSON using the appropriate resource
 def ai_fix_json(json_data):
     prompt = f"Fix the following JSON data and make it valid:\n\n{json_data}\n\nFixed JSON:"
-    response = client.text_generation(prompt, max_new_tokens=1024)
-    
-    fixed_json = response[0]['generated_text'].split("Fixed JSON:")[-1].strip()
     
     try:
-        parsed_data = json.loads(fixed_json)
-        pretty_json = json.dumps(parsed_data, indent=4)
-        return pretty_json, "JSON fixed using AI."
-    except json.JSONDecodeError as e:
-        return None, f"Failed to fix JSON: {str(e)}"
+        # Check if the space is running on GPU
+        if torch.cuda.is_available():
+            response = client_gpu.text_generation(prompt, max_new_tokens=1024)
+            return response[0]['generated_text'].split("Fixed JSON:")[-1].strip(), "JSON fixed using AI on GPU."
+        else:
+            raise RuntimeError("GPU not available, falling back to CPU.")
+    
+    except Exception as gpu_error:
+        # Fall back to CPU if GPU processing fails
+        print(f"Falling back to CPU due to: {gpu_error}")
+        response = client_cpu.text_generation(prompt, max_new_tokens=1024)
+        return response[0]['generated_text'].split("Fixed JSON:")[-1].strip(), "JSON fixed using AI on CPU."
 
 def process_file(uploaded_file):
     json_data = uploaded_file  # This is already the content of the file as a string
     cleaned_json, message = ai_fix_json(json_data)
     
-    if cleaned_json:
-        return cleaned_json, message, cleaned_json
-    else:
-        return None, message, None
+    try:
+        parsed_data = json.loads(cleaned_json)
+        pretty_json = json.dumps(parsed_data, indent=4)
+        return pretty_json, message, pretty_json
+    except json.JSONDecodeError as e:
+        return None, f"Failed to fix JSON: {str(e)}", None
 
 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 Dynamic Resource Allocation",
+    description="Upload a JSON file to automatically fix, remove duplicates, and download the cleaned version using AI with GPU/CPU fallback."
 )
 
 iface.launch()