Create app.py

app.py

@@ -0,0 +1,144 @@
+import yfinance as yf
+import pandas as pd
+import numpy as np
+import torch
+import joblib
+from tqdm import tqdm
+from modeling_stockllama import StockLlamaForForecasting
+from configuration_stockllama import StockLlamaConfig
+from peft import LoraConfig, get_peft_model
+from datasets import Dataset
+import os
+from transformers import Trainer, TrainingArguments
+from huggingface_hub import login, upload_file
+import wandb
+import gradio as gr
+import spaces
+HF_TOKEN = os.getenv('HF_TOKEN')
+WANDB_TOKEN = os.getenv('WANDB_TOKEN')
+
+
+@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):
+    try:
+        stock_data = yf.download(stock_symbol, start=start_date, end=end_date, progress=False)
+    except Exception as e:
+        print(f"Error downloading data for {stock_symbol}: {e}")
+        return
+    
+    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)
+
+    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).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)
+    config = LoraConfig(
+        r=64,
+        lora_alpha=32,
+        target_modules=["q_proj", "v_proj", "o_proj", "k_proj"],
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    model = get_peft_model(model, config)
+
+    login(token=HF_TOKEN)
+    wandb.login(key=WANDB_TOKEN)
+    dct = {"input_ids": seq_tensors, "label": target_tensors}
+    dataset = Dataset.from_dict(dct)
+    dataset.push_to_hub(f"Q-bert/{stock_symbol}-{start_date}_{end_date}")
+    trainer = Trainer(
+        model=model,
+        train_dataset=dataset,
+        args=TrainingArguments(
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=4,
+            num_train_epochs=epochs,
+            warmup_steps=5,
+            save_steps=100,
+            learning_rate=learning_rate,
+            fp16=True,
+            logging_steps=1,
+            push_to_hub=True,
+            report_to="wandb",
+            optim="adamw_torch",
+            weight_decay=0.01,
+            lr_scheduler_type="linear",
+            seed=3407,
+            output_dir=f"StockLlama-LoRA-{stock_symbol}",
+        ),
+    )
+    
+    trainer.train()
+
+    model = model.merge_and_unload()
+    model.push_to_hub(f"Q-bert/StockLlama-tuned-{stock_symbol}")
+    scaler_path = "scaler.joblib"
+    joblib.dump(scaler, scaler_path)
+    upload_file(
+        path_or_fileobj=scaler_path,
+        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)
+    train_stock_model(
+        stock_symbol=stock_symbol,
+        start_date=start_date,
+        end_date=end_date,
+        feature_range=feature_range,
+        data_seq_length=data_seq_length,
+        epochs=epochs,
+        batch_size=batch_size,
+        learning_rate=learning_rate
+    )
+    return f"Training initiated for {stock_symbol} from {start_date} to {end_date}."
+
+iface = gr.Interface(
+    fn=gradio_train_stock_model,
+    inputs=[
+        gr.Textbox(label="Stock Symbol", value="LUNC-USD"),
+        gr.Date(label="Start Date", value="2023-01-01"),
+        gr.Date(label="End Date", value="2024-08-24"),
+        gr.Slider(minimum=0, maximum=100, step=1, label="Feature Range Min", value=10),
+        gr.Slider(minimum=0, maximum=100, step=1, label="Feature Range Max", value=100),
+        gr.Slider(minimum=1, maximum=512, step=1, label="Data Sequence Length", value=256),
+        gr.Slider(minimum=1, maximum=50, step=1, label="Epochs", value=10),
+        gr.Slider(minimum=1, maximum=64, step=1, label="Batch Size", value=16),
+        gr.Slider(minimum=1e-5, maximum=1e-1, step=1e-5, label="Learning Rate", value=2e-4)
+    ],
+    outputs="text",
+    live=True
+)
+
+iface.launch()