Spaces:
Running
on
Zero
Running
on
Zero
| 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, hf_hub_download | |
| import wandb | |
| import gradio as gr | |
| import spaces | |
| from huggingface_hub import HfApi | |
| hf_api = HfApi() | |
| HF_TOKEN = os.getenv('HF_TOKEN') | |
| WANDB_TOKEN = os.getenv('WANDB_TOKEN') | |
| login(token=HF_TOKEN) | |
| wandb.login(key=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 | |
| def check_existing_model(stock_symbol, start_date, end_date): | |
| repo_id = f"Q-bert/StockLlama-tuned-{stock_symbol}-{stock_symbol}-{start_date}_{end_date}" | |
| state = repo_id in [model.modelId for model in hf_api.list_models()] | |
| return state | |
| 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): | |
| repo_id = f"Q-bert/StockLlama-tuned{stock_symbol}-{start_date}_{end_date}" | |
| if check_existing_model(stock_symbol, start_date, end_date): | |
| return f"Model for {stock_symbol} from {start_date} to {end_date} already exists." | |
| 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"] | |
| 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) 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("StockLlama/StockLlama-base-v1").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) | |
| dct = {"input_ids": seq_tensors, "label": target_tensors} | |
| dataset = Dataset.from_dict(dct) | |
| dataset.push_to_hub(f"StockLlama/{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=10, | |
| 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/StockLlama-LoRA-{stock_symbol}-{start_date}_{end_date}", | |
| ), | |
| ) | |
| trainer.train() | |
| model = model.merge_and_unload() | |
| model.push_to_hub(f"StockLlama/StockLlama-tuned-{stock_symbol}-{start_date}_{end_date}") | |
| 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"StockLlama/StockLlama-tuned-{stock_symbol}-{start_date}_{end_date}" | |
| ) | |
| return f"Training completed and model saved for {stock_symbol} from {start_date} to {end_date}." | |
| 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) | |
| result = 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 result | |
| title = "StockLlama-TrainOnAnyStock" | |
| description = """ | |
| ## StockLlama | |
|  | |
| ### Description | |
| StockLlama is a time series forecasting pre-trained model based on Llama, enhanced with custom embeddings for improved accuracy. | |
| ### How It Works | |
| **Data Collection:** The model retrieves historical stock price data using the yfinance library. Users specify the stock symbol, date range, and other parameters through a Gradio interface. | |
| **Data Preprocessing:** The collected stock prices are scaled to a specified range using a custom Scaler class. The data is then divided into sequences of a defined length, with each sequence serving as input to the model and the next data point as the target. | |
| **Model Architecture:** StockLlama is a modified version of the Llama model, specifically tailored for time series forecasting. The model is enhanced with custom embeddings and fine-tuned using a LoRA (Low-Rank Adaptation) configuration, allowing for efficient training on the specific stock data. | |
| **Training Process:** The training is managed using the Hugging Face Trainer class. The model learns to predict the next data point in the sequence, optimizing its weights over multiple epochs. The training process can be monitored via Weights & Biases integration. | |
| **Deployment:** After training, the model is pushed to the Hugging Face Hub, making it accessible for future predictions. The scaler used for data normalization is also saved and uploaded, ensuring that new data can be correctly transformed and predictions can be accurately descaled. | |
| ### Contributing | |
| Contributions to this project are welcome! If you find any issues or want to add new features, feel free to open an issue or submit a pull request. | |
| ### License | |
| This project is licensed under the [Apache 2.0 License](https://opensource.org/license/apache-2-0). | |
| ### Credits | |
| The StockLlama model used in this project is based on the work by [Talha Rüzgar Akkuş](https://www.linkedin.com/in/talha-r%C3%BCzgar-akku%C5%9F-1b5457264/). | |
| """ | |
| iface = gr.Interface( | |
| fn=gradio_train_stock_model, | |
| inputs=[ | |
| gr.Textbox(label="Stock Symbol", value="BTC-USD"), | |
| gr.Textbox(label="Start Date", value="2023-01-01"), | |
| gr.Textbox(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) | |
| ], | |
| description=description, | |
| title=title, | |
| outputs="text", | |
| ) | |
| iface.launch() |