Spaces:
Configuration error
Configuration error
File size: 1,733 Bytes
151ec26 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | from pathlib import Path
import torch
from transformers import (
TrOCRProcessor,
VisionEncoderDecoderModel,
)
BASE_DIR = Path(__file__).resolve().parents[1]
MODEL_PATH = (
BASE_DIR
/ "models"
/ "pretrained"
/ "trocr"
)
print("=" * 70)
print("SANJEEVANI - TrOCR LOAD TEST")
print("=" * 70)
print("Model:", MODEL_PATH)
print(
"CUDA:",
torch.cuda.is_available()
)
if torch.cuda.is_available():
print(
"GPU:",
torch.cuda.get_device_name(0)
)
# ------------------------------------------------------------
# PROCESSOR
# ------------------------------------------------------------
print("\nLoading processor...")
processor = TrOCRProcessor.from_pretrained(
str(MODEL_PATH)
)
print("Processor loaded successfully.")
# ------------------------------------------------------------
# MODEL
# ------------------------------------------------------------
print("\nLoading model...")
model = VisionEncoderDecoderModel.from_pretrained(
str(MODEL_PATH)
)
print("Model loaded successfully.")
# ------------------------------------------------------------
# GPU
# ------------------------------------------------------------
if torch.cuda.is_available():
model = model.to("cuda")
print(
"\nModel moved to GPU successfully."
)
# ------------------------------------------------------------
# SUMMARY
# ------------------------------------------------------------
print("\n" + "=" * 70)
print("SUCCESS")
print("=" * 70)
print("TrOCR processor: OK")
print("TrOCR model: OK")
if torch.cuda.is_available():
print("CUDA: OK")
print(
"GPU: "
+torch.cuda.get_device_name(0)
)
print("=" * 70)
|