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 @spaces.GPU(duration=120) 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}." @spaces.GPU(duration=120) 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 ![The Logo](https://i.ibb.co/qBfBDr5/Whats-App-G-rsel-2024-08-24-saat-00-56-37-f466d65d.jpg) ### 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()