(cls, pretrained_model_name_or_path, *args, **kwargs)
| 145 | |
| 146 | @classmethod |
| 147 | def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): |
| 148 | # Allow loading models directly from the hub, or using s3 |
| 149 | if not pretrained_model_name_or_path.startswith(cls.s3_prefix): |
| 150 | return super().from_pretrained( |
| 151 | pretrained_model_name_or_path, *args, **kwargs |
| 152 | ) |
| 153 | |
| 154 | local_path = cls.get_local_path(pretrained_model_name_or_path) |
| 155 | pretrained_model_name_or_path = pretrained_model_name_or_path.replace( |
| 156 | cls.s3_prefix, "" |
| 157 | ) |
| 158 | |
| 159 | # Retry logic for downloading the model folder |
| 160 | retries = 3 |
| 161 | delay = 5 |
| 162 | attempt = 0 |
| 163 | success = False |
| 164 | while not success and attempt < retries: |
| 165 | try: |
| 166 | download_directory(pretrained_model_name_or_path, local_path) |
| 167 | success = True # If download succeeded |
| 168 | except Exception as e: |
| 169 | logger.error( |
| 170 | f"Error downloading model from {pretrained_model_name_or_path}. Attempt {attempt + 1} of {retries}. Error: {e}" |
| 171 | ) |
| 172 | attempt += 1 |
| 173 | if attempt < retries: |
| 174 | logger.info(f"Retrying in {delay} seconds...") |
| 175 | time.sleep(delay) # Wait before retrying |
| 176 | else: |
| 177 | logger.error( |
| 178 | f"Failed to download {pretrained_model_name_or_path} after {retries} attempts." |
| 179 | ) |
| 180 | raise e # Reraise exception after max retries |
| 181 | |
| 182 | return super().from_pretrained(local_path, *args, **kwargs) |
no test coverage detected