Update app.py

app.py

@@ -1,51 +1,103 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-from peft import PeftModel
-
-class DeepSeekLoraCPUInference:
-    def __init__(self, base_model="deepseek-ai/deepseek-r1", fine_tuned_model="./deepseek_lora_finetuned"):
-        self.tokenizer = AutoTokenizer.from_pretrained(fine_tuned_model)
-        
-        # Load model in 4-bit on CPU (if no GPU is available)
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        quant_config = BitsAndBytesConfig(
-            load_in_4bit=True if device == "cuda" else False,  # Use 4-bit only if GPU is available
-            bnb_4bit_compute_dtype=torch.bfloat16,
+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"
+
+# ✅ Function to start training
+def train_model(dataset_url, model_url, epochs):
+    try:
+        # Load the tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_url)
+
+        # ✅ 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
         )
 
-        self.model = AutoModelForCausalLM.from_pretrained(
-            base_model, 
-            quantization_config=quant_config if device == "cuda" else None,
+        model = AutoModelForCausalLM.from_pretrained(
+            model_url,
+            quantization_config=bnb_config if device == "cuda" else None,
             device_map=device
         )
 
-        # Load fine-tuned LoRA model
-        self.model = PeftModel.from_pretrained(self.model, fine_tuned_model)
-        self.model.to(device)
-        self.model.eval()
-
-    def generate_text(self, prompt, max_length=200):
-        """Generates text efficiently using CPU or GPU."""
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(device)
-        
-        with torch.no_grad():
-            output = self.model.generate(
-                **inputs, 
-                max_length=max_length, 
-                temperature=0.7, 
-                top_p=0.9, 
-                repetition_penalty=1.1
-            )
-
-        return self.tokenizer.decode(output[0], skip_special_tokens=True)
-
-if __name__ == "__main__":
-    model = DeepSeekLoraCPUInference()
+        # ✅ Apply LoRA for efficient training
+        lora_config = LoraConfig(
+            task_type=TaskType.CAUSAL_LM,
+            r=8,
+            lora_alpha=32,
+            lora_dropout=0.1,
+            target_modules=["q_proj", "v_proj"]
+        )
+
+        model = get_peft_model(model, lora_config)
+        model.to(device)
+
+        # ✅ Load dataset
+        dataset = load_dataset(dataset_url)
+
+        # ✅ Tokenization function
+        def tokenize_function(examples):
+            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=256)
+
+        tokenized_datasets = dataset.map(tokenize_function, batched=True)
+        train_dataset = tokenized_datasets["train"]
+
+        # ✅ Training Arguments
+        training_args = TrainingArguments(
+            output_dir="./deepseek_lora_cpu",
+            evaluation_strategy="epoch",
+            learning_rate=5e-4,
+            per_device_train_batch_size=1,
+            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,
+            optim="adamw_torch",
+            report_to="none"
+        )
+
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset
+        )
+
+        # ✅ Start Training
+        trainer.train()
+
+        # ✅ Save the Fine-Tuned Model
+        model.save_pretrained("./deepseek_lora_finetuned")
+        tokenizer.save_pretrained("./deepseek_lora_finetuned")
+
+        return "✅ Training Completed! Model saved successfully."
     
-    prompt = "The implications of AI in the next decade are"
-    generated_text = model.generate_text(prompt)
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+# ✅ Gradio UI
+with gr.Blocks() as app:
+    gr.Markdown("# 🚀 AutoTrain DeepSeek R1 (CPU)")
+
+    dataset_url = gr.Textbox(label="Dataset URL (Hugging Face)", placeholder="e.g. samsum")
+    model_url = gr.Textbox(label="Model URL (Hugging Face)", placeholder="e.g. deepseek-ai/deepseek-r1")
+    epochs = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of Training Epochs")
+
+    train_button = gr.Button("Start Training")
+    output_text = gr.Textbox(label="Training Output")
+
+    train_button.click(train_model, inputs=[dataset_url, model_url, epochs], outputs=output_text)
 
-    print("\nGenerated Text:\n", generated_text)
+# ✅ Launch the app
+app.launch(server_name="0.0.0.0", server_port=7860)