File size: 3,311 Bytes
6677176 |
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 |
# model_loader.py
import os
from pathlib import Path
import torch
def load_model_and_tokenizer(model_repo_dir_or_local_path: str, base_model_id: str = "google/muril-base-cased", device: str = None):
"""
Robust loader that:
- loads full model if model.safetensors or pytorch_model.bin exists in model_repo_dir_or_local_path
- otherwise loads base_model_id then applies PEFT adapter from adapter_model.safetensors (if present)
Returns (tokenizer, model_on_device, backend_str)
"""
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
model_dir = Path(model_repo_dir_or_local_path)
# prefer safetensors full model if present
full_model_files = ["model.safetensors", "pytorch_model.bin"]
adapter_files = ["adapter_model.safetensors", "adapter_config.json", "adapter.safetensors"]
# Import delayed to allow environment to control package errors
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
tokenizer = None
model = None
# 1) Try to load full model from model_dir
for f in full_model_files:
full_path = model_dir / f
if full_path.exists():
tokenizer = AutoTokenizer.from_pretrained(str(model_dir), use_fast=True)
try:
model = AutoModel.from_pretrained(str(model_dir), trust_remote_code=False)
backend = "full-AutoModel"
except Exception:
model = AutoModelForSequenceClassification.from_pretrained(str(model_dir), trust_remote_code=False)
backend = "full-AutoModelForSequenceClassification"
model.to(device)
model.eval()
return tokenizer, model, backend
# 2) If full model not present, check for adapter files
has_adapter = any((model_dir / af).exists() for af in adapter_files)
if has_adapter:
try:
tokenizer = AutoTokenizer.from_pretrained(str(model_dir), use_fast=True)
except Exception:
tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_fast=True)
try:
base = AutoModel.from_pretrained(str(model_dir))
base_name = str(model_dir)
except Exception:
base = AutoModel.from_pretrained(base_model_id)
base_name = base_model_id
base.to(device)
base.eval()
try:
from peft import PeftModel
adapter_dir = str(model_dir)
# If adapter files exist, PeftModel.from_pretrained will pick them up
peft_model = PeftModel.from_pretrained(base, adapter_dir, is_trainable=False)
peft_model.to(device)
peft_model.eval()
return tokenizer, peft_model, "peft-attached"
except Exception as e:
raise RuntimeError(f"Failed to load/apply PEFT adapter from {model_dir}: {e}") from e
# 3) Fallback: try direct load (may fail)
try:
tokenizer = AutoTokenizer.from_pretrained(str(model_dir), use_fast=True)
model = AutoModel.from_pretrained(str(model_dir))
model.to(device)
model.eval()
return tokenizer, model, "auto-fallback"
except Exception as e:
raise RuntimeError(f"Unable to load model or adapters from {model_dir}. Error: {e}") from e
|