Update app.py

app.py

@@ -13,13 +13,15 @@ os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
 @spaces.GPU(duration=120)  # Use GPU for this function, if available
 def ai_fix_json(model_id, json_data):
     client = InferenceClient(model=model_id)
-
     prompt = f"Fix the following JSON data and make it valid:\n\n{json_data}\n\nFixed JSON:"
     
     try:
         # Check if GPU is available
         if torch.cuda.is_available():
             response = client.text_generation(prompt, max_new_tokens=1024)
+            # Handle cases where the response might be malformed or too short
+            if not response or 'generated_text' not in response[0]:
+                return None, f"Failed to process JSON with model {model_id}. Response was invalid."
             fixed_json = response[0]['generated_text'].split("Fixed JSON:")[-1].strip()
             return fixed_json, f"JSON fixed using AI on GPU with model: {model_id}"
         else:
@@ -29,21 +31,27 @@ def ai_fix_json(model_id, json_data):
         print(f"Falling back to CPU due to: {gpu_error}")
         try:
             response = client.text_generation(prompt, max_new_tokens=1024)
+            # Handle cases where the response might be malformed or too short
+            if not response or 'generated_text' not in response[0]:
+                return None, f"Failed to process JSON with model {model_id}. Response was invalid."
             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
+            return None, f"Failed to process with model {model_id}. Error: {str(e)}"
 
 def process_file(model_id, uploaded_file):
-    json_data = uploaded_file  # This is already the content of the file as a string
+    json_data = uploaded_file.read().decode("utf-8")  # Ensure the file is read correctly as a string
     cleaned_json, message = ai_fix_json(model_id, json_data)
     
-    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
+    if cleaned_json:
+        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
+    else:
+        return None, message, None
 
 # List of available models
 model_options = [