Update app.py

app.py

@@ -23,6 +23,7 @@ from accelerate import Accelerator
 import subprocess
 import sys
 import json
+import shutil
 
 # --- Configuration ---
 YOUR_HF_USERNAME = "Twelve2five"
@@ -333,19 +334,52 @@ def train_model(
     learning_rate=1e-4,
     progress=gr.Progress()
 ):
-    progress(0, desc="Installing dependencies...")
+    progress(0, desc="Setting up environment...")
     log = []
     
-    # Force reinstallation of transformers with specific version
-    log.append("Installing dependencies with specific versions...")
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", "transformers==4.36.2"])
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-U", "accelerate", "bitsandbytes", "peft", "datasets", "huggingface_hub", "deepspeed"])
+    # Completely clean up transformers installation
+    log.append("Completely reinstalling transformers and dependencies...")
+    
+    # First uninstall any existing transformers
+    subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", "transformers"])
+    
+    # Clean any cached files that might be causing issues
+    cache_dirs = [
+        os.path.expanduser("~/.cache/huggingface"),
+        os.path.expanduser("~/.cache/pip")
+    ]
+    
+    for cache_dir in cache_dirs:
+        if os.path.exists(cache_dir):
+            log.append(f"Cleaning cache directory: {cache_dir}")
+            try:
+                shutil.rmtree(cache_dir)
+            except Exception as e:
+                log.append(f"Warning: Could not clean {cache_dir}: {e}")
+    
+    # Install a stable version of transformers known to work with Llama models
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "transformers==4.35.2", "sentencepiece"])
+    
+    # Install other dependencies
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", 
+                          "accelerate", "bitsandbytes==0.41.1", "peft==0.6.1", 
+                          "datasets", "huggingface_hub", "deepspeed==0.12.3"])
     
     # Now import everything after installation to ensure we use the correct versions
-    from datasets import Dataset
-    from huggingface_hub import snapshot_download
-    from transformers import AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig, TrainingArguments, Trainer
-    from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
+    try:
+        from datasets import Dataset
+        from huggingface_hub import snapshot_download
+        import torch
+        import transformers
+        from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM
+        from transformers import BitsAndBytesConfig, TrainingArguments, Trainer
+        from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
+        
+        log.append(f"Transformers version: {transformers.__version__}")
+        log.append(f"PyTorch version: {torch.__version__}")
+    except ImportError as e:
+        log.append(f"Error importing libraries: {e}")
+        return "\n".join(log)
     
     # --- Configuration ---
     progress(0.05, desc="Setting up configuration...")
@@ -371,8 +405,10 @@ def train_model(
     progress(0.1, desc="Loading base model...")
     try:
         # First try to download the repo without loading the model
-        # to see what files are available
         local_model_path = "./model_files"
+        if os.path.exists(local_model_path):
+            shutil.rmtree(local_model_path)  # Clean up any previous files
+            
         snapshot_download(
             repo_id=hf_model_repo_id,
             local_dir=local_model_path,
@@ -387,29 +423,31 @@ def train_model(
                 config_data = json.load(f)
                 log.append(f"Model architecture type: {config_data.get('model_type', 'unknown')}")
                 
-                # Force model_type to llama if needed
-                if "architectures" in config_data and "LlamaForCausalLM" in config_data["architectures"]:
-                    config_data["model_type"] = "llama"
-                    with open(os.path.join(local_model_path, "config.json"), "w") as f:
-                        json.dump(config_data, f)
-                    log.append("Updated config.json to use llama model_type")
+                # Force model_type to llama
+                config_data["model_type"] = "llama"
+                if "architectures" in config_data:
+                    config_data["architectures"] = ["LlamaForCausalLM"]
+                
+                with open(os.path.join(local_model_path, "config.json"), "w") as f:
+                    json.dump(config_data, f)
+                log.append("Updated config.json to use llama model_type")
         
-        # Now try to load the config and model from local path
-        config = AutoConfig.from_pretrained(
+        # Now try to load with explicit Llama classes
+        config = LlamaConfig.from_pretrained(
             local_model_path,
-            trust_remote_code=False  # Set to False to avoid custom model code loading
+            trust_remote_code=False
         )
         
         log.append(f"Successfully loaded config: {config.model_type}")
         
-        # Load model with the config
-        model = AutoModelForCausalLM.from_pretrained(
+        # Load model with specific Llama class
+        model = LlamaForCausalLM.from_pretrained(
             local_model_path,
             config=config,
             quantization_config=bnb_config,
             device_map="auto",
-            trust_remote_code=False,
-            torch_dtype=torch.bfloat16
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True
         )
         
         log.append(f"Loaded model vocab size: {model.config.vocab_size}")
@@ -417,7 +455,20 @@ def train_model(
     except Exception as e:
         error_msg = f"Error loading model: {str(e)}"
         log.append(error_msg)
-        return "\n".join(log)
+        
+        # Try a fallback approach
+        try:
+            log.append("Trying fallback approach with AutoModelForCausalLM...")
+            model = AutoModelForCausalLM.from_pretrained(
+                local_model_path,
+                device_map="auto",
+                torch_dtype=torch.bfloat16,
+                low_cpu_mem_usage=True
+            )
+            log.append(f"Fallback model loaded successfully")
+        except Exception as e2:
+            log.append(f"Fallback approach also failed: {str(e2)}")
+            return "\n".join(log)
     
     # --- Prepare for K-bit Training & Apply LoRA ---
     progress(0.15, desc="Preparing model for fine-tuning...")