Spaces:
Sleeping
Sleeping
| # inference.py (final robust) | |
| import pandas as pd | |
| import numpy as np | |
| from statsmodels.tsa.arima.model import ARIMA | |
| from sklearn.preprocessing import MinMaxScaler | |
| from tensorflow.keras.models import Sequential | |
| from tensorflow.keras.layers import LSTM, Dense | |
| def arima_forecast(ts_data, order=(5,1,0), steps=5): | |
| ts_series = pd.Series(ts_data) | |
| model = ARIMA(ts_series, order=order) | |
| model_fit = model.fit() | |
| forecast = model_fit.forecast(steps=steps) | |
| return forecast.tolist() | |
| def lstm_forecast(ts_data, look_back=60, steps=5, epochs=20): | |
| # Adjust look_back if input is too short | |
| if len(ts_data) < look_back + 1: | |
| look_back = max(1, len(ts_data) - 1) | |
| # Normalize | |
| scaler = MinMaxScaler(feature_range=(0, 1)) | |
| scaled_data = scaler.fit_transform(np.array(ts_data).reshape(-1,1)) | |
| # Create sequences | |
| def create_sequences(dataset, look_back): | |
| X, Y = [], [] | |
| for i in range(len(dataset) - look_back): | |
| X.append(dataset[i:(i+look_back), 0]) | |
| Y.append(dataset[i + look_back, 0]) | |
| return np.array(X), np.array(Y) | |
| X, y = create_sequences(scaled_data, look_back) | |
| if X.size == 0: | |
| # fallback: train on the whole series as single sample (not ideal but avoids crash) | |
| X = scaled_data[:-1].reshape(1, -1, 1) | |
| y = np.array([scaled_data[-1,0]]) | |
| else: | |
| X = X.reshape((X.shape[0], X.shape[1], 1)) | |
| # Build LSTM | |
| model = Sequential() | |
| # prefer Input layer to avoid warnings | |
| model.add(LSTM(50, return_sequences=True, input_shape=(X.shape[1], 1))) | |
| model.add(LSTM(50)) | |
| model.add(Dense(1)) | |
| model.compile(optimizer='adam', loss='mean_squared_error') | |
| # Train | |
| model.fit(X, y, epochs=epochs, batch_size=32, verbose=0) | |
| # Forecast future steps | |
| last_seq = scaled_data[-look_back:].reshape(1, look_back, 1) | |
| predictions = [] | |
| for _ in range(steps): | |
| pred = model.predict(last_seq, verbose=0) # pred shape can vary by TF version | |
| val = float(np.asarray(pred).reshape(-1)[0]) # extract scalar safely | |
| predictions.append(val) | |
| # make pred into shape (1,1,1) | |
| pred_reshaped = np.asarray(pred).reshape(1,1,1) | |
| last_seq = np.concatenate([last_seq[:,1:,:], pred_reshaped], axis=1) | |
| predictions = scaler.inverse_transform(np.array(predictions).reshape(-1,1)) | |
| return predictions.flatten().tolist() | |
| def infer(model_type: str, input_data: list, steps: int = 5, epochs: int = 20): | |
| """ | |
| model_type: 'arima' or 'lstm' | |
| input_data: list of recent stock prices | |
| steps: number of future days to forecast | |
| epochs: LSTM training epochs (only used for LSTM) | |
| """ | |
| if model_type.lower() == 'arima': | |
| return arima_forecast(input_data, steps=steps) | |
| elif model_type.lower() == 'lstm': | |
| return lstm_forecast(input_data, steps=steps, epochs=epochs) | |
| else: | |
| return {"error": "Invalid model_type. Use 'arima' or 'lstm'."} | |