| from transformers import LlamaConfig, AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments | |
| from datasets import Dataset | |
| import pandas as pd | |
| import numpy as np | |
| import torch | |
| from pathlib import Path | |
| def generate_test_candles(n_samples=50): | |
| timestamps = pd.date_range(start="2024-01-01", periods=n_samples, freq="D") | |
| data = { | |
| "time": timestamps.astype(np.int64) // 10**9, | |
| "open": np.random.uniform(100, 150, n_samples), | |
| "high": np.random.uniform(150, 200, n_samples), | |
| "low": np.random.uniform(50, 100, n_samples), | |
| "close": np.random.uniform(100, 150, n_samples), | |
| "volume": np.random.uniform(1000, 5000, n_samples) | |
| } | |
| return pd.DataFrame(data) | |
| def prepare_training_data(examples): | |
| features = np.column_stack([ | |
| examples['time'], | |
| examples['open'], | |
| examples['high'], | |
| examples['low'], | |
| examples['close'], | |
| examples['volume'] | |
| ]) | |
| return { | |
| "input_ids": torch.tensor(features, dtype=torch.float32), | |
| "labels": torch.tensor(examples['close'], dtype=torch.float32) | |
| } | |
| Path("./results").mkdir(exist_ok=True) | |
| Path("./logs").mkdir(exist_ok=True) | |
| test_data = generate_test_candles() | |
| dataset = Dataset.from_pandas(test_data) | |
| dataset = dataset.map(prepare_training_data, batched=True) | |
| training_args = TrainingArguments( | |
| output_dir="./results", | |
| num_train_epochs=3, | |
| per_device_train_batch_size=4, | |
| learning_rate=2e-4, | |
| fp16=True, | |
| save_steps=100, | |
| logging_dir="./logs", | |
| logging_steps=10, | |
| remove_unused_columns=False | |
| ) | |
| config = LlamaConfig() | |
| model = AutoModelForCausalLM.from_pretrained( | |
| "tmmdev2/llama-vision-charts", | |
| config=config, | |
| trust_remote_code=True | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| "tmmdev2/llama-vision-charts", | |
| trust_remote_code=True | |
| ) | |
| trainer = Trainer( | |
| model=model, | |
| args=training_args, | |
| train_dataset=dataset, | |
| tokenizer=tokenizer | |
| ) | |
| results = trainer.train() | |
| model.save_pretrained("./results/final_model") | |
| tokenizer.save_pretrained("./results/final_model") | |
| print("Training completed!") | |
| print(f"Final loss: {results.training_loss}") | |