Update app.py

app.py

@@ -4,8 +4,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 
 # Model configuration
-# Loading from current directory since model files are uploaded to Space root
-MODEL_NAME = "."  # Current directory contains model.safetensors and tokenizer files
+# 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
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -24,24 +25,53 @@ def load_model():
     print(f"Loading model from: {MODEL_NAME}")
     print(f"Using device: {DEVICE}")
     
+    # Check if required files exist
+    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:
-        # Load tokenizer
+        # 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")
         
         # Set pad token if not set
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
         
-        # Load model with appropriate settings
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_NAME,
-            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
-            device_map="auto" if DEVICE == "cuda" else None,
-            trust_remote_code=True
-        )
-        
         if DEVICE == "cpu":
             model = model.to(DEVICE)
+        elif DEVICE == "cuda":
+            model = model.to(DEVICE)
         
         print("✅ Model and tokenizer loaded successfully!")
         
@@ -54,9 +84,13 @@ def load_model():
     except Exception as e:
         print(f"❌ Error loading model: {e}")
         print("\n🔧 Troubleshooting:")
-        print("1. If using HF repository: Make sure MODEL_NAME is 'username/model-name'")
-        print("2. If using local files: Make sure model files are in the correct folder")
-        print("3. For private repos: Add authentication token")
+        print("1. Make sure you have uploaded ALL required files:")
+        print("   - model.safetensors (✅ you have this)")
+        print("   - config.json (❓ might be missing)")
+        print("   - tokenizer.json or vocab.json + merges.txt (❓ might be missing)")
+        print("   - tokenizer_config.json (✅ you have this)")
+        print("2. Files should be in the Space root directory")
+        print("3. Check if the model was saved correctly from your notebook")
         raise e
 
 # Initialize model and tokenizer