Update app.py

app.py

@@ -5,6 +5,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 import math
 import os
+import sys
+import transformers
 
 class RMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-5):
@@ -191,98 +193,172 @@ model_id = "jatingocodeo/SmolLM2"
 
 def load_model():
     try:
-        print("Loading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        print("Tokenizer loaded successfully")
+        print("\n=== Starting model loading process ===")
+        print(f"Model ID: {model_id}")
         
-        # Ensure the tokenizer has the necessary special tokens
-        special_tokens = {
-            'pad_token': '[PAD]',
-            'eos_token': '</s>',
-            'bos_token': '<s>'
-        }
-        print("Adding special tokens...")
-        tokenizer.add_special_tokens(special_tokens)
+        print("\n1. Loading tokenizer...")
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            print("✓ Tokenizer loaded successfully")
+            print(f"Tokenizer type: {type(tokenizer)}")
+            print(f"Vocabulary size: {len(tokenizer)}")
+        except Exception as e:
+            print(f"× Error loading tokenizer: {str(e)}")
+            raise
         
-        print("Loading model from Hugging Face Hub...")
-        # Create config first
-        config = SmolLM2Config(
-            pad_token_id=tokenizer.pad_token_id,
-            bos_token_id=tokenizer.bos_token_id,
-            eos_token_id=tokenizer.eos_token_id
-        )
+        print("\n2. Adding special tokens...")
+        try:
+            special_tokens = {
+                'pad_token': '[PAD]',
+                'eos_token': '</s>',
+                'bos_token': '<s>'
+            }
+            num_added = tokenizer.add_special_tokens(special_tokens)
+            print(f"✓ Added {num_added} special tokens")
+            print(f"Special tokens: {tokenizer.special_tokens_map}")
+        except Exception as e:
+            print(f"× Error adding special tokens: {str(e)}")
+            raise
         
-        # Load model from Hub
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            config=config,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            trust_remote_code=True,
-            low_cpu_mem_usage=True
-        )
+        print("\n3. Creating model configuration...")
+        try:
+            config = SmolLM2Config(
+                pad_token_id=tokenizer.pad_token_id,
+                bos_token_id=tokenizer.bos_token_id,
+                eos_token_id=tokenizer.eos_token_id
+            )
+            print("✓ Configuration created successfully")
+            print(f"Config: {config}")
+        except Exception as e:
+            print(f"× Error creating configuration: {str(e)}")
+            raise
+        
+        print("\n4. Loading model from Hub...")
+        try:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                config=config,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+                local_files_only=False  # Force download from Hub
+            )
+            print("✓ Model loaded successfully")
+            print(f"Model type: {type(model)}")
+        except Exception as e:
+            print(f"× Error loading model: {str(e)}")
+            print("Attempting to print model files in Hub repo...")
+            from huggingface_hub import list_repo_files
+            try:
+                files = list_repo_files(model_id)
+                print(f"Files in repo: {files}")
+            except Exception as hub_e:
+                print(f"Error listing repo files: {str(hub_e)}")
+            raise
         
-        # Move model to device manually
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"Moving model to device: {device}")
-        model = model.to(device)
+        print("\n5. Moving model to device...")
+        try:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            print(f"Selected device: {device}")
+            model = model.to(device)
+            print("✓ Model moved to device successfully")
+        except Exception as e:
+            print(f"× Error moving model to device: {str(e)}")
+            raise
         
-        # Resize token embeddings to match new tokenizer
-        print("Resizing token embeddings...")
-        model.resize_token_embeddings(len(tokenizer))
+        print("\n6. Resizing token embeddings...")
+        try:
+            old_size = model.get_input_embeddings().weight.shape[0]
+            model.resize_token_embeddings(len(tokenizer))
+            new_size = model.get_input_embeddings().weight.shape[0]
+            print(f"✓ Token embeddings resized from {old_size} to {new_size}")
+        except Exception as e:
+            print(f"× Error resizing token embeddings: {str(e)}")
+            raise
         
-        print("Model loaded successfully!")
+        print("\n=== Model loading completed successfully! ===")
         return model, tokenizer
+        
     except Exception as e:
