Spaces:
Runtime error
Runtime error
File size: 1,354 Bytes
c5ad565 211d8d1 c5ad565 211d8d1 c5ad565 211d8d1 c5ad565 211d8d1 c5ad565 211d8d1 c5ad565 211d8d1 c5ad565 211d8d1 c5ad565 211d8d1 c5ad565 211d8d1 c5ad565 211d8d1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
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')
|