train_job_base.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
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()