Update app.py

app.py

@@ -12,16 +12,38 @@ model.eval()
 
 # Define the prediction function
 def predict_water_usage(state_idx, target_year, structured_data):
-    tensor_data = torch.tensor(np.array(list(structured_data.values())), dtype=torch.float32).to(device)
+    # Ensure structured_data is correctly formatted
+    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()]
+    tensor_data = torch.tensor(data_values, dtype=torch.float32).to(device)
+    
+    # Get model output
     with torch.no_grad():
         output = model(tensor_data)
-    return output.tolist()
+    
+    return {"prediction": output.tolist()}
 
-# Update Gradio interface using gr.components
+# Configure Gradio interface with structured JSON input example
 inputs = [
     gr.components.Number(label="State Index"),
     gr.components.Number(label="Target Year"),
-    gr.components.JSON(label="Structured Data")  # Expects JSON input
+    gr.components.JSON(
+        label="Structured Data", 
+        placeholder="""
+{
+  "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")
 
@@ -29,4 +51,4 @@ interface = gr.Interface(fn=predict_water_usage, inputs=inputs, outputs=outputs)
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    interface.launch(share=True)
+    interface.launch()