Enhance model loading in transformers with improved error handling and fallback mechanisms; implement compatibility checks for trust_remote_code parameter

app.py

@@ -290,15 +290,25 @@ def load_model_direct(model_type):
     
     try:
         # Load config from temp file
-        config = AutoConfig.from_pretrained(config_path, trust_remote_code=True)
+        try:
+            config = AutoConfig.from_pretrained(config_path)
+        except Exception:
+            # Try without trust_remote_code
+            config = AutoConfig.from_pretrained(config_path)
         
         # Load model with manual config
-        model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
-            f"PascalNotin/Tranception_{model_type}",
-            config=config,
-            trust_remote_code=True,
-            ignore_mismatched_sizes=True
-        )
+        try:
+            model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
+                f"PascalNotin/Tranception_{model_type}",
+                config=config,
+                ignore_mismatched_sizes=True
+            )
+        except TypeError:
+            # Fallback without newer parameters
+            model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
+                f"PascalNotin/Tranception_{model_type}",
+                config=config
+            )
         
         os.unlink(config_path)  # Clean up temp file
         return model
@@ -337,13 +347,28 @@ def load_model_cached(model_type):
         os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
         os.environ["HF_HUB_DISABLE_EXPERIMENTAL_WARNING"] = "1"
         
-        model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
-            model_path,
-            cache_dir=cache_dir,
-            force_download=True,
-            trust_remote_code=True,
-            resume_download=False
-        )
+        # Try loading without trust_remote_code first (compatibility issue)
+        try:
+            model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
+                model_path,
+                cache_dir=cache_dir,
+                force_download=True,
+                resume_download=False
+            )
+        except Exception as e1:
+            print(f"Loading without trust_remote_code failed: {e1}")
+            # Fallback: try with trust_remote_code for older transformers versions
+            try:
+                model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
+                    model_path,
+                    cache_dir=cache_dir,
+                    force_download=True,
+                    trust_remote_code=True,
+                    resume_download=False
+                )
+            except Exception as e2:
+                print(f"Loading with trust_remote_code also failed: {e2}")
+                raise e1  # Raise the original exception
         MODEL_CACHE[model_type] = model
         print(f"{model_type} model loaded and cached")
         return model
@@ -360,10 +385,17 @@ def load_model_cached(model_type):
         try:
             # Use full URL to bypass any path resolution issues
             full_url = f"https://huggingface.co/PascalNotin/Tranception_{model_type}"
-            model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
-                full_url,
-                cache_dir=cache_dir
-            )
+            try:
+                model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
+                    full_url,
+                    cache_dir=cache_dir
+                )
+            except TypeError:
+                # Try without trust_remote_code if it's not supported
+                model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
+                    full_url,
+                    cache_dir=cache_dir
+                )
             MODEL_CACHE[model_type] = model
             print(f"{model_type} model loaded successfully with full URL")
             return model
@@ -400,12 +432,18 @@ def load_model_cached(model_type):
                 print("Medium model failed, trying Small model...")
                 # Try Small model as last resort
                 try:
-                    model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
-                        "PascalNotin/Tranception_Small",
-                        trust_remote_code=True,
-                        force_download=True,
-                        cache_dir="/tmp/huggingface/small"
-                    )
+                    try:
+                        model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
+                            "PascalNotin/Tranception_Small",
+                            force_download=True,
+                            cache_dir="/tmp/huggingface/small"
+                        )
+                    except TypeError:
+                        model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
+                            "PascalNotin/Tranception_Small",
+                            force_download=True,
+                            cache_dir="/tmp/huggingface/small"
+                        )
                     MODEL_CACHE["Small"] = model
                     print("Small model loaded as fallback")
                     return model