lockin2025

.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

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"""