Update app.py

app.py

@@ -1,35 +1,44 @@
 import torch
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Trainer, TrainingArguments
 from datasets import load_dataset
 
-# Load the dataset
-dataset = load_dataset('your_dataset_name')
+# Load your dataset
+dataset = load_dataset('your_dataset_name')  # Replace with your dataset name
 
-# Initialize the model and processor
+# Initialize the processor and model
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
 model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
 
-# Prepare the dataset for training
+# Preprocess the data
 def preprocess_data(example):
+    # Process images and texts
     pixel_values = processor(images=example['image'], return_tensors="pt").pixel_values
     labels = processor(text=example['text'], return_tensors="pt").input_ids
     return {'pixel_values': pixel_values, 'labels': labels}
 
+# Map preprocessing to the train dataset
 train_dataset = dataset['train'].map(preprocess_data)
 
-# Fine-tune the model
-training_args = {
-    'per_device_train_batch_size': 8,
-    'num_train_epochs': 3,
-    'logging_steps': 100,
-    'save_steps': 500,
-    'evaluation_strategy': 'steps',
-}
+# Training arguments
+training_args = TrainingArguments(
+    output_dir='./results',
+    per_device_train_batch_size=8,
+    num_train_epochs=3,
+    logging_steps=100,
+    save_steps=500,
+    evaluation_strategy='steps',
+)
 
+# Trainer setup
 trainer = Trainer(
     model=model,
     args=training_args,
     train_dataset=train_dataset,
 )
 
+# Train the model
 trainer.train()
+
+# Save the model and processor after training
+model.save_pretrained('./your_model_name')
+processor.save_pretrained('./your_model_name')