Update app.py

app.py

@@ -1,89 +1,93 @@
 import os
 from datasets import load_dataset
-from transformers import (
-    TrOCRProcessor,
-    VisionEncoderDecoderModel,
-    Seq2SeqTrainer,
-    Seq2SeqTrainingArguments,
-    default_data_collator,
-)
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator
 
 # Check if model already exists
 if os.path.exists("trained_model"):
     print("✅ Model already exists. Skipping training.")
-else:
-    print("🚀 Starting training...")
-
-    # Load dataset (only 100 samples for speed)
-    ds = load_dataset("Azu/Handwritten-Mathematical-Expression-Convert-LaTeX", split="train[:100]")
-
-    processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
-
-    # First, filter out examples without valid labels
-    def is_valid(example):
-        label = example.get("label")
-        return isinstance(label, str) or (isinstance(label, dict) and "latex" in label)
-
-    ds = ds.filter(is_valid)
-
-    # Define preprocessing
-    def preprocess(example):
-        img = example["image"].convert("RGB")
-        inputs = processor(images=img, return_tensors="pt")
-
-        label_data = example["label"]
-        label_str = label_data["latex"] if isinstance(label_data, dict) else label_data
-
-        labels = processor.tokenizer(
-            label_str,
-            truncation=True,
-            padding="max_length",
-            max_length=128
-        ).input_ids
-
-        return {
-            "pixel_values": inputs.pixel_values[0],
-            "labels": labels
-        }
-
-    # Apply preprocessing
-    ds = ds.map(preprocess, remove_columns=["image", "label"])
-
-    # Final safety check
-    if len(ds) == 0:
-        raise RuntimeError("❌ No usable training samples after preprocessing.")
-
-    # Load model
-    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 arguments
-    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,  # Required False for CPU-only env
-        push_to_hub=False,
-    )
-
-    # Trainer
-    trainer = Seq2SeqTrainer(
-        model=model,
-        args=training_args,
-        train_dataset=ds,
-        tokenizer=processor.tokenizer,
-        data_collator=default_data_collator,
-    )
-
-    # Start training
-    trainer.train()
-    print("✅ Training completed")
-
-    # Save model + processor
-    model.save_pretrained("trained_model")
-    processor.save_pretrained("trained_model")
-    print("✅ Model saved to trained_model/")
+    exit()
+
+print("🚀 Starting training...")
+
+# Load only 100 samples for faster CPU training
+ds = load_dataset("Azu/Handwritten-Mathematical-Expression-Convert-LaTeX", split="train[:100]")
+
+# DEBUG: Inspect a few labels
+print("\n🔍 Sample labels from dataset:")
+for i in range(5):
+    print(f"{i}: {ds[i]['label']} (type: {type(ds[i]['label'])})")
+
+processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+
+# Safely extract label string from possible dict or str
+def safe_get_label(example):
+    label = example.get("label")
+    if isinstance(label, dict) and "latex" in label:
+        return label["latex"]
+    elif isinstance(label, str):
+        return label
+    else:
+        return None
+
+def preprocess(example):
+    label_str = safe_get_label(example)
+    if not isinstance(label_str, str) or label_str.strip() == "":
+        return {}  # Skip if label is invalid
+
+    # Convert image to RGB
+    img = example["image"].convert("RGB")
+    inputs = processor(images=img, return_tensors="pt")
+
+    # Tokenize label
+    labels = processor.tokenizer(
+        label_str,
+        truncation=True,
+        padding="max_length",
+        max_length=128
+    ).input_ids
+
+    return {
+        "pixel_values": inputs.pixel_values[0],
+        "labels": labels
+    }
+
+# Preprocess and filter
+ds = ds.map(preprocess, remove_columns=["image", "label"])
+ds = ds.filter(lambda ex: "labels" in ex and ex["labels"] is not None)
+
+# Check number of remaining examples
+print(f"✅ Total usable training samples: {len(ds)}")
+if len(ds) == 0:
+    raise RuntimeError("❌ No usable training samples after preprocessing.")
+
+# Model setup
+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,
+)
+
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=ds,
+    tokenizer=processor.tokenizer,
+    data_collator=default_data_collator,
+)
+
+trainer.train()
+print("✅ Training completed")
+
+# Save model
+model.save_pretrained("trained_model")
+processor.save_pretrained("trained_model")
+print("✅ Model saved to trained_model/")