File size: 2,242 Bytes
6574d6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | 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}")
|