Upload finetune.py

finetune.py

@@ -0,0 +1,136 @@
+import pandas as pd
+from datasets import load_dataset
+from sklearn.model_selection import train_test_split
+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainingArguments, Seq2SeqTrainer
+from PIL import Image
+import io
+import numpy as np
+
+device = 'mps:0'
+# Load the dataset and filter for Latin entries
+dataset = load_dataset("CATMuS/medieval", split='train')
+# latin_dataset = dataset.filter(lambda example: example['language'] == 'Latin')
+latin_dataset = dataset.filter(lambda example: example['language'] == 'Latin' and example['script_type'] == 'Caroline')
+
+print(latin_dataset)
+# Convert to pandas DataFrame for easier manipulation
+df = pd.DataFrame(latin_dataset)
+
+# Split the data into training and testing sets
+train_df, test_df = train_test_split(df, test_size=0.2)
+train_df.reset_index(drop=True, inplace=True)
+test_df.reset_index(drop=True, inplace=True)
+
+# Define the dataset class
+class HandwrittenTextDataset(Dataset):
+    def __init__(self, df, processor, max_target_length=128):
+        self.df = df
+        self.processor = processor
+        self.max_target_length = max_target_length
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        image_data = self.df['im'][idx]
+        text = self.df['text'][idx]
+
+        # Convert array to PIL image
+        image = Image.fromarray(np.array(image_data)).convert("RGB")
+
+        # Prepare image (i.e., resize + normalize)
+        pixel_values = self.processor(images=image, return_tensors="pt").pixel_values
+
+        # Add labels (input_ids) by encoding the text
+        labels = self.processor.tokenizer(text, 
+                                          padding="max_length", 
+                                          max_length=self.max_target_length,
+                                          truncation=True).input_ids
+        # Important: make sure that PAD tokens are ignored by the loss function
+        labels = [label if label != self.processor.tokenizer.pad_token_id else -100 for label in labels]
+
+        encoding = {"pixel_values": pixel_values.squeeze(), "labels": torch.tensor(labels)}
+        return encoding
+# Instantiate processor and dataset
+processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+train_dataset = HandwrittenTextDataset(df=train_df, processor=processor)
+eval_dataset = HandwrittenTextDataset(df=test_df, processor=processor)
+
+# Create corresponding dataloaders
+train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)
+eval_dataloader = DataLoader(eval_dataset, batch_size=4)
+
+# Load the model
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+
+# Set special tokens used for creating the decoder_input_ids from the labels
+model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
+model.config.pad_token_id = processor.tokenizer.pad_token_id
+
+# Make sure vocab size is set correctly
+model.config.vocab_size = model.config.decoder.vocab_size
+
+# Set beam search parameters
+model.config.eos_token_id = processor.tokenizer.sep_token_id
+model.config.max_length = 64
+model.config.early_stopping = True
+model.config.no_repeat_ngram_size = 3
+model.config.length_penalty = 2.0
+model.config.num_beams = 4
+
+# Training arguments
+training_args = Seq2SeqTrainingArguments(
+    output_dir="./results",
+    per_device_train_batch_size=4,
+    num_train_epochs=10,
+    logging_steps=1000,
+    save_steps=1000,
+    evaluation_strategy="steps",
+    save_total_limit=2,
+    predict_with_generate=True,
+    fp16=False,  # Set to True if using a compatible GPU
+)
+
+# Trainer
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+)
+
+# Train the model
+trainer.train()
+
+# After training, save both the model and the processor
+model.save_pretrained("./finetuned_model")
+processor.save_pretrained("./finetuned_model")
+
+from datasets import load_metric
+
+cer_metric = load_metric("cer")
+
+def compute_cer(pred_ids, label_ids):
+    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
+    label_ids[label_ids == -100] = processor.tokenizer.pad_token_id
+    label_str = processor.batch_decode(label_ids, skip_special_tokens=True)
+
+    cer = cer_metric.compute(predictions=pred_str, references=label_str)
+
+    return cer
+
+# Evaluation
+model.eval()
+valid_cer = 0.0
+with torch.no_grad():
+    for batch in eval_dataloader:
+        # Run batch generation
+        outputs = model.generate(batch["pixel_values"].to(device))
+        # Compute metrics
+        cer = compute_cer(pred_ids=outputs, label_ids=batch["labels"])
+        valid_cer += cer
+
+print("Validation CER:", valid_cer / len(eval_dataloader))
+