ankahi / ankahi_bundle /backend /model_loader.py

Add files using upload-large-folder tool

49c4f78 verified 6 months ago

3.9 kB

import torch
import time
import os
from transformers import AutoModelForCausalLM, AutoProcessor
from peft import PeftModel
from .config import MODEL_PATH, ADAPTERS_BASE

class ModelLoader:
    def __init__(self):
        self.model = None
        self.processor = None
        self.active_persona = None
        self.active_adapter = None
        self.load_start_time = None
        self.chat_template = None

    def load_base_model(self):
        print(f"Loading base model from {MODEL_PATH}...")
        self.load_start_time = time.time()
        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                MODEL_PATH,
                device_map="auto",
                torch_dtype=torch.bfloat16,
            )
            self.processor = AutoProcessor.from_pretrained(MODEL_PATH)
            self.model.eval()
            
            # Load chat template if exists
            chat_template_path = os.path.join(MODEL_PATH, "chat_template.jinja")
            if os.path.exists(chat_template_path):
                with open(chat_template_path, 'r') as f:
                    self.chat_template = f.read()
            
            print(f"Base model loaded in {time.time() - self.load_start_time:.2f}s")
            print(f"GPU memory allocated: {self.get_gpu_memory_gb():.2f} GB")
        except Exception as e:
            print(f"Error loading model: {e}")
            self.model = None
            raise

    def load_persona_adapter(self, persona_id):
        if persona_id == self.active_persona:
            return

        if not self.model:
            raise RuntimeError("Base model not loaded")

        adapter_path = os.path.join(ADAPTERS_BASE, persona_id, f"persona-{persona_id}.lora")
        
        # In some cases the directory structure might be different based on the artifacts
        # Checking if persona-persona_id.lora directory exists, or if it's just the adapter safely
        if not os.path.exists(adapter_path):
            # Fallback for common structure: artifacts/stage2/persona_id/adapter_model.safetensors
            # Wait, the instruction says: /artifacts/stage2/ananya/persona-ananya.lora/adapter_model.safetensors
            adapter_path = os.path.join(ADAPTERS_BASE, persona_id, f"persona-{persona_id}.lora")

        print(f"Switching to persona adapter: {persona_id} from {adapter_path}")
        start_time = time.time()
        
        try:
            # If an adapter is already loaded, we need to handle it.
            # PEFT models allow multiple adapters, but for simplicity we will unload/reload or just switch.
            if self.active_persona:
                # Merge and unload current adapter before loading new one to keep memory clean
                # or just use set_adapter if already loaded.
                # However, the instructions imply loading from disk.
                # PeftModel.from_pretrained on an already wrapped model adds a new adapter.
                if isinstance(self.model, PeftModel):
                    self.model = self.model.unload() # Unload back to base model

            self.model = PeftModel.from_pretrained(
                self.model,
                adapter_path,
                adapter_name=persona_id
            )
            self.active_persona = persona_id
            print(f"Adapter {persona_id} loaded in {time.time() - start_time:.2f}s")
        except Exception as e:
            print(f"Error loading adapter: {e}")
            raise

    def unload_adapter(self):
        if isinstance(self.model, PeftModel):
            self.model = self.model.unload()
            self.active_persona = None

    def get_gpu_memory_gb(self):
        if torch.cuda.is_available():
            return torch.cuda.memory_allocated() / 1e9
        return 0.0

    def is_ready(self):
        return self.model is not None

# Singleton instance
model_loader = ModelLoader()

	import torch
	import time
	import os
	from transformers import AutoModelForCausalLM, AutoProcessor
	from peft import PeftModel
	from .config import MODEL_PATH, ADAPTERS_BASE

	class ModelLoader:
	def __init__(self):
	self.model = None
	self.processor = None
	self.active_persona = None
	self.active_adapter = None
	self.load_start_time = None
	self.chat_template = None

	def load_base_model(self):
	print(f"Loading base model from {MODEL_PATH}...")
	self.load_start_time = time.time()
	try:
	self.model = AutoModelForCausalLM.from_pretrained(
	MODEL_PATH,
	device_map="auto",
	torch_dtype=torch.bfloat16,
	)
	self.processor = AutoProcessor.from_pretrained(MODEL_PATH)
	self.model.eval()

	# Load chat template if exists
	chat_template_path = os.path.join(MODEL_PATH, "chat_template.jinja")
	if os.path.exists(chat_template_path):
	with open(chat_template_path, 'r') as f:
	self.chat_template = f.read()

	print(f"Base model loaded in {time.time() - self.load_start_time:.2f}s")
	print(f"GPU memory allocated: {self.get_gpu_memory_gb():.2f} GB")
	except Exception as e:
	print(f"Error loading model: {e}")
	self.model = None
	raise

	def load_persona_adapter(self, persona_id):
	if persona_id == self.active_persona:
	return

	if not self.model:
	raise RuntimeError("Base model not loaded")

	adapter_path = os.path.join(ADAPTERS_BASE, persona_id, f"persona-{persona_id}.lora")

	# In some cases the directory structure might be different based on the artifacts
	# Checking if persona-persona_id.lora directory exists, or if it's just the adapter safely
	if not os.path.exists(adapter_path):
	# Fallback for common structure: artifacts/stage2/persona_id/adapter_model.safetensors
	# Wait, the instruction says: /artifacts/stage2/ananya/persona-ananya.lora/adapter_model.safetensors
	adapter_path = os.path.join(ADAPTERS_BASE, persona_id, f"persona-{persona_id}.lora")

	print(f"Switching to persona adapter: {persona_id} from {adapter_path}")
	start_time = time.time()

	try:
	# If an adapter is already loaded, we need to handle it.
	# PEFT models allow multiple adapters, but for simplicity we will unload/reload or just switch.
	if self.active_persona:
	# Merge and unload current adapter before loading new one to keep memory clean
	# or just use set_adapter if already loaded.
	# However, the instructions imply loading from disk.
	# PeftModel.from_pretrained on an already wrapped model adds a new adapter.
	if isinstance(self.model, PeftModel):
	self.model = self.model.unload() # Unload back to base model

	self.model = PeftModel.from_pretrained(
	self.model,
	adapter_path,
	adapter_name=persona_id
	)
	self.active_persona = persona_id
	print(f"Adapter {persona_id} loaded in {time.time() - start_time:.2f}s")
	except Exception as e:
	print(f"Error loading adapter: {e}")
	raise

	def unload_adapter(self):
	if isinstance(self.model, PeftModel):
	self.model = self.model.unload()
	self.active_persona = None

	def get_gpu_memory_gb(self):
	if torch.cuda.is_available():
	return torch.cuda.memory_allocated() / 1e9
	return 0.0

	def is_ready(self):
	return self.model is not None

	# Singleton instance
	model_loader = ModelLoader()