diff --git a/.gitignore b/.gitignore index 5634417..6f206fc 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,50 @@ +# Model cache and checkpoints +models/cache/ +models/checkpoints/ +*.pt +*.bin +*.safetensors + +# Logs +logs/ +*.log + +# Environment variables .env -venv + +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# Virtual Environment +venv/ +env/ +ENV/ + +# IDE +.idea/ +.vscode/ +*.swp +*.swo + config/*.local.yaml -__pycache__ .mypy_cache/ .pytest_cache \ No newline at end of file diff --git a/models/llm/fingpt.py b/models/llm/fingpt.py index 6946d5a..631635d 100644 --- a/models/llm/fingpt.py +++ b/models/llm/fingpt.py @@ -1,9 +1,11 @@ import os +from pathlib import Path from typing import Any, Dict, List, Optional import torch from peft import PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer +from accelerate import init_empty_weights, load_checkpoint_and_dispatch from .base import BaseLLM @@ -18,22 +20,72 @@ def __init__(self, config: Dict[str, Any]): raise ValueError("HUGGINGFACE_TOKEN not found in environment") self.model = None self.tokenizer = None - self.base_model = "falcon" - self.peft_model = "FinGPT/fingpt-mt_falcon-7b_lora" - self.from_remote = True + # Update model paths to use base Falcon model + self.base_model = "tiiuae/falcon-7b" + self.peft_model = "FinGPT/fingpt-sentiment-falcon-7b" + self.model_cache_dir = Path("models/cache") + self.offload_folder = self.model_cache_dir / "fingpt_offload" + self.checkpoint_dir = self.model_cache_dir / "checkpoints" + + # Create necessary directories + for dir_path in [self.model_cache_dir, self.offload_folder, self.checkpoint_dir]: + dir_path.mkdir(parents=True, exist_ok=True) + self._load_model() def _load_model(self): - self.model = AutoModelForCausalLM.from_pretrained( - "tiiuae/falcon-7b", # Use specific model ID - token=self.token, - trust_remote_code=True, - device_map="auto", - ).to(self.device) + """Initialize model with proper device handling""" + try: + # 1. Initialize tokenizer from base model + self.tokenizer = AutoTokenizer.from_pretrained( + self.base_model, + token=self.token, + trust_remote_code=True, + cache_dir=str(self.model_cache_dir) + ) - self.tokenizer = AutoTokenizer.from_pretrained( - "tiiuae/falcon-7b", token=self.token - ) + # 2. Initialize model with accelerate + with init_empty_weights(): + model = AutoModelForCausalLM.from_pretrained( + self.base_model, + token=self.token, + trust_remote_code=True, + torch_dtype=torch.float16, + cache_dir=str(self.model_cache_dir) + ) + + # 3. Download PEFT model if needed and get local path + peft_path = self._ensure_peft_model_downloaded() + + # 4. Load and dispatch model with offload folder + self.model = load_checkpoint_and_dispatch( + model, + checkpoint=peft_path, + device_map="auto", + no_split_module_classes=["FalconDecoderLayer"], + dtype=torch.float16, + offload_folder=str(self.offload_folder) + ) + + except Exception as e: + raise RuntimeError(f"Failed to load FinGPT model: {str(e)}") + + def _ensure_peft_model_downloaded(self) -> str: + """Download PEFT model and return local path""" + from huggingface_hub import snapshot_download + + try: + # Download model files to checkpoint directory + local_path = snapshot_download( + repo_id=self.peft_model, + token=self.token, + cache_dir=str(self.checkpoint_dir), + local_files_only=False # Force check for updates + ) + return local_path + + except Exception as e: + raise RuntimeError(f"Failed to download PEFT model: {str(e)}") async def load_model(self) -> None: """Load FinGPT Falcon model"""