-        print(f"Error loading model: {str(e)}")
-        print(f"Error type: {type(e)}")
+        print("\n!!! ERROR IN MODEL LOADING !!!")
+        print(f"Error type: {type(e).__name__}")
+        print(f"Error message: {str(e)}")
+        print("\nFull traceback:")
         import traceback
         traceback.print_exc()
+        print("\nAdditional debug info:")
+        print(f"Python version: {sys.version}")
+        print(f"PyTorch version: {torch.__version__}")
+        print(f"Transformers version: {transformers.__version__}")
+        print(f"CUDA available: {torch.cuda.is_available()}")
+        if torch.cuda.is_available():
+            print(f"CUDA version: {torch.version.cuda}")
         raise
 
 def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
     try:
-        print(f"\nGenerating text for prompt: {prompt}")
-        # Load model and tokenizer (caching them for subsequent calls)
+        print("\n=== Starting text generation ===")
+        print(f"Input prompt: {prompt}")
+        print(f"Parameters: max_length={max_length}, temperature={temperature}, top_k={top_k}")
+        
         if not hasattr(generate_text, "model"):
-            print("First call - loading model...")
+            print("\n1. First call - loading model...")
             generate_text.model, generate_text.tokenizer = load_model()
         
-        # Ensure the prompt is not empty
         if not prompt.strip():
+            print("× Empty prompt received")
             return "Please enter a prompt."
         
-        # Add BOS token if needed
+        print("\n2. Processing prompt...")
         if not prompt.startswith(generate_text.tokenizer.bos_token):
             prompt = generate_text.tokenizer.bos_token + prompt
+            print("Added BOS token to prompt")
         
-        print("Encoding prompt...")
-        # Encode the prompt
-        input_ids = generate_text.tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=2048)
-        input_ids = input_ids.to(generate_text.model.device)
+        print("\n3. Encoding prompt...")
+        try:
+            input_ids = generate_text.tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=2048)
+            print(f"Encoded shape: {input_ids.shape}")
+            input_ids = input_ids.to(generate_text.model.device)
+            print("✓ Encoding successful")
+        except Exception as e:
+            print(f"× Error encoding prompt: {str(e)}")
+            raise
         
-        print("Generating text...")
-        # Generate text
-        with torch.no_grad():
-            output_ids = generate_text.model.generate(
-                input_ids,
-                max_length=min(max_length + len(input_ids[0]), 2048),
-                temperature=temperature,
-                top_k=top_k,
-                do_sample=True,
-                pad_token_id=generate_text.tokenizer.pad_token_id,
-                eos_token_id=generate_text.tokenizer.eos_token_id,
-                num_return_sequences=1
-            )
+        print("\n4. Generating text...")
+        try:
+            with torch.no_grad():
+                output_ids = generate_text.model.generate(
+                    input_ids,
+                    max_length=min(max_length + len(input_ids[0]), 2048),
+                    temperature=temperature,
+                    top_k=top_k,
+                    do_sample=True,
+                    pad_token_id=generate_text.tokenizer.pad_token_id,
+                    eos_token_id=generate_text.tokenizer.eos_token_id,
+                    num_return_sequences=1
+                )
+            print(f"Generation shape: {output_ids.shape}")
+        except Exception as e:
+            print(f"× Error during generation: {str(e)}")
+            raise
+        
+        print("\n5. Decoding output...")
+        try:
+            generated_text = generate_text.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+            print("✓ Decoding successful")
+            print(f"Output length: {len(generated_text)}")
+        except Exception as e:
+            print(f"× Error decoding output: {str(e)}")
+            raise
         
-        print("Decoding generated text...")
-        # Decode and return the generated text
-        generated_text = generate_text.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        print("Generation completed successfully!")
+        print("\n=== Generation completed successfully! ===")
         return generated_text.strip()
     
     except Exception as e:
-        print(f"Error during generation: {str(e)}")
-        print(f"Error type: {type(e)}")
+        print("\n!!! ERROR IN TEXT GENERATION !!!")
+        print(f"Error type: {type(e).__name__}")
+        print(f"Error message: {str(e)}")
+        print("\nFull traceback:")
         import traceback
         traceback.print_exc()
         return f"An error occurred: {str(e)}"