Update app.py

app.py

@@ -17,6 +17,26 @@ import spaces
 HF_TOKEN = os.getenv('HF_TOKEN')
 WANDB_TOKEN = os.getenv('WANDB_TOKEN')
 
+class Scaler:
+    def __init__(self, feature_range):
+        self.feature_range = feature_range
+        self.min_df = None
+        self.max_df = None
+
+    def fit(self, df: pd.Series):
+        self.min_df = df.min()
+        self.max_df = df.max()
+
+    def transform(self, df: pd.Series) -> pd.Series:
+        min_val, max_val = self.feature_range
+        scaled_df = (df - self.min_df) / (self.max_df - self.min_df)
+        scaled_df = scaled_df * (max_val - min_val) + min_val
+        return scaled_df
+
+    def inverse_transform(self, X: np.ndarray) -> np.ndarray:
+        min_val, max_val = self.feature_range
+        min_x, max_x = np.min(X), np.max(X)
+        return (X - min_x) / (max_x - min_x) * (max_val - min_val) + min_val
 
 @spaces.GPU
 def train_stock_model(stock_symbol, start_date, end_date, feature_range=(10, 100), data_seq_length=256, epochs=10, batch_size=16, learning_rate=2e-4):
@@ -28,27 +48,6 @@ def train_stock_model(stock_symbol, start_date, end_date, feature_range=(10, 100
     
     data = stock_data["Close"]
 
-    class Scaler:
-        def __init__(self, feature_range):
-            self.feature_range = feature_range
-            self.min_df = None
-            self.max_df = None
-
-        def fit(self, df: pd.Series):
-            self.min_df = df.min()
-            self.max_df = df.max()
-
-        def transform(self, df: pd.Series) -> pd.Series:
-            min_val, max_val = self.feature_range
-            scaled_df = (df - self.min_df) / (self.max_df - self.min_df)
-            scaled_df = scaled_df * (max_val - min_val) + min_val
-            return scaled_df
-
-        def inverse_transform(self, X: np.ndarray) -> np.ndarray:
-            min_val, max_val = self.feature_range
-            min_x, max_x = np.min(X), np.max(X)
-            return (X - min_x) / (max_x - min_x) * (max_val - min_val) + min_val
-
     scaler = Scaler(feature_range)
     scaler.fit(data)
     scaled_data = scaler.transform(data)
@@ -56,12 +55,11 @@ def train_stock_model(stock_symbol, start_date, end_date, feature_range=(10, 100
     seq = [np.array(scaled_data[i:i + data_seq_length]) for i in range(len(scaled_data) - data_seq_length)]
     target = [np.array(scaled_data[i + data_seq_length:i + data_seq_length + 1]) for i in range(len(scaled_data) - data_seq_length)]
 
-    seq_tensors = [torch.tensor(s, dtype=torch.float32) for s in seq]
+    seq_tensors = [torch.tensor(s, dtype=torch.float32).unsqueeze(0) for s in seq]
     target_tensors = [t[0] for t in target]
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model = StockLlamaForForecasting.from_pretrained("Q-bert/StockLlama").to(device)
-    print("Model Installed.")
     config = LoraConfig(
         r=64,
         lora_alpha=32,
@@ -71,7 +69,7 @@ def train_stock_model(stock_symbol, start_date, end_date, feature_range=(10, 100
         task_type="CAUSAL_LM",
     )
     model = get_peft_model(model, config)
-    print("Model pefted.")
+
     login(token=HF_TOKEN)
     wandb.login(key=WANDB_TOKEN)
     dct = {"input_ids": seq_tensors, "label": target_tensors}
@@ -110,6 +108,7 @@ def train_stock_model(stock_symbol, start_date, end_date, feature_range=(10, 100
         path_in_repo=f"scalers/{scaler_path}",
         repo_id=f"Q-bert/StockLlama-tuned-{stock_symbol}"
     )
+
 @spaces.GPU
 def gradio_train_stock_model(stock_symbol, start_date, end_date, feature_range_min, feature_range_max, data_seq_length, epochs, batch_size, learning_rate):
     feature_range = (feature_range_min, feature_range_max)