Update app.py

app.py

@@ -7,6 +7,8 @@ import math
 import os
 import sys
 import transformers
+from dotenv import load_dotenv
+from huggingface_hub import login
 
 class RMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-5):
@@ -226,16 +228,22 @@ def load_model():
         print("\n=== Starting model loading process ===")
         print(f"Model ID: {model_id}")
         
-        # Check HuggingFace access first
-        print("\nChecking HuggingFace access...")
-        if not check_huggingface_access():
-            raise Exception("Unable to access HuggingFace repository. Please check your token and repository permissions.")
+        # Load environment variables
+        load_dotenv()
+        
+        # Get HF token
+        hf_token = os.getenv('HF_TOKEN')
+        if not hf_token:
+            raise ValueError("HF_TOKEN not found in environment variables")
+            
+        # Login to Hugging Face
+        login(hf_token)
         
         print("\n1. Loading tokenizer...")
         try:
             tokenizer = AutoTokenizer.from_pretrained(
                 model_id,
-                use_auth_token=os.getenv("HF_TOKEN"),
+                use_auth_token=hf_token,
                 trust_remote_code=True
             )
             print("✓ Tokenizer loaded successfully")
@@ -259,37 +267,14 @@ def load_model():
             print(f"× Error adding special tokens: {str(e)}")
             raise
         
-        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...")
+        print("\n3. Loading model...")
         try:
-            # First try to list files in the repository
-            from huggingface_hub import list_repo_files
-            try:
-                files = list_repo_files(model_id, token=os.getenv("HF_TOKEN"))
-                print(f"Files in repository: {files}")
-            except Exception as hub_e:
-                print(f"Warning: Could not list repository files: {str(hub_e)}")
-            
             model = AutoModelForCausalLM.from_pretrained(
                 model_id,
-                config=config,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                use_auth_token=hf_token,
                 trust_remote_code=True,
-                use_auth_token=os.getenv("HF_TOKEN"),
-                low_cpu_mem_usage=True,
-                local_files_only=False
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                low_cpu_mem_usage=True
             )
             print("✓ Model loaded successfully")
             print(f"Model type: {type(model)}")
@@ -297,7 +282,7 @@ def load_model():
             print(f"× Error loading model: {str(e)}")
             raise
         
-        print("\n5. Moving model to device...")
+        print("\n4. Moving model to device...")
         try:
             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
             print(f"Selected device: {device}")
@@ -307,16 +292,6 @@ def load_model():
             print(f"× Error moving model to device: {str(e)}")
             raise
         
-        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("\n=== Model loading completed successfully! ===")
         return model, tokenizer