Spaces:
Running
on
Zero
Running
on
Zero
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,144 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import yfinance as yf
|
2 |
+
import pandas as pd
|
3 |
+
import numpy as np
|
4 |
+
import torch
|
5 |
+
import joblib
|
6 |
+
from tqdm import tqdm
|
7 |
+
from modeling_stockllama import StockLlamaForForecasting
|
8 |
+
from configuration_stockllama import StockLlamaConfig
|
9 |
+
from peft import LoraConfig, get_peft_model
|
10 |
+
from datasets import Dataset
|
11 |
+
import os
|
12 |
+
from transformers import Trainer, TrainingArguments
|
13 |
+
from huggingface_hub import login, upload_file
|
14 |
+
import wandb
|
15 |
+
import gradio as gr
|
16 |
+
import spaces
|
17 |
+
HF_TOKEN = os.getenv('HF_TOKEN')
|
18 |
+
WANDB_TOKEN = os.getenv('WANDB_TOKEN')
|
19 |
+
|
20 |
+
|
21 |
+
@spaces.GPU
|
22 |
+
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):
|
23 |
+
try:
|
24 |
+
stock_data = yf.download(stock_symbol, start=start_date, end=end_date, progress=False)
|
25 |
+
except Exception as e:
|
26 |
+
print(f"Error downloading data for {stock_symbol}: {e}")
|
27 |
+
return
|
28 |
+
|
29 |
+
data = stock_data["Close"]
|
30 |
+
|
31 |
+
class Scaler:
|
32 |
+
def __init__(self, feature_range):
|
33 |
+
self.feature_range = feature_range
|
34 |
+
self.min_df = None
|
35 |
+
self.max_df = None
|
36 |
+
|
37 |
+
def fit(self, df: pd.Series):
|
38 |
+
self.min_df = df.min()
|
39 |
+
self.max_df = df.max()
|
40 |
+
|
41 |
+
def transform(self, df: pd.Series) -> pd.Series:
|
42 |
+
min_val, max_val = self.feature_range
|
43 |
+
scaled_df = (df - self.min_df) / (self.max_df - self.min_df)
|
44 |
+
scaled_df = scaled_df * (max_val - min_val) + min_val
|
45 |
+
return scaled_df
|
46 |
+
|
47 |
+
def inverse_transform(self, X: np.ndarray) -> np.ndarray:
|
48 |
+
min_val, max_val = self.feature_range
|
49 |
+
min_x, max_x = np.min(X), np.max(X)
|
50 |
+
return (X - min_x) / (max_x - min_x) * (max_val - min_val) + min_val
|
51 |
+
|
52 |
+
scaler = Scaler(feature_range)
|
53 |
+
scaler.fit(data)
|
54 |
+
scaled_data = scaler.transform(data)
|
55 |
+
|
56 |
+
seq = [np.array(scaled_data[i:i + data_seq_length]) for i in range(len(scaled_data) - data_seq_length)]
|
57 |
+
target = [np.array(scaled_data[i + data_seq_length:i + data_seq_length + 1]) for i in range(len(scaled_data) - data_seq_length)]
|
58 |
+
|
59 |
+
seq_tensors = [torch.tensor(s, dtype=torch.float32).unsqueeze(0) for s in seq]
|
60 |
+
target_tensors = [t[0] for t in target]
|
61 |
+
|
62 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
63 |
+
model = StockLlamaForForecasting.from_pretrained("Q-bert/StockLlama").to(device)
|
64 |
+
config = LoraConfig(
|
65 |
+
r=64,
|
66 |
+
lora_alpha=32,
|
67 |
+
target_modules=["q_proj", "v_proj", "o_proj", "k_proj"],
|
68 |
+
lora_dropout=0.05,
|
69 |
+
bias="none",
|
70 |
+
task_type="CAUSAL_LM",
|
71 |
+
)
|
72 |
+
model = get_peft_model(model, config)
|
73 |
+
|
74 |
+
login(token=HF_TOKEN)
|
75 |
+
wandb.login(key=WANDB_TOKEN)
|
76 |
+
dct = {"input_ids": seq_tensors, "label": target_tensors}
|
77 |
+
dataset = Dataset.from_dict(dct)
|
78 |
+
dataset.push_to_hub(f"Q-bert/{stock_symbol}-{start_date}_{end_date}")
|
79 |
+
trainer = Trainer(
|
80 |
+
model=model,
|
81 |
+
train_dataset=dataset,
|
82 |
+
args=TrainingArguments(
|
83 |
+
per_device_train_batch_size=batch_size,
|
84 |
+
gradient_accumulation_steps=4,
|
85 |
+
num_train_epochs=epochs,
|
86 |
+
warmup_steps=5,
|
87 |
+
save_steps=100,
|
88 |
+
learning_rate=learning_rate,
|
89 |
+
fp16=True,
|
90 |
+
logging_steps=1,
|
91 |
+
push_to_hub=True,
|
92 |
+
report_to="wandb",
|
93 |
+
optim="adamw_torch",
|
94 |
+
weight_decay=0.01,
|
95 |
+
lr_scheduler_type="linear",
|
96 |
+
seed=3407,
|
97 |
+
output_dir=f"StockLlama-LoRA-{stock_symbol}",
|
98 |
+
),
|
99 |
+
)
|
100 |
+
|
101 |
+
trainer.train()
|
102 |
+
|
103 |
+
model = model.merge_and_unload()
|
104 |
+
model.push_to_hub(f"Q-bert/StockLlama-tuned-{stock_symbol}")
|
105 |
+
scaler_path = "scaler.joblib"
|
106 |
+
joblib.dump(scaler, scaler_path)
|
107 |
+
upload_file(
|
108 |
+
path_or_fileobj=scaler_path,
|
109 |
+
path_in_repo=f"scalers/{scaler_path}",
|
110 |
+
repo_id=f"Q-bert/StockLlama-tuned-{stock_symbol}"
|
111 |
+
)
|
112 |
+
@spaces.GPU
|
113 |
+
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):
|
114 |
+
feature_range = (feature_range_min, feature_range_max)
|
115 |
+
train_stock_model(
|
116 |
+
stock_symbol=stock_symbol,
|
117 |
+
start_date=start_date,
|
118 |
+
end_date=end_date,
|
119 |
+
feature_range=feature_range,
|
120 |
+
data_seq_length=data_seq_length,
|
121 |
+
epochs=epochs,
|
122 |
+
batch_size=batch_size,
|
123 |
+
learning_rate=learning_rate
|
124 |
+
)
|
125 |
+
return f"Training initiated for {stock_symbol} from {start_date} to {end_date}."
|
126 |
+
|
127 |
+
iface = gr.Interface(
|
128 |
+
fn=gradio_train_stock_model,
|
129 |
+
inputs=[
|
130 |
+
gr.Textbox(label="Stock Symbol", value="LUNC-USD"),
|
131 |
+
gr.Date(label="Start Date", value="2023-01-01"),
|
132 |
+
gr.Date(label="End Date", value="2024-08-24"),
|
133 |
+
gr.Slider(minimum=0, maximum=100, step=1, label="Feature Range Min", value=10),
|
134 |
+
gr.Slider(minimum=0, maximum=100, step=1, label="Feature Range Max", value=100),
|
135 |
+
gr.Slider(minimum=1, maximum=512, step=1, label="Data Sequence Length", value=256),
|
136 |
+
gr.Slider(minimum=1, maximum=50, step=1, label="Epochs", value=10),
|
137 |
+
gr.Slider(minimum=1, maximum=64, step=1, label="Batch Size", value=16),
|
138 |
+
gr.Slider(minimum=1e-5, maximum=1e-1, step=1e-5, label="Learning Rate", value=2e-4)
|
139 |
+
],
|
140 |
+
outputs="text",
|
141 |
+
live=True
|
142 |
+
)
|
143 |
+
|
144 |
+
iface.launch()
|