Initialize and return an image model and its feature extractor. Args: device (str): The device to run the model on. Returns: tuple: A tuple containing the feature extractor and the image model.
(device: str = None)
| 70 | |
| 71 | |
| 72 | def get_image_model(device: str = None): |
| 73 | """ |
| 74 | Initialize and return an image model and its feature extractor. |
| 75 | |
| 76 | Args: |
| 77 | device (str): The device to run the model on. |
| 78 | |
| 79 | Returns: |
| 80 | tuple: A tuple containing the feature extractor and the image model. |
| 81 | """ |
| 82 | model_base = "google/vit-base-patch16-224-in21k" |
| 83 | return ( |
| 84 | AutoFeatureExtractor.from_pretrained( |
| 85 | model_base, |
| 86 | torch_dtype=torch.float16, |
| 87 | device_map=device, |
| 88 | ), |
| 89 | AutoModel.from_pretrained( |
| 90 | model_base, |
| 91 | torch_dtype=torch.float16, |
| 92 | device_map=device, |
| 93 | ).eval(), |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | def parse_pdf( |
no outgoing calls
no test coverage detected