Spaces:
Running
Running
File size: 2,147 Bytes
3c0c995 0e85735 3c0c995 0e85735 3c0c995 0e85735 3c0c995 0e85735 3c0c995 0e85735 3c0c995 0e85735 3c0c995 0e85735 3c0c995 0e85735 3c0c995 0e85735 3c0c995 0e85735 | 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 | import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import gc
class ModelLoader:
def __init__(self):
self.text_model = None
self.tokenizer = None
self.device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {self.device}")
def load_text_model(self):
"""Safe Phi-3 model loading"""
if self.text_model is None:
model_id = "microsoft/Phi-3-mini-4k-instruct"
# === Proper 4-bit Quantization Config ===
if self.device == "cuda":
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
)
torch_dtype = torch.float16
else:
# CPU fallback - no quantization
quantization_config = None
torch_dtype = torch.float32
print(f"Loading {model_id} with {'4-bit quantization' if quantization_config else 'CPU mode'}...")
self.tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
self.text_model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=quantization_config, # ← Correct way
torch_dtype=torch_dtype,
device_map="auto" if self.device == "cuda" else None,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
if self.device == "cpu":
self.text_model = self.text_model.to(self.device)
return self.text_model, self.tokenizer
def clear_memory(self):
if self.text_model is not None:
del self.text_model
self.text_model = None
if self.tokenizer is not None:
del self.tokenizer
self.tokenizer = None
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache() |