Update app.py

app.py

@@ -1,26 +1,39 @@
 from datasets import load_dataset
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator
 
-# Load the handwritten math dataset
+# Load the handwritten math dataset (1000 examples)
 ds = load_dataset("Azu/Handwritten-Mathematical-Expression-Convert-LaTeX", split="train[:1000]")
 
+# Load processor and model
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
 
+# Preprocess function
 def preprocess(ex):
     img = ex["image"].convert("RGB")
     inputs = processor(images=img, return_tensors="pt")
-    labels = processor.tokenizer(ex["label"], truncation=True, padding="max_length", max_length=128).input_ids
+
+    # Convert label index to actual LaTeX string
+    label_str = ds.features["label"].int2str(ex["label"])
+    labels = processor.tokenizer(
+        label_str,
+        truncation=True,
+        padding="max_length",
+        max_length=128
+    ).input_ids
+
     ex["pixel_values"] = inputs.pixel_values[0]
     ex["labels"] = labels
     return ex
 
+# Apply preprocessing
 ds = ds.map(preprocess, remove_columns=["image", "label"])
 
-
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+# Model config
 model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
 model.config.pad_token_id = processor.tokenizer.pad_token_id
 
+# Training arguments
 training_args = Seq2SeqTrainingArguments(
     output_dir="trained_model",
     per_device_train_batch_size=2,
@@ -32,6 +45,7 @@ training_args = Seq2SeqTrainingArguments(
     push_to_hub=False,
 )
 
+# Trainer
 trainer = Seq2SeqTrainer(
     model=model,
     args=training_args,
@@ -40,6 +54,10 @@ trainer = Seq2SeqTrainer(
     data_collator=default_data_collator,
 )
 
-trainer.train()
-model.save_pretrained("trained_model")
-processor.save_pretrained("trained_model")
+# Train and save
+if __name__ == "__main__":
+    print("🚀 Training started")
+    trainer.train()
+    print("✅ Training completed")
+    model.save_pretrained("trained_model")
+    processor.save_pretrained("trained_model")