Model pipeline

app.py

@@ -1,65 +1,56 @@
 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
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model_path = "./water_forecast_2.pth"
+model_path = "./water_forecast_2.pth"  # Path to the model file
 model = LSTM(input_size=8, lstm_layer_sizes=[128, 128, 128], output_size=3).to(device)
 model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
 model.eval()
 
 # Define the prediction function
 def predict_water_usage(state_idx, target_year, structured_data):
-    # Convert structured data JSON string to dictionary
-    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:
+    # structured_data is now a dictionary directly, no need to parse it
+    if len(structured_data) < 5:
         return {"error": "Structured data must include 5 years of data for the specified state."}
 
     # Convert structured data for model input (extract values for model)
-    data_values = [list(values) for year, values in structured_data[state_idx].items()]
+    data_values = [list(values) for values in structured_data.values()]
+    print("Structured Data:", structured_data)
+    print("Data Values:", data_values)
     
     # 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."}
-    
+
+    # Check if data_values contains only numeric data
+    for year_data in data_values:
+        if not all(isinstance(val, (int, float)) for val in year_data):
+            return {"error": "All values in structured data should be numeric."}
+
+    # Convert data_values to tensor
     tensor_data = torch.tensor(data_values, dtype=torch.float32).to(device)
-    
+
     # Get model output
     with torch.no_grad():
         output = model(tensor_data)
-    
+
     return {"prediction": output.tolist()}
 
 # Configure Gradio interface
 inputs = [
     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]
-            }
-        }"""
-    )
+    gr.JSON(label="Structured Data")  # JSON input for structured data
 ]
 
 outputs = gr.JSON(label="Prediction")
 
+# Set up the Gradio Interface
 interface = gr.Interface(fn=predict_water_usage, inputs=inputs, outputs=outputs)
 
-# Launch the Gradio app
+# Launch Gradio
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch(show_error=True)