Update train.py

train.py

@@ -1,33 +1,44 @@
-from model import get_model
-from pix2tex.dataset.dataset import Im2LatexDataset
-from pix2tex.trainer import Trainer
-import os
-
-os.makedirs('trained_model', exist_ok=True)
-
-# Training parameters
-config = {
-    "batch_size": 4,
-    "epochs": 1,
-    "max_seq_len": 150,
-    "warmup_steps": 10,
-    "lr": 1e-4,
-    "device": "cpu",
-    "save_dir": "trained_model",
-    "resume": False
-}
-
-# Dataset path
-dataset = Im2LatexDataset(
-    data_root='handwritten_dataset',
-    transform=None,
-    max_length=config["max_seq_len"]
+from datasets import load_dataset
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator
+
+# Load the handwritten math dataset
+ds = load_dataset("Azu/Handwritten-Mathematical-Expression-Convert-LaTeX", split="train[:1000]")
+
+processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+
+def preprocess(ex):
+    img = ex["image"].convert("RGB")
+    inputs = processor(images=img, return_tensors="pt")
+    labels = processor.tokenizer(ex["text"], truncation=True, padding="max_length", max_length=128).input_ids
+    ex["pixel_values"] = inputs.pixel_values[0]
+    ex["labels"] = labels
+    return ex
+
+ds = ds.map(preprocess, remove_columns=["image", "text"])
+
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
+model.config.pad_token_id = processor.tokenizer.pad_token_id
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir="trained_model",
+    per_device_train_batch_size=2,
+    num_train_epochs=1,
+    learning_rate=5e-5,
+    logging_steps=10,
+    save_steps=500,
+    fp16=False,
+    push_to_hub=False,
 )
 
-# Initialize model and trainer
-model, tokenizer = get_model()
-trainer = Trainer(model, tokenizer, config)
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=ds,
+    tokenizer=processor.tokenizer,
+    data_collator=default_data_collator,
+)
 
-print("🧠 Starting training...")
-trainer.train(dataset)
-print("✅ Training complete. Model saved to 'trained_model/'")
+trainer.train()
+model.save_pretrained("trained_model")
+processor.save_pretrained("trained_model")