🔧🐢 Lazy Transformer Patch ⚡️✨

# Edit description:
# - Switched from `transformers` dynamic loading to direct BitNet package imports
# - Imported BitNetConfig, BitNetTokenizer, BitNetForCausalLM from `bitnet`
# - Removed AutoConfig/AutoTokenizer/AutoModelForCausalLM calls (left imports for no reason; orphaned)
# - Note to update requirements: `pip install bitnet`

https://pypi.org/project/bitnet/

app.py

@@ -2,6 +2,11 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 
+# use the official bitnet package to supply the missing code
+from bitnet.configuration_bitnet import BitNetConfig
+from bitnet.modeling_bitnet import BitNetForCausalLM
+from bitnet.tokenization_bitnet import BitNetTokenizer
+
 # Singleton for model and tokenizer
 _model = None
 _tokenizer = None
@@ -10,19 +15,13 @@ def load_model():
     global _model, _tokenizer
     if _model is None or _tokenizer is None:
         model_id = "microsoft/bitnet-b1.58-2B-4T"
-        _tokenizer = AutoTokenizer.from_pretrained(
-            model_id,
-            trust_remote_code=True
-        )
-        config = AutoConfig.from_pretrained(
-            model_id,
-            trust_remote_code=True
-        )
-        _model = AutoModelForCausalLM.from_pretrained(
+        # load tokenizer, config, and model from the bitnet pip package
+        _tokenizer = BitNetTokenizer.from_pretrained(model_id)
+        config     = BitNetConfig.from_pretrained(model_id)
+        _model     = BitNetForCausalLM.from_pretrained(
             model_id,
             config=config,
-            torch_dtype=torch.bfloat16,
-            trust_remote_code=True
+            torch_dtype=torch.bfloat16
         )
     return _model, _tokenizer