Update app.py

app.py

@@ -4,7 +4,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 
 # Model configuration
-MODEL_NAME = "model.safetensors"  # Replace with your actual HF model repo
+# Loading from current directory since model files are uploaded to Space root
+MODEL_NAME = "."  # Current directory contains model.safetensors and tokenizer files
+
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Global variables for model caching
@@ -22,34 +24,48 @@ def load_model():
     print(f"Loading model from: {MODEL_NAME}")
     print(f"Using device: {DEVICE}")
     
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    
-    # 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)
-    
-    print("✅ Model and tokenizer loaded successfully!")
-    
-    # Cache the loaded model and tokenizer
-    _model = model
-    _tokenizer = tokenizer
+    try:
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        
+        # 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)
+        
+        print("✅ Model and tokenizer loaded successfully!")
+        
+        # Cache the loaded model and tokenizer
+        _model = model
+        _tokenizer = tokenizer
+        
+        return model, tokenizer
     
-    return model, tokenizer
+    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")
+        raise e
 
 # Initialize model and tokenizer
-model, tokenizer = load_model()
+try:
+    model, tokenizer = load_model()
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    # Create dummy objects to prevent further errors
+    model, tokenizer = None, None
 
 def generate_code(pseudocode, indent=1, line=1, temperature=0.7, top_p=0.9, max_length=128):
     """
@@ -67,6 +83,20 @@ def generate_code(pseudocode, indent=1, line=1, temperature=0.7, top_p=0.9, max_
         Generated code string
     """
     try:
+        # Check if model is loaded
+        if model is None or tokenizer is None:
+            return """❌ Model not loaded. Please check:
+            
+1. MODEL_NAME in app.py - should be either:
+   - Your HF repository: "username/model-name"
+   - Local path: "./model" (if files uploaded to Space)
+
+2. If using HF repository, make sure it exists and is public
+
+3. If using local files, ensure model files are in correct folder
+
+Current MODEL_NAME: """ + MODEL_NAME
+        
         # Validate inputs
         if not pseudocode.strip():
             return "❌ Error: Please enter some pseudocode."