train_job.py

# /// 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()