Delete train_ocr.py

train_ocr.py

@@ -1,85 +0,0 @@
-import os
-os.environ["HF_TOKEN"] = "lol nooo"
-import torch
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator
-from datasets import load_dataset
-from torch.utils.data import Dataset
-
-# 1. Hardware Check
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"--- GPU STATUS: {device.upper()} ---")
-
-# 2. Download Dataset
-print("Downloading Handwriting Dataset (~260MB)...")
-dataset = load_dataset("Teklia/IAM-line", split="train").train_test_split(test_size=0.1)
-
-# 3. Download Model
-print("Downloading TrOCR Base Model (~1.5GB)...")
-model_id = "microsoft/trocr-base-handwritten"
-processor = TrOCRProcessor.from_pretrained(model_id)
-model = VisionEncoderDecoderModel.from_pretrained(model_id).to(device)
-
-# 4. Prepare Dataset
-class HandwritingDataset(Dataset):
-    def __init__(self, hf_dataset, processor):
-        self.dataset = hf_dataset
-        self.processor = processor
-    def __len__(self):
-        return len(self.dataset)
-    def __getitem__(self, idx):
-        item = self.dataset[idx]
-        image = item["image"].convert("RGB")
-        text = item["text"]
-        pixel_values = self.processor(image, return_tensors="pt").pixel_values.squeeze()
-        labels = self.processor.tokenizer(text, padding="max_length", max_length=64, truncation=True).input_ids
-        labels = [label if label != self.processor.tokenizer.pad_token_id else -100 for label in labels]
-        return {"pixel_values": pixel_values, "labels": torch.tensor(labels)}
-
-train_dataset = HandwritingDataset(dataset['train'], processor)
-eval_dataset = HandwritingDataset(dataset['test'], processor)
-
-# --- CRITICAL OPTIMIZATIONS FOR 8GB VRAM ---
-model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
-model.config.pad_token_id = processor.tokenizer.pad_token_id
-model.config.vocab_size = model.config.decoder.vocab_size
-
-# Enable Gradient Checkpointing (The "Magic" VRAM Saver)
-model.gradient_checkpointing_enable() 
-
-# 5. Training Configuration
-training_args = Seq2SeqTrainingArguments(
-    predict_with_generate=True,
-    eval_strategy="steps",
-    per_device_train_batch_size=2,    # Small batches are safer
-    per_device_eval_batch_size=1,     # Extra safety during evaluation
-    gradient_accumulation_steps=8,    # Effectively a batch size of 16
-    output_dir="./working_checkpoints",
-    logging_steps=10,
-    save_steps=400,                   # Save more frequently just in case
-    eval_steps=400,
-    fp16=True,                        # Required for RTX 3060 Ti
-    max_steps=2000,
-    learning_rate=4e-5,
-    save_total_limit=3,               # Keep the 3 most recent checkpoints
-    dataloader_num_workers=0,         # Prevents Windows multi-threading errors
-    report_to="none"                  # Stops it from asking for a login (WandB/etc)
-)
-
-# 6. Start Training
-trainer = Seq2SeqTrainer(
-    model=model,
-    processing_class=processor.image_processor,
-    args=training_args,
-    train_dataset=train_dataset,
-    eval_dataset=eval_dataset,
-    data_collator=default_data_collator,
-)
-
-print("Starting training! Once the progress bar starts, you can safely walk away.")
-trainer.train()
-
-# 7. Save Final Model
-print("Saving final output...")
-trainer.save_model("./final_handwriting_model")
-processor.save_pretrained("./final_handwriting_model")
-print("DONE. See you tomorrow morning.")