# /// script # dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "pandas", "numpy", "datasets", "transformers", "torch", "accelerate", "huggingface_hub"] # /// import os import json import pandas as pd import numpy as np from datetime import datetime, timedelta from datasets import Dataset, load_dataset from transformers import AutoTokenizer, AutoModelForSeq2SeqLM from trl import SFTTrainer, SFTConfig import trackio from huggingface_hub import HfApi import torch from tqdm import tqdm # --- DATA PREPARATION --- def process_data(): print("Loading USDJPY data (streaming)...") ds_jpy = load_dataset("huggingXG/forex_USDJPY", split="train", streaming=True) data = [] max_ticks = 2000000 for i, row in enumerate(ds_jpy): data.append(row) if i >= max_ticks: break df_jpy = pd.DataFrame(data) df_jpy['timestamp'] = pd.to_datetime(df_jpy['timestamp'], format='ISO8601') df_jpy.set_index('timestamp', inplace=True) print("Resampling to hourly OHLCV...") df_jpy['mid'] = (df_jpy['ask'] + df_jpy['bid']) / 2 resampled = df_jpy['mid'].resample('1H').ohlc() resampled.dropna(inplace=True) print("Loading Calendar data...") ds_cal = load_dataset("Ehsanrs2/Forex_Factory_Calendar", split="train") df_cal = ds_cal.to_pandas() df_cal['DateTime'] = pd.to_datetime(df_cal['DateTime'], utc=True) df_cal = df_cal[df_cal['Currency'].isin(['USD', 'JPY'])] print("Generating training examples...") examples = [] context_size = 168 horizon_size = 24 if resampled.index.tz is None: resampled.index = resampled.index.tz_localize('UTC') else: resampled.index = resampled.index.tz_convert('UTC') indices = range(context_size, len(resampled) - horizon_size, 1) for i in indices: context_window = resampled.iloc[i-context_size:i]['close'].tolist() target_window = resampled.iloc[i:i+horizon_size]['close'].tolist() start_time = resampled.index[i] end_time = resampled.index[i+horizon_size-1] events_in_horizon = df_cal[(df_cal['DateTime'] >= start_time) & (df_cal['DateTime'] <= end_time)] macro_events = [] for _, event in events_in_horizon.iterrows(): macro_events.append({ "time": event['DateTime'].strftime("%Y-%m-%dT%H:%M:%SZ"), "event": event['Event'], "impact": event['Impact'].lower().split()[0] }) start_price = context_window[-1] end_price = target_window[-1] change = (end_price - start_price) / start_price if change > 0.0005: direction = "up" elif change < -0.0005: direction = "down" else: direction = "sideways" confidence = min(0.99, abs(change) * 50 + 0.5) high_impact = [e['event'] for e in macro_events if 'high' in e['impact']] macro_note = f"{high_impact[0]} in window" if high_impact else "No high impact events" input_data = { "context_window": [round(float(c), 4) for c in context_window], "macro_events": macro_events } output_data = { "forecast": [round(float(f), 4) for f in target_window], "direction": direction, "confidence": round(float(confidence), 2), "macro_context": macro_note } examples.append({ "prompt": json.dumps(input_data), "completion": json.dumps(output_data) }) split_idx = int(len(examples) * 0.8) return Dataset.from_list(examples[:split_idx]), Dataset.from_list(examples[split_idx:]) # --- EVALUATION --- def run_custom_eval(model, tokenizer, eval_dataset): print("Running custom evaluation metrics...") model.eval() results = [] # Sample a portion of eval dataset for speed if it's too large eval_subset = eval_dataset.select(range(min(200, len(eval_dataset)))) for example in tqdm(eval_subset): inputs = tokenizer(example['prompt'], return_tensors="pt").to(model.device) with torch.no_grad(): outputs = model.generate(**inputs, max_new_tokens=512) prediction_text = tokenizer.decode(outputs[0], skip_special_tokens=True) actual_output = json.loads(example['completion']) # Parse prediction is_valid_json = False schema_pass = False pred_json = {} try: pred_json = json.loads(prediction_text) is_valid_json = True # Check schema required_keys = {"forecast", "direction", "confidence", "macro_context"} if all(k in pred_json for k in required_keys): schema_pass = True except: pass results.append({ "valid_json": is_valid_json, "schema_pass": schema_pass, "actual": actual_output, "predicted": pred_json if schema_pass else None }) # Calculate metrics valid_json_rate = np.mean([r['valid_json'] for r in results]) schema_pass_rate = np.mean([r['schema_pass'] for r in results]) # Filters only for schema passes valid_results = [r for r in results if r['schema_pass']] mae = 0 dir_acc = 0 conf_corr = 0 if valid_results: maes = [] dir_matches = [] confidences = [] actual_accuracies = [] for r in valid_results: a = r['actual'] p = r['predicted'] # MAE on forecast maes.append(np.mean(np.abs(np.array(a['forecast']) - np.array(p['forecast'])))) # Directional Accuracy match = 1 if a['direction'] == p['direction'] else 0 dir_matches.append(match) confidences.append(p['confidence']) actual_accuracies.append(match) mae = np.mean(maes) dir_acc = np.mean(dir_matches) if len(confidences) > 1: conf_corr = np.corrcoef(confidences, actual_accuracies)[0, 1] metrics = { "mae": float(mae), "directional_accuracy": float(dir_acc), "valid_json_rate": float(valid_json_rate), "schema_pass_rate": float(schema_pass_rate), "confidence_calibration": float(conf_corr) } return metrics def main(): trackio.init(project="usdjpy-forecasting", name="t5-efficient-small-sft") train_ds, eval_ds = process_data() model_id = "google/t5-efficient-small" tokenizer = AutoTokenizer.from_pretrained(model_id) training_args = SFTConfig( output_dir="./t5-usdjpy-sft", max_steps=1000, learning_rate=2e-4, lr_scheduler_type="linear", warmup_ratio=0.03, per_device_train_batch_size=4, gradient_accumulation_steps=4, eval_strategy="steps", eval_steps=100, save_strategy="steps", save_steps=500, logging_steps=10, push_to_hub=True, hub_model_id="AbdelrehmanFouad/t5-efficient-small-usdjpy-forecaster-sft", report_to="trackio", max_length=512, fp16=True if torch.cuda.is_available() else False, ) trainer = SFTTrainer( model=model_id, train_dataset=train_ds, eval_dataset=eval_ds, args=training_args, processing_class=tokenizer, ) trainer.train() trainer.push_to_hub() # Custom Evaluation final_metrics = run_custom_eval(trainer.model, tokenizer, eval_ds) print("Final Custom Metrics:") print(json.dumps(final_metrics, indent=2)) # Log to Trackio trackio.log(final_metrics) if __name__ == "__main__": main()