adding files

lstm_predictor.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+
+class LSTMPredictor(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, forecast_horizon=3, n_layers=2, dropout=0.2):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.n_layers = n_layers
+        self.forecast_horizon = forecast_horizon
+
+        self.lstm = nn.LSTM(input_dim, hidden_dim, n_layers, batch_first=True, dropout=dropout)
+        self.fc = nn.Linear(hidden_dim, output_dim * forecast_horizon)
+
+    def forward(self, x):
+        lstm_out, _ = self.lstm(x)
+        predictions = self.fc(lstm_out[:, -1, :])
+        return predictions.view(-1, self.forecast_horizon, predictions.shape[-1] // self.forecast_horizon)

predict.py

@@ -0,0 +1,77 @@
+import json
+import torch
+import numpy as np
+from torch.utils.data import DataLoader, Dataset
+from lstm_predictor import LSTMPredictor
+from huggingface_hub import hf_hub_download
+
+# Dataset Class
+class TimeSeriesDataset(Dataset):
+    def __init__(self, data, seq_length):
+        self.data = data
+        self.seq_length = seq_length
+
+    def __len__(self):
+        return len(self.data) - self.seq_length
+
+    def __getitem__(self, idx):
+        return torch.tensor(self.data[idx:idx + self.seq_length], dtype=torch.float32)
+
+# Load Config
+def load_config(config_path):
+    with open(config_path, 'r') as file:
+        config = json.load(file)
+    return config
+
+# Load Model from Hugging Face
+def load_model(config):
+    # Download model from Hugging Face
+    model_file = hf_hub_download(repo_id=config["repo_id"], filename=config["model_path"])
+    
+    # Load the model architecture
+    model = LSTMPredictor(
+        input_dim=config["input_dim"],
+        hidden_dim=config["hidden_dim"],
+        output_dim=config["output_dim"],
+        forecast_horizon=config["forecast_horizon"],
+        n_layers=config["n_layers"],
+        dropout=config["dropout"]
+    )
+    # Load weights
+    model.load_state_dict(torch.load(model_file, map_location=torch.device(config["device"])))
+    model.to(config["device"])
+    model.eval()
+    return model
+
+# Prediction Function
+def predict(model, dataloader, config):
+    predictions = []
+    with torch.no_grad():
+        for batch in dataloader:
+            batch = batch.to(config["device"])
+            output = model(batch)
+            predictions.append(output.cpu().numpy())
+    return np.vstack(predictions)
+
+# Main Function
+def main():
+    config_path = "config.json"  # Path to config file
+    config = load_config(config_path)
+
+    # Load test data
+    raw_data = np.load(config["data_path"])
+    dataset = TimeSeriesDataset(raw_data, seq_length=config["seq_length"])
+    dataloader = DataLoader(dataset, batch_size=config["batch_size"], shuffle=False)
+
+    # Load model
+    model = load_model(config)
+
+    # Predict
+    predictions = predict(model, dataloader, config)
+
+    # Save predictions
+    np.save(config["output_path"], predictions)
+    print(f"Predictions saved to {config['output_path']}")
+
+if __name__ == "__main__":
+    main()