Update app.py

app.py

@@ -9,14 +9,9 @@ import spaces
 os.environ['CUDA_HOME'] = '/usr/local/cuda'
 os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
 
-def initialize_zerogpu():
-    # Example function to initialize ZeroGPU, replace with actual initialization if needed
-    pass
-
 # Function to perform JSON correction using a chosen model
 @spaces.GPU(duration=120)  # Use GPU for this function, if available
 def ai_fix_json(model_id, json_data):
-    # Initialize the InferenceClient with the chosen model
     client = InferenceClient(model=model_id)
 
     prompt = f"Fix the following JSON data and make it valid:\n\n{json_data}\n\nFixed JSON:"
@@ -31,11 +26,13 @@ def ai_fix_json(model_id, json_data):
             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.text_generation(prompt, max_new_tokens=1024)
-        fixed_json = response[0]['generated_text'].split("Fixed JSON:")[-1].strip()
-        return fixed_json, f"JSON fixed using AI on CPU with model: {model_id}"
+        try:
+            response = client.text_generation(prompt, max_new_tokens=1024)
+            fixed_json = response[0]['generated_text'].split("Fixed JSON:")[-1].strip()
+            return fixed_json, f"JSON fixed using AI on CPU with model: {model_id}"
+        except Exception as e:
+            return None, f"Failed to process with model {model_id}. Error: {str(e)}", None
 
 def process_file(model_id, uploaded_file):
     json_data = uploaded_file  # This is already the content of the file as a string
@@ -48,20 +45,29 @@ def process_file(model_id, uploaded_file):
     except json.JSONDecodeError as e:
         return None, f"Failed to fix JSON: {str(e)}", None
 
+# List of available models
+model_options = [
+    "EleutherAI/gpt-neo-2.7B",
+    "gpt2",
+    "facebook/opt-1.3b",
+    "EleutherAI/gpt-j-6B",
+    "google/flan-t5-base"
+]
+
 iface = gr.Interface(
     fn=process_file,
     inputs=[
-        gr.Textbox(label="Model ID (e.g., gpt-3.5-turbo)", value="gpt-3.5-turbo"),
+        gr.Dropdown(label="Select Model", choices=model_options, value="EleutherAI/gpt-neo-2.7B"),
         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 with Model Selection",
-    description="Upload a JSON file to automatically fix, remove duplicates, and download the cleaned version using AI with GPU/CPU fallback. Select any model from the Hugging Face Hub."
+    description="Upload a JSON file to automatically fix, remove duplicates, and download the cleaned version using AI with GPU/CPU fallback. Select any model from the dropdown list."
 )
 
 if __name__ == "__main__":
     try:
-        initialize_zerogpu()
+        # Initialize any necessary services, e.g., ZeroGPU
         print("ZeroGPU initialized.")
     except Exception as e:
         print(f"ZeroGPU initialization failed: {e}. Falling back to CPU.")