Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 import gradio as gr
 import os
+import logging
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -8,26 +9,96 @@ from transformers import (
     Trainer,
     DataCollatorForLanguageModeling
 )
+from datasets import load_dataset
 
-# Force CPU mode
+# Force CPU-only mode
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.environ["BITSANDBYTES_NOWELCOME"] = "1"
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+
 def train():
-    model = AutoModelForCausalLM.from_pretrained(
-        "microsoft/phi-2",
-        device_map="auto",
-        trust_remote_code=True,
-        load_in_4bit=False  # Disable quantization
-    )
+    try:
+        # Load model and tokenizer
+        model_name = "microsoft/phi-2"
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="cpu",
+            trust_remote_code=True,
+            load_in_4bit=False  # Disable quantization
+        )
+
+        # Add padding token
+        tokenizer.pad_token = tokenizer.eos_token
+
+        # Load sample dataset
+        dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
+
+        # Tokenization function
+        def tokenize_function(examples):
+            return tokenizer(
+                examples["text"],
+                padding="max_length",
+                truncation=True,
+                max_length=256,
+                return_tensors="pt",
+            )
+
+        tokenized_dataset = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["text"]
+        )
+
+        # Data collator
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False
+        )
+
+        # Training arguments
+        training_args = TrainingArguments(
+            output_dir="./results",
+            per_device_train_batch_size=2,
+            per_device_eval_batch_size=2,
+            num_train_epochs=1,  # Reduced for testing
+            logging_dir="./logs",
+            fp16=False,
+            bf16=False,
+            use_cpu=True  # Explicit CPU usage
+        )
+
+        # Trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_dataset["train"],
+            data_collator=data_collator,
+        )
+
+        # Start training
+        logging.info("Starting training...")
+        trainer.train()
+        logging.info("Training completed!")
+        
+        return "✅ Training successful! Model saved."
+
+    except Exception as e:
+        logging.error(f"Error: {str(e)}")
+        return f"❌ Training failed: {str(e)}"
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Phi-2 CPU Training")
+    start_btn = gr.Button("Start Training")
+    output = gr.Textbox()
     
-    training_args = TrainingArguments(
-        output_dir="./results",
-        per_device_train_batch_size=2,
-        num_train_epochs=3,
-        use_cpu=True,  # Explicit CPU usage
-        fp16=False,
-        bf16=False,
+    start_btn.click(
+        fn=train,
+        outputs=output
     )
-    
-    # Rest of training code...
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)