File size: 2,251 Bytes
aa77a41 9ab0f08 aa77a41 | 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 | import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# Define the model path (adjust as needed)
MODEL_PATH = "model_files" # Your fine-tuned model path
# System prompt for guiding model behavior
DEFAULT_PROMPT = """<|system|>
You are a compassionate listener. Respond with:
- Short, natural sentences
- Occasional empathetic sounds ("Oh...", "I see")
- Open-ended questions when appropriate
- Validation before advice
- Clear crisis handoff when needed
Examples of good responses:
1. "That sounds really overwhelming. Can you tell me more about what's been happening?"
2. "I'm hearing a lot of pain in what you're sharing. Have you talked to anyone about this?"
3. "This seems really important. Let's focus on how you're feeling right now."
</s>"""
def load_model():
"""
Loads the fine-tuned model and tokenizer with optimizations for memory and performance.
Returns:
model: The loaded Hugging Face model.
tokenizer: The corresponding tokenizer.
device: The device (CPU/GPU) the model is loaded on.
"""
print(f"🔍 Loading model from: {MODEL_PATH}")
# 1. Load Tokenizer
tokenizer = AutoTokenizer.from_pretrained(
MODEL_PATH,
cache_dir="./cache", # Cache directory for faster reloads
use_fast=True, # Use the fast tokenizer for better performance
padding_side="left" # Ensure padding is consistent for generation
)
# 2. Load Model with Memory Optimization
model = AutoModelForCausalLM.from_pretrained(
MODEL_PATH,
cache_dir="./cache",
trust_remote_code=True, # Allow custom model code
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, # Use FP16 on GPU
device_map="auto", # Automatically map model to available devices
load_in_4bit=True if torch.cuda.is_available() else False # Quantize to 4-bit on GPU
)
# 3. Set Device
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print("✅ Model successfully loaded.")
return model, tokenizer, device
# Test the loader
if __name__ == "__main__":
model, tokenizer, device = load_model()
print("Model and tokenizer successfully loaded.") |