text generation model

trocr.py

@@ -0,0 +1,30 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class IAMDataset(Dataset):
+    def __init__(self, crops, processor):
+        self.crops = crops
+        self.processor = processor
+
+    def __len__(self):
+        return len(self.crops)
+
+    def __getitem__(self, idx):
+        crp = self.crops[idx]
+        pixel_values = self.processor(crp, return_tensors="pt").pixel_values
+        encoding = {"pixel_values": pixel_values.squeeze()}
+        return encoding
+
+def get_processor_model(checkpoint:str):
+    rec_processor = TrOCRProcessor.from_pretrained('trocr_printed_processor/')
+    rec_model = VisionEncoderDecoderModel.from_pretrained('trocr_printed_model/')
+    rec_model.config.eos_token_id = 2
+    rec_model.config.pad_token_id = 2
+    rec_model.to(device)
+    rec_model.eval()
+    return rec_processor, rec_model