Update app.py

app.py

@@ -35,16 +35,16 @@ dataset = project.version(1).download("folder")
 subprocess.run(['wget', '--no-check-certificate', 'https://docs.google.com/uc?export=download&id=12reT7rxiRqTERYqeKYx7WGz5deMXjnEo', '-O', 'filetxt'])
 subprocess.run(['unzip', 'filetxt'])
 
-# def seed_everything(seed_value):
-#     np.random.seed(seed_value)
-#     torch.manual_seed(seed_value)
-#     torch.cuda.manual_seed_all(seed_value)
-#     torch.backends.cudnn.deterministic = True
-#     torch.backends.cudnn.benchmark = False
+def seed_everything(seed_value):
+    np.random.seed(seed_value)
+    torch.manual_seed(seed_value)
+    torch.cuda.manual_seed_all(seed_value)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
 
-# seed_everything(42)
+seed_everything(42)
 
-# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 def download_and_unzip(url, save_path):
     print(f"Downloading and extracting assets....", end="")
@@ -164,7 +164,7 @@ valid_dataset = CustomOCRDataset(
 )
 
 model = VisionEncoderDecoderModel.from_pretrained(ModelConfig.MODEL_NAME)
-#model.to(device)
+model.to(device)
 print(model)
 # Total parameters and trainable parameters.
 total_params = sum(p.numel() for p in model.parameters())
@@ -237,8 +237,7 @@ trainer = Seq2SeqTrainer(
 res = trainer.train()
 
 processor = TrOCRProcessor.from_pretrained(ModelConfig.MODEL_NAME)
-trained_model = VisionEncoderDecoderModel.from_pretrained('seq2seq_model_printed/checkpoint-'+str(res.global_step))
-#.to(device)
+trained_model = VisionEncoderDecoderModel.from_pretrained('seq2seq_model_printed/checkpoint-'+str(res.global_step)).to(device)
 
 def read_and_show(image_path):
     """
@@ -262,8 +261,7 @@ def ocr(image, processor, model):
         generated_text: the OCR'd text string.
     """
     # We can directly perform OCR on cropped images.
-    pixel_values = processor(image, return_tensors='pt').pixel_values
-    #.to(device)
+    pixel_values = processor(image, return_tensors='pt').pixel_values.to(device)
     generated_ids = model.generate(pixel_values)
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return generated_text