Update app.py

app.py

@@ -4,9 +4,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 
 # Model configuration
-# Alternative: Load base model and then load your safetensors weights
-MODEL_NAME = "gpt2"  # Use base GPT-2 model
-CUSTOM_WEIGHTS_PATH = "./model.safetensors"  # Your custom weights
+# Since you have all model files in Space root, try loading directly
+MODEL_NAME = "."  # Load from current directory with all your uploaded files
+CUSTOM_WEIGHTS_PATH = "./model.safetensors"  # Backup: your custom weights
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -25,55 +25,59 @@ def load_model():
     print(f"Loading model from: {MODEL_NAME}")
     print(f"Using device: {DEVICE}")
     
-    # Check if required files exist
+    # List available files for debugging
     import os
-    required_files = [
-        "model.safetensors",
-        "tokenizer_config.json"
-    ]
-    
-    missing_files = []
-    for file in required_files:
-        file_path = os.path.join(MODEL_NAME if MODEL_NAME != "." else "", file)
-        if not os.path.exists(file_path):
-            missing_files.append(file)
-    
-    if missing_files:
-        print(f"❌ Missing required files: {missing_files}")
-        print("Available files in directory:")
-        try:
-            files = os.listdir(MODEL_NAME if MODEL_NAME != "." else ".")
-            for f in files:
-                print(f"  - {f}")
-        except:
-            print("  Could not list directory contents")
-        raise FileNotFoundError(f"Missing required files: {missing_files}")
+    try:
+        current_files = os.listdir(".")
+        print("Available files in current directory:")
+        for f in current_files:
+            print(f"  - {f}")
+    except Exception as e:
+        print(f"Could not list directory: {e}")
     
     try:
-        # Load base model and tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-        
-        # Try to load your custom weights if they exist
-        if os.path.exists(CUSTOM_WEIGHTS_PATH):
-            print(f"Loading custom weights from: {CUSTOM_WEIGHTS_PATH}")
-            from safetensors.torch import load_file
-            custom_weights = load_file(CUSTOM_WEIGHTS_PATH)
-            model.load_state_dict(custom_weights, strict=False)
-            print("✅ Custom weights loaded!")
-        else:
-            print(f"⚠️ Custom weights not found at {CUSTOM_WEIGHTS_PATH}, using base model")
+        # First try to load directly from your uploaded files
+        print("Attempting to load model directly from uploaded files...")
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+            model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+            print("✅ Successfully loaded model directly from your uploaded files!")
+        except Exception as direct_load_error:
+            print(f"Direct load failed: {direct_load_error}")
+            print("Falling back to base model + custom weights...")
+            
+            # Fallback: Load base model and add custom weights
+            tokenizer = AutoTokenizer.from_pretrained("gpt2")
+            model = AutoModelForCausalLM.from_pretrained("gpt2")
+            
+            # Try to load your custom weights
+            if os.path.exists(CUSTOM_WEIGHTS_PATH):
+                print(f"Loading custom weights from: {CUSTOM_WEIGHTS_PATH}")
+                try:
+                    from safetensors.torch import load_file
+                    custom_weights = load_file(CUSTOM_WEIGHTS_PATH)
+                    
+                    # Load the weights into the model
+                    missing_keys, unexpected_keys = model.load_state_dict(custom_weights, strict=False)
+                    
+                    if missing_keys:
+                        print(f"⚠️ Missing keys: {len(missing_keys)} (this might be normal for LoRA models)")
+                    if unexpected_keys:
+                        print(f"⚠️ Unexpected keys: {len(unexpected_keys)}")
+                    
+                    print("✅ Custom weights loaded successfully!")
+                except Exception as e:
+                    print(f"⚠️ Could not load custom weights: {e}")
+                    print("Using base GPT-2 model instead")
         
         # Set pad token if not set
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
         
-        if DEVICE == "cpu":
-            model = model.to(DEVICE)
-        elif DEVICE == "cuda":
-            model = model.to(DEVICE)
+        # Move model to device
+        model = model.to(DEVICE)
         
-        print("✅ Model and tokenizer loaded successfully!")
+        print(f"✅ Model loaded successfully on {DEVICE}!")
         
         # Cache the loaded model and tokenizer
         _model = model