Update app.py

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
 from peft import get_peft_model, LoraConfig, TaskType
 from datasets import load_dataset
-from bitsandbytes import BitsAndBytesConfig
 
 # ✅ Check if a GPU is available, otherwise use CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -11,24 +10,11 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # ✅ Function to start training
 def train_model(dataset_url, model_url, epochs):
     try:
-        # Load the tokenizer
+        # Load tokenizer and model
         tokenizer = AutoTokenizer.from_pretrained(model_url)
+        model = AutoModelForCausalLM.from_pretrained(model_url).to(device)
 
-        # ✅ Load model with 4-bit quantization for CPU efficiency
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True if device == "cuda" else False,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_use_double_quant=True
-        )
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_url,
-            quantization_config=bnb_config if device == "cuda" else None,
-            device_map=device
-        )
-
-        # ✅ Apply LoRA for efficient training
+        # ✅ Apply LoRA (Reduces trainable parameters)
         lora_config = LoraConfig(
             task_type=TaskType.CAUSAL_LM,
             r=8,
@@ -50,20 +36,20 @@ def train_model(dataset_url, model_url, epochs):
         tokenized_datasets = dataset.map(tokenize_function, batched=True)
         train_dataset = tokenized_datasets["train"]
 
-        # ✅ Training Arguments
+        # ✅ Training Arguments (Optimized for CPU)
         training_args = TrainingArguments(
             output_dir="./deepseek_lora_cpu",
             evaluation_strategy="epoch",
             learning_rate=5e-4,
-            per_device_train_batch_size=1,
+            per_device_train_batch_size=1,  # Keeps memory low
             per_device_eval_batch_size=1,
             num_train_epochs=int(epochs),
             save_strategy="epoch",
             save_total_limit=2,
             logging_dir="./logs",
             logging_steps=10,
-            fp16=False,
-            gradient_checkpointing=True,
+            fp16=False,  # Disable FP16 for CPU
+            gradient_checkpointing=True,  # Saves memory
             optim="adamw_torch",
             report_to="none"
         )