Finetune_Models / app.py
arshadrana's picture
Update app.py
211d8d1 verified
import torch
from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Trainer, TrainingArguments
from datasets import load_dataset
# Load your dataset
dataset = load_dataset('your_dataset_name') # Replace with your dataset name
# Initialize the processor and model
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
# 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)
# 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')