🔧🐢 Lazy Transformer Patch ⚡️✨ (#2)

- 🔧🐢 Lazy Transformer Patch ⚡️✨ (e11a1c01a7e673f9928b03ed640fd0d3de9ca6a9)


Co-authored-by: AI-B <[email protected]>

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