Update train.py

train.py

@@ -1,30 +1,30 @@
-# train.py
-import time
-import yaml
+import os
+import torch
+from pix2tex.dataset.latex_dataset import Im2LatexDataset
+from pix2tex.models import get_model
+from pix2tex.trainer import build_trainer
+from pix2tex.utils import set_seed, get_config
+from pix2tex.tokenizer import LatexTokenizer
 
-def train_model(config_file="train.yaml"):
-    print("🧠 Starting training...")
+# Load config
+config = get_config("train.yaml")
 
-    try:
-        with open(config_file, "r") as f:
-            config = yaml.safe_load(f)
-    except Exception as e:
-        print("⚠️ Failed to load config:", str(e))
-        return f"❌ Failed to load config: {str(e)}"
+# Set CPU-only if CUDA not available or forced
+config["device"] = "cpu"
+torch.set_default_tensor_type('torch.FloatTensor')
 
-    model_name = config.get("model", {}).get("name", "default_model")
-    epochs = config.get("training", {}).get("epochs", 5)
-    lr = config.get("training", {}).get("learning_rate", 0.001)
-    batch_size = config.get("training", {}).get("batch_size", 32)
-    device = config.get("training", {}).get("device", "cpu")
+set_seed(config.get("seed", 42))
 
-    print(f"📦 Model: {model_name}")
-    print(f"🔧 Device: {device}")
-    print(f"📚 Epochs: {epochs}, Batch Size: {batch_size}, Learning Rate: {lr}")
+# Tokenizer
+tokenizer = LatexTokenizer(config["tokenizer_path"])
 
-    for epoch in range(1, epochs + 1):
-        print(f"🌀 Epoch {epoch}/{epochs} ...")
-        time.sleep(1)  # Simulate work
+# Dataset
+trainset = Im2LatexDataset(config["train_csv"], tokenizer, config)
+valset = Im2LatexDataset(config["val_csv"], tokenizer, config, is_val=True)
 
-    print("✅ Training complete.")
-    return f"✅ Dummy training for `{model_name}` finished on `{device}`!"
+# Model
+model = get_model(config, tokenizer)
+
+# Trainer
+trainer = build_trainer(model, tokenizer, config, trainset=trainset, valset=valset)
+trainer.train()