Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 import gradio as gr
 import numpy as np
+import json  # Import json for safer parsing
 from src.model import LSTM  # Adjust to your model path
 
 # Load the model
@@ -13,13 +14,21 @@ model.eval()
 # Define the prediction function
 def predict_water_usage(state_idx, target_year, structured_data):
     # Convert structured data JSON string to dictionary
-    structured_data = eval(structured_data) if isinstance(structured_data, str) else structured_data
-
+    try:
+        structured_data = json.loads(structured_data) if isinstance(structured_data, str) else structured_data
+    except json.JSONDecodeError:
+        return {"error": "Invalid JSON format for structured data."}
+    
     if state_idx not in structured_data or len(structured_data[state_idx]) < 5:
         return {"error": "Structured data must include 5 years of data for the specified state."}
 
-    # Convert structured data for model input
-    data_values = [values for year, values in structured_data[state_idx].items()]
+    # Convert structured data for model input (extract values for model)
+    data_values = [list(values) for year, values in structured_data[state_idx].items()]
+    
+    # Ensure the data has the right shape for the model
+    if len(data_values) != 5:  # Check if there are exactly 5 years of data
+        return {"error": "Structured data should have 5 years of data."}
+    
     tensor_data = torch.tensor(data_values, dtype=torch.float32).to(device)
     
     # Get model output
@@ -30,23 +39,24 @@ def predict_water_usage(state_idx, target_year, structured_data):
 
 # Configure Gradio interface
 inputs = [
-    gr.components.Number(label="State Index"),
-    gr.components.Number(label="Target Year"),
-    gr.components.Textbox(
+    gr.Number(label="State Index"),  # Numeric input for state index
+    gr.Number(label="Target Year"),  # Numeric input for target year
+    gr.Textbox(
         label="Structured Data (JSON format)",
         lines=10,
         placeholder="""{
-  "state_idx": {
-    "2020": [value1, value2, ..., value8],
-    "2021": [value1, value2, ..., value8],
-    "2022": [value1, value2, ..., value8],
-    "2023": [value1, value2, ..., value8],
-    "2024": [value1, value2, ..., value8]
-  }
-}"""
+            "state_idx": {
+                "2020": [value1, value2, ..., value8],
+                "2021": [value1, value2, ..., value8],
+                "2022": [value1, value2, ..., value8],
+                "2023": [value1, value2, ..., value8],
+                "2024": [value1, value2, ..., value8]
+            }
+        }"""
     )
 ]
-outputs = gr.components.JSON(label="Prediction")
+
+outputs = gr.JSON(label="Prediction")
 
 interface = gr.Interface(fn=predict_water_usage, inputs=inputs, outputs=outputs)