# /// 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 from datasets import Dataset, load_dataset from transformers import ( AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments, DataCollatorForSeq2Seq ) import trackio import torch from tqdm import tqdm # --- DATA PREPARATION --- def process_data(tokenizer): print("Loading USDJPY data (streaming)...") ds_jpy = load_dataset("huggingXG/forex_USDJPY", split="train", streaming=True) data = [] # Collect 3M ticks (~2-3 months) max_ticks = 3000000 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 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({ "input_text": "forecast usdjpy: " + json.dumps(input_data), "target_text": json.dumps(output_data) }) def tokenize_fn(batch): model_inputs = tokenizer(batch["input_text"], max_length=1024, truncation=True) labels = tokenizer(text_target=batch["target_text"], max_length=512, truncation=True) model_inputs["labels"] = labels["input_ids"] return model_inputs full_ds = Dataset.from_list(examples) split_idx = int(len(full_ds) * 0.8) train_ds = full_ds.select(range(split_idx)) eval_ds = full_ds.select(range(split_idx, len(full_ds))) tokenized_train = train_ds.map(tokenize_fn, batched=True, remove_columns=full_ds.column_names) tokenized_eval = eval_ds.map(tokenize_fn, batched=True, remove_columns=full_ds.column_names) return tokenized_train, tokenized_eval, eval_ds # --- EVALUATION --- def run_custom_eval(model, tokenizer, raw_eval_ds): print("Running custom evaluation...") model.eval() results = [] eval_subset = raw_eval_ds.select(range(min(100, len(raw_eval_ds)))) for example in tqdm(eval_subset): inputs = tokenizer(example['input_text'], return_tensors="pt", max_length=1024, truncation=True).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 = json.loads(example['target_text']) is_valid = False schema_pass = False pred_json = {} try: pred_json = json.loads(prediction_text) is_valid = True if all(k in pred_json for k in ["forecast", "direction", "confidence"]): schema_pass = True except: pass results.append({"valid": is_valid, "schema": schema_pass, "actual": actual, "pred": pred_json if schema_pass else None}) valid_rate = np.mean([r['valid'] for r in results]) schema_rate = np.mean([r['schema'] for r in results]) valid_results = [r for r in results if r['schema']] mae, dir_acc = 0.0, 0.0 if valid_results: maes = [np.mean(np.abs(np.array(r['actual']['forecast']) - np.array(r['pred']['forecast']))) for r in valid_results] accs = [1 if r['actual']['direction'] == r['pred']['direction'] else 0 for r in valid_results] mae, dir_acc = np.mean(maes), np.mean(accs) return {"mae": float(mae), "directional_accuracy": float(dir_acc), "valid_json_rate": float(valid_rate), "schema_pass_rate": float(schema_rate)} def main(): trackio.init(project="usdjpy-forecasting", name="t5-efficient-base-sft") model_id = "google/t5-efficient-base" tokenizer = AutoTokenizer.from_pretrained(model_id) train_ds, tokenized_eval, raw_eval_ds = process_data(tokenizer) model = AutoModelForSeq2SeqLM.from_pretrained(model_id) args = Seq2SeqTrainingArguments( output_dir="./t5-base-usdjpy", max_steps=1000, learning_rate=1e-4, # Slightly lower for larger model stability per_device_train_batch_size=4, gradient_accumulation_steps=4, eval_strategy="steps", eval_steps=200, save_strategy="steps", save_steps=500, logging_steps=10, push_to_hub=True, hub_model_id="AbdelrehmanFouad/t5-efficient-base-usdjpy-forecaster", report_to="trackio", predict_with_generate=True, fp16=False, # Use FP32 for maximum stability on T4 ) trainer = Seq2SeqTrainer( model=model, args=args, train_dataset=train_ds, eval_dataset=tokenized_eval, processing_class=tokenizer, data_collator=DataCollatorForSeq2Seq(tokenizer, model=model), ) trainer.train() trainer.push_to_hub() final_metrics = run_custom_eval(trainer.model, tokenizer, raw_eval_ds) print(json.dumps(final_metrics, indent=2)) trackio.log(final_metrics) if __name__ == "__main__": main()