Update app.py

app.py

@@ -15,11 +15,8 @@ from zeta.optim import StableAdamWUnfused
 import gradio as gr
 import os
 import subprocess
+os.system("pip install git+https://github.com/shumingma/transformers.git")
 
-def install(package):
-    subprocess.check_call([os.sys.executable, "-m", "pip", "install", "--ignore-installed", package])
-install("zetascale=2.8.2")
-install("git+https://github.com/shumingma/transformers.git#egg=transformers")
 # Suppress TorchDynamo errors (this will fallback to eager mode)
 import torch._dynamo
 torch._dynamo.config.suppress_errors = True
@@ -66,8 +63,12 @@ transformers.utils.logging.enable_explicit_format()
 model_id = "microsoft/bitnet-b1.58-2B-4T-bf16"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 hf_save_dir = "./bitnet"
-model = AutoModelForCausalLM.from_pretrained(model_id)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+device = model.device
 if torch.cuda.is_available():
     print("CUDA is available. Using GPU:", torch.cuda.get_device_name(0))
 else:
@@ -124,6 +125,7 @@ val_dataset = torch.utils.data.Subset(processed_dataset, list(range(split_idx, l
 # ---------------------------------------------------------------------------------
 # Collate function for DataLoader
 # ---------------------------------------------------------------------------------
+@spaces.GPU
 def sft_collate_fn(batch):
     """
     Collate a list of examples by padding them to the maximum sequence length in the batch.
@@ -159,6 +161,7 @@ optim = StableAdamWUnfused(model.parameters(), lr=LEARNING_RATE)
 # ---------------------------------------------------------------------------------
 # Define training function for Gradio UI
 # ---------------------------------------------------------------------------------
+@spaces.GPU
 def train_model():
     """
     Runs a training loop for a fixed number of batches and returns training logs.